1 Introduction↩︎

The remarkable success of deep learning relies on the ability of overparameterized neural networks to generalize well to unseen data, even when they possess enough capacity to memorize the training set with random labels [3]. Classical learning theory, which relies on uniform convergence and capacity measures like VC-dimension [4], [5], often fails to explain this phenomenon [3], [6]. Consequently, attention has shifted toward the concept of implicit bias [6], [7], which posits that gradient-based optimization methods favor solutions that generalize well.

The implicit bias in linear models is well-understood—gradient descent converges to the minimum Euclidean norm solution in regression [3], [8] and max-margin solution in classification [9]. As for nonlinear model, diagonal linear network, where each layer applies a diagonal linear transformation to the input, is a simple testbed for studying implicit bias. In [10], the authors showed that for a two-layer diagonal linear network, the implicit bias of gradient descent is the \(\ell_1\) norm under infinitesimal initialization, and to the \(L_2\) norm under large (infinite) initialization. They further demonstrated that similar results hold for deep diagonal linear networks under infinitesimal initialization of a specific direction.

Beyond characterizing convergence points, there is growing interest in understanding the complete trajectory of the dynamics. It is observed that under small initialization, saddle-to-saddle dynamics is present in many models, such as the diagonal linear network [1], [11], [12], matrix factorization [13], [14], ReLU network [15], [16]. The most relevant work is [1], which, building on the mirror flow technique from [17], proves that under infinitesimal initialization, the gradient flow trajectory of a two-layer diagonal linear network successively transitions from one saddle point to another, ultimately converging to the minimum \(\ell_1\)-norm solution. These saddle-to-saddle transitions are further characterized through a recursive algorithm.

The analysis in [1] relies on the explicit form of mirror flow. In [18], the author establishes the existence of mirror flow in a broader setting known as commuting parametrization, which encompasses diagonal linear networks as a special case. This suggests that the results of [1] could potentially be extended to more general architectures, although such a generalization remains an open question.

In this paper, we achieve this generalization via a novel technique called the Structural Invariant Manifold [2], which investigates data-independent properties of parametric models. We extend the analysis of saddle-to-saddle dynamics in [1] to both deep diagonal linear networks and general two-layer diagonal linear networks (defined in Definition 1). As a direct corollary, we characterize the implicit bias of the two types of diagonal linear network.

1.1 Contributions↩︎

Training Trajectory: Theorem 1 in [1] addresses the saddle-to-saddle dynamics of two-layer diagonal linear networks and deep networks restricted to specific initialization directions. We generalize this in Theorem 3 by extending the analysis to general two-layer networks (Definition 1) and establishing results for deep networks under generic initialization directions.

Implicit bias: [10], [19] established that the implicit bias corresponds to \(\ell_1\)-norm minimization for standard two-layer diagonal linear networks and deep diagonal linear networks under specific initialization directions. We generalize these findings by characterizing the implicit bias as a modified \(\ell_1\)-norm minimization for general two-layer diagonal linear networks and deep diagonal linear networks under generic initialization directions (Corollary 1).

Mechanism: We provide a theoretical analysis of the underlying mechanism that leads to the equivalence between the dynamics and the recursive algorithm, identifying the Structural Invariant Manifold [2] as the key for this behavior (Section 6).

2 Related Works↩︎

Diagonal Linear Network: Early work on diagonal linear networks focused on establishing the implicit bias of gradient flow [10], [17], [19] and mirror flow [17], [20]. Subsequent research has elucidated the geometric nature of these optimization paths, specifically saddle-to-saddle dynamics [1], [11], [12], [21]. More recently, literature has addressed the algorithmic impact of discrete hyperparameters, such as step size [22], [23], stochastic noise [24], and momentum [25].

Matrix Factorization: Since diagonal linear networks can be viewed as a special case of matrix factorization, related literature is particularly relevant. [19] prove that for commuting observations, gradient descent exhibits an implicit bias toward low nuclear norm solutions. [13] provide both theoretical and empirical evidence that gradient flow with infinitesimally small initialization behaves like a greedy low-rank algorithm. However, [1] demonstrate that this greedy low-rank behavior does not extend to diagonal linear networks. [14] further identify the “connectivity” of observations as a key factor influencing the dynamics. Extending the technique in this paper to the matrix factorization setting is an interesting direction for future research.

3 Preliminary↩︎

Analytic Parametric Model and Its Training: In this paper, we use \(F(\vtheta)(\vx)\) to denote a parametric model. \(\vtheta\in \sR^M\) is the parameter of the model, \(\vx\in \sR^d\) is the input of the model. The output of the model is \(F(\vtheta)(\vx)\in \sR\). If \(F\) is real-analytic for \((\vtheta,\vx)\in \sR^{M}\times \sR^d\), we say \(F\) is an analytic parametric model. To train the parametric model, we collect data \(S=\{(\vx_i,y_i)\}_{i=1}^n\). We define MSE loss \(L(\vtheta)=\sum_{i=1}^n (F(\vtheta)(\vx_i)-y_i)^2\). We consider its gradient flow \[\frac{\rd \vtheta}{\rd t}=-\nabla_{\vtheta} L(\vtheta), \vtheta(0)=\vtheta_0. \label{eq:gradient32flow}\tag{1}\] Here, \(\vtheta_0\) is the initialization.

Diagonal Linear Network: An \(L\)-layer diagonal linear model is \(F(\vtheta)(\vx)=\sum_{i=1}^n a_{i,1}a_{i,2}\cdots a_{i,L} x_i\), where \(\vx=(x_1,\ldots,x_n)\in \sR^n\), \(\vtheta=(a_{i,1},a_{i,2}\ldots ,a_{i,L})_{i=1}^n\in \sR^{nL}\). The training data of diagonal linear model is denoted as \(\vX,\vy\), where \(\vX\) is an \(m\times n\) matrix, \(\vy\) is an \(m\times 1\) vector. Each row of \(\vX\) and \(\vy\) is a sample. The MSE loss is \(L(\vtheta)=\|\vX\vk(\vtheta)-\vy\|_2^2\). Here, \(\vk(\vtheta)=(k_i(\vtheta_i))_{i=1}^n,\) where \(k_i(\vtheta_i)=a_{i,1}\times a_{i,2}\times \cdots \times a_{i,L}.\) If \(L\geq 3\), we call it the deep diagonal linear network.

Notation: We use \([n]\) to denote the set \(\{1,\ldots,n\}\). We use bold letters to denote vectors, and non-bold letters to denote scalars. For a vector \(\vv\), we use \(v_i\) denotes the i-th coordinate of \(\vv\). For a matrix \(\vX\), we use \(\vX_{ij}\) to denote the element in the \(i\)-th row and \(j\)-th column, \(\vX_i\) to denote the \(i\)-th row of \(\vX\), and \(\vX_{:,j}\) to denote the \(j\)-th column of \(\vX\).

Infinitesimal initialization: Let \(\vtheta_0\) denote a fixed base parameter vector. We set the initialization of the model weights as: \(\vtheta(0) = s \vtheta_0\) where \(s > 0\) is a scalar multiplier. The “infinitesimal initialization" refers to the asymptotic limit as \(s \to 0\).

\(\ell_1\) Implicit Bias: Two-layer diagonal linear model is known to exhibit implicit regularization of \(\ell_1\) norm under infinitesimal initialization [19]. Mathematically, if we denote the solution of Equation 1 as \(\phi(\vtheta_0,t)\), then we have \[\lim_{s \to 0}\lim_{t\to +\infty} \phi(s\vtheta_0,t)=\vk^*,\] where \(\vk^*\) is a solution of the following \(\ell_1\) minimization problem: \[\min _{\vk\in \sR^n} \|\vk\|_1 \quad s.t. \quad \vX\vk=\vy.\]

4 Main Results↩︎

Summary of Main Results: Briefly speaking, in Corollary 1, we prove the implicit bias of deep diagonal linear network and general two-layer diagonal linear network (defined in Definition 1) under infinitesimal initialization. The implicit bias is a modified \(\ell_1\) norm: \[R(\vk)=\sum_{i=1}^n (t_i^+k_i^+ +t_i^-k_i^-).\]

The proof of Corollary 1 consists of two steps. In Theorem 3, we prove that the dynamics of the two models is equivalent to Algorithm 1. In Theorem 1, we prove that Algorithm 1 will converge to the solution of \[\min _{\vk} R(\vk) \quad s.t. \quad \vX\vk=\vy.\] Then the implicit bias is a direct corollary.

4.1 The Algorithm↩︎

4.2 Dynamics of diagonal linear network↩︎

5 Intuition of the Dynamics↩︎

In this section, we use a three-layer diagonal linear network as an illustrative example. The proof of Theorem 3 serves as a rigorous formalization of the intuition developed in this section.

Specifically, we analyze the gradient flow dynamics of the model \(F(\vtheta)(\vx)=\sum_{i=1}^n a_ib_ic_ix_i\) under a squared error loss \(L(\vtheta)\). For a single neuron (dropping the index \(i\)), the dynamics for its parameters \((a,b,c)\) are given by: \[\begin{align} \frac{\rd a}{\rd t} &= bc \cdot v, & \frac{\rd b}{\rd t} &= ac \cdot v, & \frac{\rd c}{\rd t} &= ab \cdot v, \end{align}\] where \(v = \sum_{k=1}^m \vX_{k}(y_k - F(\vtheta)(\vX_k))\) is a term related to the correlation between the neuron’s feature and the prediction error.

5.1 The Feature Selection Phase and Distance-Speed Argument↩︎

In the early phase of training, with infinitesimally small initial parameters \(\vtheta(0)\), the model output \(F(\vtheta)(\vx)\) is negligible. Consequently, \(v(t)\) remains approximately constant, \(v(t) \approx v_0 = \sum_{k=1}^m \vX_{k}y_k\).

This constant, \(v_0\), can be interpreted as the speed of the neuron’s evolution. By rescaling time as \(\tau = v_0 t\), we can isolate the geometry of the parameter trajectory: \[\begin{align} \frac{\rd a}{\rd \tau} &= bc, & \frac{\rd b}{\rd \tau} &= ac, & \frac{\rd c}{\rd \tau} &= ab. \end{align}\] These dynamics possess two conserved quantities: \[\begin{align} a^2(\tau) - c^2(\tau) &= a^2(0) - c^2(0) := \mu_1^2 \\ b^2(\tau) - c^2(\tau) &= b^2(0) - c^2(0) := \mu_2^2 \end{align}\] Without loss of generality, let us assume \(a(0)^2 > b(0)^2 > c(0)^2 > 0\) and \(a(0)b(0) > 0\). Substituting the conserved quantities into the dynamic for \(c\) yields: \[\frac{\rd c}{\rd \tau} = \sqrt{c^2 + \mu_1^2} \sqrt{c^2 + \mu_2^2}.\] Now, consider an initialization of scale \(s \ll 1\): \(a(0)=a_0 s, b(0)=b_0 s, c(0)=c_0 s\). The conserved quantities become \(\mu_1^2 = (a_0^2-c_0^2)s^2 := \tilde{\mu}_1^2 s^2\) and \(\mu_2^2 = (b_0^2-c_0^2)s^2 := \tilde{\mu}_2^2 s^2\). By defining a scaled parameter \(\tilde{c}(\tau) = c(\tau)/s\), its dynamic is: \[\frac{\rd \tilde{c}}{\rd \tau} = s \sqrt{\tilde{c}^2 + \tilde{\mu}_1^2} \sqrt{\tilde{c}^2 + \tilde{\mu}_2^2}, \quad \text{with} \quad \tilde{c}(0) = c_0.\] A neuron becomes “significant" when its parameters grow from \(\mathcal{O}(s)\) to \(\mathcal{O}(1)\). In the scaled frame, this corresponds to \(\tilde{c}\) growing from \(c_0\) to \(\pm\infty\). The time required for this growth (the blow-up time) can be found by separating variables. For instance, the time to grow to positive significance is: \[\begin{align} \tau^+ &= \int_{c_0}^{\infty} \frac{\rd \tilde{c}}{s \sqrt{\tilde{c}^2 + \tilde{\mu}_1^2} \sqrt{\tilde{c}^2 + \tilde{\mu}_2^2}}\\& = \frac{1}{s} \int_{c_0}^{\infty} \frac{\rd x}{\sqrt{x^2 + \tilde{\mu}_1^2} \sqrt{x^2 + \tilde{\mu}_2^2}}. \end{align}\] This structure motivates the core intuition. We define the distance a neuron must travel as the integral, which depends only on the initial direction \((a_0, b_0, c_0)\): \[D^+ = \int_{c_0}^{\infty} \frac{\rd x}{\sqrt{x^2 + \tilde{\mu}_1^2} \sqrt{x^2 + \tilde{\mu}_2^2}}.\] The time to grow to significance in the original time \(t\) is obtained by un-scaling \(\tau^+\): \[t^+ = \frac{\tau^+}{|v_0|} = \frac{1}{|v_0|s} D^+. \label{eq:distance32speed}\tag{3}\]

Distance-Speed Argument: Equation 3 reveals that the time for a neuron to become significant is proportional to a “distance” \(D^+\) (determined by its initialization) and inversely proportional to its “speed” \(|v_0|\) (determined by its feature’s correlation with the target). The neurons with shortest time will be \(\fO(1)\) first, and therefore it will be chosen as the grown feature.

5.2 The Learning Phase and Sign-Locking of Features↩︎

Learning Phase: Once a feature, say \(k_j = a_jb_jc_j\), has grown to a significant magnitude at time \(t_{\min}\), the dynamics of the system change. The approximation that the learned vector \(\vk\) is near zero no longer holds. The system now enters a “learning phase" where the gradient descent dynamics actively adjust the significant features to minimize the loss function \(L(\vtheta)\).

Mismatch of Time Scales of Two Phases: There exists a mismatch in the time scales of the feature selection and learning phases. Specifically, the feature selection phase persists for a duration on the order of \(\mathcal{O}\left(\frac{1}{s}\right)\), whereas the duration of the learning phase is independent of \(s\). Consequently, when \(s\) is sufficiently small, the time spent in the learning phase becomes negligible in comparison to that of the feature selection phase. This implies that, during training, no new features are developed away from zero.

Sign-Locking: The key insight is that once \(k_j(t)\) has become significantly non-zero, its sign is effectively locked for a very long time. Consider the case where \(k_j(t)\) has grown to be positive, meaning \(a_j(t),b_j(t),c_j(t)\) are all large. Since we have conserved quantities \[\begin{align} a_j^2(\tau) - c_j^2(\tau) &= a_j^2(0) - c_j^2(0) \\ b_j^2(\tau) - c_j^2(\tau) &= b_j^2(0) - c_j^2(0), \end{align}\] therefore, the parameter \(\vtheta_i(\tau):=(a_j(\tau),b_j(\tau),c_j(\tau))\) is restricted to a curve during training. As a consequence, for \(k_j(t)\) to flip its sign and become negative, \(a_j(t),b_j(t),c_j(t)\) would first have to decrease back to its initialization, which is \(\fO(s)\) scale. So the time for \(k_j(t)\) to flip its sign and become negative is of order \(\fO(\frac{1}{s})\). Since we assume \(s\) is a sufficiently small number representing the scale of initialization, this required time is enormous.

This reveals a mismatch in time scales during the learning dynamics. Consider the second learning phase. Let \(k_j\) and \(k_l\) denote the feature chosen at the first and second feature selection phase, respectively. We compare the following two time scales:

• The time required for \(k_j\) to flip its sign is of order \(\frac{1}{s}\), which diverges to infinity as \(s \to 0\).

• The time required for \(k_j\) and \(k_l\) to reach a minimum of the loss function \(L(\boldsymbol{\theta})\) without changing their signs is approximately independent of \(s\), since both coordinates have already moved away from zero.

Consequently, when \(s\) is sufficiently small, \(k_j\) and \(k_l\) reach the minimum long before \(k_j\) is able to flip its sign. This results in a learning phase that can be effectively modeled as a sign-constrained optimization problem. For instance, if \(k_j\) becomes positive during the first learning phase, then the second learning phase can be described by the following constrained optimization: \[\min_{k_j, k_l} \left\| \vX_{:,j} k_j + \vX_{:,l} k_l - \vy \right\|_2^2 \quad \text{subject to} \quad k_j \geq 0.\]

5.3 Comprehensive Dynamics↩︎

Inter-Phase Continuity of Feature Selection: Although certain neurons are not selected during a given feature selection phase, they still undergo non-trivial movement and settle at specific positions within the interval \((-t^-, t^+)\), as characterized by Equation 3 . Given that the time scale of the subsequent feature learning phase is negligible, it does not significantly affect the final positions of these neurons. As a result, in the next feature selection phase, each non-activated neuron resumes its dynamics from the position reached at the end of the previous selection phase.

Feature Selection Drives Subsequent Learning: The training process alternates between feature selection and feature learning phases, each exerting influence on the other. During a feature selection phase, a new feature is typically activated. Consequently, the subsequent learning phase jointly optimizes this newly selected feature along with those previously selected.

Learning Shapes Future Selection Dynamics: Conversely, each learning phase modifies the correlation vector \(\vv = \vX^\T(\vy-\vX\vk^{(p)} )\), which governs the “speed” of neurons in the ensuing feature selection phase. As a result, the selection behavior in each feature selection phase is inherently influenced by the convergence point of the preceding learning phase.

Comprehensive Dynamics: The training process proceeds iteratively through alternating phases. At each iteration \(p\), the following steps are performed:

1. Feature Selection Phase: A new neuron \(j_p\) is selected from the inactive set based on the minimal activation time, computed as the ratio of “distance” to “speed”. The distance corresponds to the remaining gap inherited from the end of the previous selection phase, while the speed is determined by the correlation vector \(\vv = \vX^\T(\vy-\vX\vk^{(p)})\).

2. Learning Phase: The current active set of neurons \(\mathcal{A}_p = \{j_1, \dots, j_p\}\) jointly minimizes the loss function, subject to individual sign constraints on their corresponding coefficients.

3. Termination: The procedure repeats until convergence, i.e., the gradient flow converges to its global minimum.

This framework of sequential feature activation and joint optimization defines the overall training trajectory and is exactly Algorithm 1.

6 Mechanism Behind the Dynamics↩︎

In Section 5, we used a three-layer diagonal linear network as an illustrative example to explain the validity of Theorem 3. In this section, we aim to uncover a more general underlying mechanism that helps us understand the dynamics of diagonal linear networks.

6.1 Structural Invariant Manifold (SIM)↩︎

6.2 Implications of SIM in Dynamics↩︎

SIM Induces Two-Phase Dynamics: In Theorem 7, we show that for each \(i \in [n]\), if the initial parameter satisfies \(\vtheta_i = \vzero\), then \(O_{\fF_i}(\vtheta_i) = \{\vzero\}\). This invariance implies that parameters initialized exactly at zero remain stationary under gradient flow indefinitely. In our setting, however, parameters are initialized infinitesimally close to zero rather than exactly at zero. Then such parameters require an infinite amount of time to move significantly away from the origin. Consequently, the learning dynamics exhibit a natural separation into two distinct phases: an initial feature selection phase, during which parameters remain near zero for an extended period, followed by a learning phase that unfolds over a finite time scale.

SIM Induces Sign Constraint During the Learning Phase: By Theorem 7, under gradient flow dynamics, if the initialization \(\vtheta^* = (\vtheta_i^*)_{i=1}^n\) satisfies \(\vtheta_i^* \neq \vzero\), then for each \(i \in [n]\), the trajectory \(\vtheta_i(t)\) evolves along a simple curve for all \(t \in \mathbb{R}\). Suppose, for contradiction, that \(\vtheta_i(t)\) undergoes a sign change during training; that is, there exist times \(t_1 < t_2\) such that \(\vtheta_i(t_1) > \delta > 0\) and \(\vtheta_i(t_2) < \delta' < 0\). Since \(\vtheta_i(t)\) evolves along a continuous simple curve, the intermediate value theorem implies the existence of some \(\hat{t} \in (t_1, t_2)\) such that \(\vtheta_i(\hat{t}) = \vtheta_i^*\).

Now, if the initialization \(\vtheta_i^*\) is chosen to be infinitesimally small, then reaching this value along the trajectory requires an infinite amount of time. However, the learning phase proceeds over a finite time horizon. As a result, such a sign change cannot occur within the learning phase. This implies that, due to the structure imposed by SIM, the sign of each \(\vtheta_i(t)\) is effectively preserved throughout the learning phase.

6.3 Mechanism of Incremental Learning↩︎

In this context, incremental learning refers to the phenomenon whereby only one neuron is selected during the feature selection phase.

Consider model of the form \(F(\vtheta)(\vx)=\sum_{i=1}^n k_i(\vtheta_i)x_i\). For each \(i \in [n]\), the gradient flow dynamics of the parameter \(\vtheta_i\) are given by: \[\frac{\rd \vtheta_i}{\rd t} = \nabla k_i(\vtheta_i) \, u_i(\vtheta),\] where \(\vu(\vtheta) = \vX^\T (\vy-\vX \vk(\vtheta))\). During the feature selection phase, the change in \(\vtheta\) is small due to the scale of initialization. Consequently, \(u_i(\vtheta)\) remains approximately constant and can be approximated by a fixed value \(u_i^*\). Therefore, the dynamics simplify to: \[\frac{\rd \vtheta_i}{\rd t} \approx \nabla k_i(\vtheta_i) \, u_i^*.\]

Under this approximation, \(u_i^*\) can be interpreted as the effective “speed” at which \(\vtheta_i\) evolves along the curve \(O_{\fF_i}(\vtheta_i)\).

As established in Section 5.1 and Lemma [lemma:32short32time32leads32to32zero], for deep diagonal linear networks and general two-layer diagonal linear networks with an initialization scale \(s\), there exists a scaling function \(h(s)\) such that \(\lim_{s \to 0} h(s) = +\infty\). This function governs the temporal dynamics required for a neuron to escape the initialization regime. Specifically, the time required for a neuron to attain an \(O(1)\) magnitude in the positive or negative direction is approximately \(t_i^+ h(s)\) and \(t_i^- h(s)\), respectively. For instance:

• \(h(s) = \frac{1}{s}\) for three-layer diagonal linear networks,

• \(h(s) = -\log s\) for general diagonal linear networks.

Define the quantity \[\delta_i := \begin{cases} \frac{t_i^+}{u_i^*}, & \text{if } u_i^* > 0, \\[8pt] \frac{t_i^-}{|u_i^*|}, & \text{if } u_i^* < 0. \end{cases}\] \(\delta_i\) represents the effective time required for neuron \(i\) to \(O(1)\), taking into account both its growth direction and speed. The neuron with the smallest \(\delta_i\) will be the one that activates first and is thus selected during the feature selection phase. Under generic conditions on the data \(\vX\) and \(\vy\), the minimum of \(\delta_i\) is attained uniquely, implying that only one neuron is selected in this phase.

7 Experiments↩︎

We consider the model \(F(\vtheta)(\vx)=\sum_{i=1}^4 k_i(\vtheta_i)x_i\), where \[k_i(\vtheta_i)=2a_ib_i+(a_i^2+b_i^2+c_i^2)c_i \, , i=1,2 ;\] \[k_i(\vtheta_i)=a_i\tanh(a_i)-(e^{b_i}-1)^2\,, i=3,4.\] One can readily verify that this model satisfies Definition 1. Ideally, to verify Theorem 3, we would require infinitesimal initialization and gradient flow, which are impractical to implement in experimental settings. As a practical alternative, we initialize the model with a sufficiently small scale of \(10^{-60}\), and employ a relatively large learning rate of \(0.1\) to accelerate the training of gradient descent.

The data matrix and label vector are given by: \[\vX = \begin{pmatrix} 1 & 0.5 & 0.7 & 0 \\ 0.5 & 1 & 0.1 & 0.7 \end{pmatrix}, \quad \vy = \begin{pmatrix} 1 \\ 0 \end{pmatrix}. \label{eq:data32matrix}\tag{4}\]

As illustrated in Figure 2, the training dynamics demonstrate transitions between successive saddle points. The three dashed lines in the second panel indicate the values of \(k_i(\cdot)\) following each learning phase. In Appendix 12.1, we compute the values of \(\vk^{(p)}\) in Algorithm 1 for \(p = 1, 2, 3\), and confirm that these computed values closely correspond to those represented by the dashed lines in the figure.

8 Limitations and Discussion↩︎

8.1 Limitations↩︎

Infinitesimal Initialization: Our analysis is confined to the regime of infinitesimal initialization. While this setting facilitates theoretical tractability, it presents two notable drawbacks: (i) Infinitesimal initialization is not practical in real-world applications, where parameters are typically initialized with finite variance; (ii) Prior works on implicit bias, such as [10], is able to characterize the implicit bias under all initialization scales.

Gradient Flow Assumption: The theoretical analysis in this paper is conducted under the assumption of continuous-time gradient flow dynamics. However, in practice, optimization is performed using discrete-time algorithms such as gradient descent or adaptive methods like Adam. The extent to which our results extend to these more realistic optimization settings remains an open question.

8.2 Discussion↩︎

8.2.0.1 Technique:

We highlight that the technique employed in this work relies on the concept of the Structural Invariant Manifold, which generally arises in nonlinear models such as matrix sensing and neural networks [2], [27]–[29]. The potential outcomes of this study may be extendable to more complex models such as matrix sensing. Investigating this generalization is an interesting direction for future research.

8.2.0.2 Similarity and difference between \(\ell_1\) and modified \(\ell_1\) minimization:

Both the standard and modified \(\ell_1\) norms fundamentally promote sparsity. Their difference is that the modified \(\ell_1\) norm can alter the recovered support (i.e., which specific features the model chooses to select). To illustrate this, consider the data matrix presented in Equation 4 .

• Standard \(\ell_1\) minimization yields the solution: \(k_3=1.43, k_4=-0.20\), with \(k_1=k_2=0\) (feature support: \(\{3, 4\}\)).

• Modified \(\ell_1\) minimization (applying a \(0.5\) penalty weight to \(|k_1|\)) shifts the optimal solution to: \(k_1=1.0, k_4=-0.71\), with \(k_2=k_3=0\) (feature support: \(\{1, 4\}\)).

This demonstrates that though the solution remains sparse, the actual features learned by the network may differ.

8.2.0.3 Gap between gradient flow (GF) and gradient descent (GD):

Whether the SIM remains exactly invariant under GD or SGD depends on its geometric curvature:

• Affine (Flat) SIMs: If a SIM is an affine subspace, it remains strictly invariant under GD/SGD. Because the tangent space is globally constant, discrete steps do not cause the trajectory to leave the manifold. For instance, in a standard diagonal linear network parameterized by \(k_i(\vtheta_i) = a_i b_i\), the affine SIM defined by \(\{ (\vtheta_i)_{i=1}^n \mid a_1 = b_1 = 0\}\) is invariant under GD and SGD.

• Curved SIMs: If a SIM is curved, it generally loses its exact invariance under GD/SGD. Taking a finite discrete step along the tangent space of a curved manifold naturally causes the parameter to diverge slightly from the exact manifold. For example, the curved SIM defined by \(\{(\vtheta_i)_{i=1}^n \mid a_1^2 - b_1^2 = 1\}\) is not exactly invariant under GD.

This theoretical gap explains a specific phenomenon observed in the training dynamics of Figure 2 (optimized via GD, learning rate \(\eta=0.1\)). Between epochs \(3000\) and \(4000\), \(k_1\) changes sign but fails to return to its exact initial scale (\(10^{-120}\)). This behavior is directly driven by the loss of exact invariance on the curved SIM caused by the discretization.

9 Conclusion↩︎

In this paper, we investigated the training dynamics of deep diagonal linear networks and general two-layer diagonal linear networks. We showed that their training dynamics can be equivalently described by Algorithm 1. Furthermore, we proved that Algorithm 1 converges to the solution of a modified \(\ell_1\) norm minimization problem. As a result, we established that the implicit bias of both types of diagonal linear networks under infinitesimal initialization corresponds to a modified \(\ell_1\) norm.

In addition, we analyzed the underlying mechanisms driving these dynamics and identified the Structural Invariant Manifold (SIM) as the key geometric structure governing the learning process.

Acknowledgements↩︎

This work was supported by the National Key R&D Program of China (Grant No. 2022YFA1008200), the National Natural Science Foundation of China (Grant No.12571567), the Natural Science Foundation of Shanghai (Grant No. 25ZR1402280).

Impact Statement↩︎

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

10 Proof of Theorems↩︎

10.1 Proof of Theorem 1↩︎

We first prove a lemma.

In the following we prove Theorem 1.

10.2 Proof of Theorem 3↩︎

Example of limit trajectory: Consider \(k_i(\vtheta)=a^2-b^2, \vtheta=(a,b)\in \sR^2\). Then \(\Gamma_i^{++}=\{(a,b)\mid a>0,b=0 \}\), \(\Gamma_i^{+-}=\{(a,b)\mid a<0,b=0 \}\), \(\Gamma_i^{-+}=\{(a,b)\mid a=0,b>0 \}\), \(\Gamma_i^{--}=\{(a,b)\mid a=0,b<0 \}.\) In this case of \(k_i(\vtheta)=a^2-b^2\), the limit trajectory can also be derived from the conserved quantity \(a(t)b(t)=a(0)b(0)\). Under infinitesimal initialization, we have \(a(t)b(t)=0\). This results in the four limit trajectories.

11 Concepts in Differential Geometry↩︎

In the appendix, we present several definitions and concepts in differential geometry that are pertinent to the content of this paper.

12 Experiments↩︎

12.1 Details of Figure 2↩︎

We consider the model \(F(\vtheta)(\vx)=\sum_{i=1}^4 k_i(\vtheta_i)\) such that: \[k_i(\vtheta_i)=2a_ib_i+(a_i^2+b_i^2+c_i^2)c_i \, , i=1,2\] \[k_i(\vtheta)=a_i\tanh(a_i)-(e^{b_i}-1)^2\,, i=3,4.\] One can readily verify that this model satisfies Definition 1, and \(t_i^+=t_i^-=1\) for all \(i\in [n]\). We have the data matrix \[\vX = \begin{pmatrix} 1 & 0.5 & 0.7 & 0 \\ 0.5 & 1 & 0.1 & 0.7 \end{pmatrix}, \quad \vy = \begin{pmatrix} 1 \\ 0 \end{pmatrix}.\]

Experiment details: The parameters are initialized uniformly in the interval \([-10^{-60},10^{-60}]\). The learning rate is \(0.1\). We use mpmath library to support calculation of high-precision. In the code, we set mp.mp.dps = 200, which means that we use \(200\) decimal places of precision for calculations.

Output of Algorithm 1: In the following we calculate the output of Algorithm 1 under the input \(\vX,\vy, t_i^+,t_i^-.\)

• \(p = 0\); The initialization is \(\vk^{(0)}=\vs^{(0)}=\vzero\).

• \(p=1\); We have \(I_2=\{1,2,3,4\}.\) By calculation, \(\vu=\vX^\T \vy=(1,0.5,0.7,0)^\T.\) So the time vector \(\vdelta=(1,2,\frac{10}{7},+\infty).\) Therefore, the first neuron will be chosen as the feature. So \(\vs^{(1)}=(1,0.5,0.7,0)^\T.\) \(k_1^{(1)}\) is the solution of \[\min_{k_1\in \sR} \| \vX_{:,1}k_1-\vy\|_2^2.\] This leads to \(k_1^{(1)}=0.8.\) So \(\vk^{(1)}=(0.8,0,0,0)^\T.\) To summarize, we have \[\vk^{(1)}=(0.8,0,0,0)^\T \, , \, \vs^{(1)}=(1,0.5,0.7,0)^\T.\]

• \(p=2\); We have \(I_2=\{2,3,4\}.\) By calculation, we have \[u_2,u_3,u_4=-0.3,0.1,-0.28.\] So we have \[\delta_2,\delta_3,\delta_4=5,3,\frac{25}{7}.\] Therefore, \(\delta_3\) is the smallest. So the third neuron is chosen. Then \(\vs^{(2)}=(1,-0.4,1,-0.84)^\T.\)

  Besides, \(k_1^{(2)}\) and \(k_3^{(2)}\) are the solution of \[\min_{(k_1,k_3)\in \sR^2} \| \vX_{:,1}k_1+\vX_{:,3}k_3-\vy\|_2^2 \quad s.t. \quad k_1\geq 0.\] The solution is \(k_1=0,k_3=1.4.\) So we have \(\vk^{(2)}=(0,0,1.4,0)^\T.\) To summarize, we have \[\vk^{(2)}=(0,0,1.4,0)^\T \, , \, \vs^{(2)}=(1,-0.4,1,-0.84)^\T.\]

• \(p=3\); We have \(I_2=\{1,2,4\}.\) By calculation, we have \[u_1,u_2,u_4=-0.05,-0.13,-0.098.\] So we have \[\delta_1,\delta_2,\delta_4=20,\frac{60}{13}, \frac{80}{49}.\] Therefore, \(\delta_4\) is the smallest. So the fourth neuron is chosen. \(k_3^{(3)}\) and \(k_4^{(3)}\) are the solution of \[\min_{(k_3,k_4)\in \sR^2} \| \vX_{:,3}k_3+\vX_{:,4}k_4-\vy\|_2^2 \quad s.t. \quad k_3\geq 0.\] The solution is \(k_3=\frac{10}{7},k_4=-\frac{10}{49}.\) So we have \(\vk^{(3)}=(0,0,\frac{10}{7},-\frac{10}{49})^\T.\) Besides, it is readily to verify that \(\vX\vk^{(3)}=\vy.\) So the iteration stops.

  To summarize, Algorithm 1 ends at \(p=3\), and we have \[\vk^{(3)}=(0,0,\frac{10}{7},-\frac{10}{49})^\T.\]

Comparison of \(\vk^{(p)}\) to the dashed line in Figure 2:

• In the first dashed line, we have \((k_1,k_2,k_3,k_4)^\T=(0.800,8.34e-70,1.51e-28,-1.35e-100)^\T.\) This value is close to \(\vk^{(1)}=(0.8,0,0,0)^\T.\)

• In the second dashed line, we have \((k_1,k_2,k_3,k_4)^\T=(-1.01e-6,-6.23e-78,1.400,-1.34e-17)^\T.\) This value is close to \(\vk^{(2)}=(0,0,1.4,0)^\T.\)

• In the second dashed line, we have \((k_1,k_2,k_3,k_4)^\T=(-1.01e-6,-6.23e-78,1.400,-1.34e-17)^\T.\) This value is close to \(\vk^{(2)}=(0,0,1.4,0)^\T.\)

• In the third dashed line, we have \((k_1,k_2,k_3,k_4)^\T=(-1.71e-6,-2.48e-57,1.429,-0.204)^\T.\) This value is close to \(\vk^{(3)}=(0,0,\frac{10}{7},-\frac{10}{49})^\T.\)

Difference between gradient flow and gradient descent: In gradient flow, as shown in Theorem 7, the parameter \(\vtheta_i\) is strictly restricted to a simple curve. As a consequence, if \(k_i(\vtheta_i)\) wants to change its sign, then \(\vtheta_i\) must must pass its initialization. Therefore, if \(k_i(\vtheta_i)\) wants to change its sign, \(|k_i|\) must decrease to its initialization scale.

However, in Figure 2, we see that the value of \(k_1\) is \(0.800\) at the first dashed line, and \(-1.01e-6\) at the second dashed line. Between the two dashed lines, the value of \(|k_1|\) does not decrease to its initialization scale \(1e-60.\) This is different from the dynamics of gradient flow. We attribute this phenomenon to the fact that Structural Invariant Manifold in Theorem 7 is not invariant under gradient descent (since \(O_{\fF_i}(\vtheta_i)\) is not a straight line). The failure of Theorem 7 in gradient descent is also evidenced by [31].

12.2 Deep Diagonal Linear Network↩︎

Figure 3 illustrates the training dynamics of a Deep Diagonal Linear Network (\(w = a \odot b \odot c\)) optimized via gradient descent. The experiment utilizes a synthetic dataset with \(M=3\) samples and \(N=8\) features. We employ a small initialization scale (\(\approx 10^{-4}\)) while maintaining a strict layer-wise hierarchy (\(a > b > c\)) to facilitate theoretical analysis.

Figures 3 (a) and  3 (b) demonstrate a clear synchronization between the descent of the loss curve and the sequential growth of neural parameters. The learning process exhibits distinct stage-wise behavior, where the system transitions between saddle points. We define the end of each stage as the point where the active parameters stabilize and the gradient of the residual becomes negligible.

Figure 3 (c) presents the theoretical “growth time" for each neuron at the onset of every stage, calculated by our algorithm as the integrated virtual distance required to escape the saddle point. Notably, these theoretical predictions align perfectly with the empirical emergence order and relative timing observed in Figure 2.

Figure 4 demonstrates the controllability of feature selection through initialization. While the network naturally prioritizes features based on data correlation, we show that this order can be manipulated. By selectively amplifying the initialization magnitude of the 5th neuron by a factor of 5 (while keeping others at the original scale), we successfully induce this neuron to emerge first, overriding the natural data-driven order. Crucially, our theoretical framework accurately captures this perturbation, correctly predicting the altered emergence sequence in Figure 4.

Figure 5 investigates the implicit bias under a specific “Zero-Balanced" initialization scheme (\(a=b, c=0\), with scale \(\approx 10^{-4}\)). This setting creates a strict saddle point at initialization. To verify the regularization properties, we computed the mathematical minimum \(\ell_1\)-norm solution for the generated dataset. As shown in Figure 4, as the loss approaches zero, the parameters converge precisely to the \(\ell_1\) solution rather than the sparse ground truth (\(\ell_0\) solution). This confirms that the \(c=0\) initialization induces a strong inductive bias towards \(\ell_1\) regularization, effectively driving the network to select the solution with the minimum norm among infinite possibilities.

References↩︎

---

1. See Lemma 2 and Definition 4 for detailed explanations. This assumption primarily fails if \(k_i(\boldsymbol{\theta}_i)\) is globally bounded from above (i.e., \(k_i(\boldsymbol{\theta}_i) < M\) for all \(\boldsymbol{\theta}_i\)) or below.↩︎

2. See Definition 8 for the definition of invariant manifold.↩︎

3. There is a detailed explanation of orbit and its properties in chapter 2 of [26].↩︎

4. Here, a simple curve refers to an analytic curve that does not intersect itself.↩︎