Gradient Flow

Table of Contents

1. Learning Objectives
2. Prerequisites
3. Notation
4. Core Intuition
5. The Linear Regression Model
6. From Gradient Descent to Gradient Flow
7. Gradient Flow for Least Squares Loss
8. Evolution of the Residual
9. Closed-Form Solution for Function Evolution
10. Trainability and Convergence Analysis
11. Evolution of Parameters
12. Comparison to the Closed-Form OLS Solution
13. Evolution of Model Predictions
14. Distribution Over Predictions
15. Common Pitfalls and Misconceptions
16. Research Perspective
17. Summary of Takeaways
18. Exercises

Learning Objectives

After completing this chapter, you should be able to:

1. Derive gradient flow as the continuous-time limit ($\alpha \to 0$) of gradient descent.
2. Write the gradient flow ODE for least squares loss.
3. Derive the ODE governing the evolution of the residual vector.
4. Solve the residual ODE via matrix exponential and interpret the solution.
5. Characterize trainability and convergence speed from the eigenvalues of $\mathbf{X}^T\mathbf{X}$.
6. Derive the closed-form parameter evolution and verify agreement with the OLS solution at $t \to \infty$.
7. Compute the evolution of predictions at new test points.
8. Describe how a Gaussian prior over initial parameters induces a distribution over predictions.
9. Explain the connection to the Neural Tangent Kernel.

Prerequisites

• Gradient Descent — discrete update rule, learning rate, convergence conditions
• Eigenvalues and Eigenvectors — spectral decomposition, matrix exponential
• Singular Value Decomposition — pseudoinverse, rank, null space
• Linear Regression — OLS solution, design matrix

Notation

• $\boldsymbol\phi \in \mathbb{R}^D$ — model parameters
• $\mathbf{x}_i \in \mathbb{R}^D$ — input vector (augmented with 1 for bias)
• $\mathbf{X} \in \mathbb{R}^{D \times I}$ — data matrix (columns are inputs)
• $\mathbf{y} \in \mathbb{R}^I$ — target vector
• $f[\mathbf{x}, \boldsymbol\phi] = \mathbf{x}^T \boldsymbol\phi$ — linear model
• $\mathbf{f}[\mathbf{X}, \boldsymbol\phi] = \mathbf{X}^T \boldsymbol\phi \in \mathbb{R}^I$ — predictions at all training points
• $L[\boldsymbol\phi]$ — loss function
• $\alpha$ — learning rate (step size)
• $t$ — continuous time variable
• $\boldsymbol\phi_t$ — parameters at time $t$
• $\boldsymbol\phi_0$ — initial parameters
• $\mathbf{A}^+$ — Moore–Penrose pseudoinverse of $\mathbf{A}$

Core Intuition

Gradient descent updates parameters in discrete steps of size $\alpha$. Gradient flow is what happens when we make the step size infinitesimally small: the discrete update becomes an ordinary differential equation (ODE). For the linear model, this ODE has an exact closed-form solution.

Why is this useful? Because the closed-form solution lets us:

• Determine whether training will converge (trainability).
• Compute how fast convergence happens (convergence speed).
• Track the exact trajectory of parameters and predictions over time.
• Understand which directions in function space are learned first.

This analysis provides the conceptual foundation for the Neural Tangent Kernel (NTK), which extends these ideas to neural networks.

Series context. This is the first part of a series on gradient-based training dynamics. Part II extends to neural networks and derives the NTK. Part III considers the Bayesian counterpart (see Bayesian Linear Regression).

The Linear Regression Model

A linear model computes a scalar output:

where $\mathbf{x} \in \mathbb{R}^D$ is the (augmented) input and $\boldsymbol\phi \in \mathbb{R}^D$ are the parameters. For an affine model, we set $\mathbf{x} \leftarrow [1, \mathbf{x}^T]^T$ so the first component of $\boldsymbol\phi$ acts as the bias.

Given $I$ training pairs $\lbrace \mathbf{x}_i, y_i \rbrace_{i=1}^I$, we minimize the least squares loss:

In matrix form, collecting all inputs as columns of $\mathbf{X} \in \mathbb{R}^{D \times I}$ and targets into $\mathbf{y} \in \mathbb{R}^I$:

Geometric picture. The loss landscape is a paraboloid in $\mathbb{R}^D$. Isocontours are ellipsoids whose axes align with the eigenvectors of $\mathbf{X}\mathbf{X}^T$.

From Gradient Descent to Gradient Flow

Discrete gradient descent applies the update:

Rearranging:

Taking the limit $\alpha \to 0$ transforms the left side into a time derivative:

Definition (Gradient Flow). Equation (6) is called gradient flow. It is an autonomous ODE describing the continuous-time evolution of parameters under infinitesimal learning rate.

Key properties:

• The loss is monotonically non-increasing along the trajectory: $\frac{dL}{dt} = \left(\frac{\partial L}{\partial \boldsymbol\phi}\right)^T \frac{d\boldsymbol\phi}{dt} = -\lVert \nabla L \rVert^2 \leq 0$.
• For convex $L$, the trajectory converges to the global minimum.
• For the linear model, the ODE is linear in $\boldsymbol\phi$ and thus exactly solvable.

Gradient Flow for Least Squares Loss

Expanding the right-hand side of (6) for the quadratic loss:

Applying the chain rule to the quadratic form:

For the linear model, $\mathbf{f}[\mathbf{X}, \boldsymbol\phi] = \mathbf{X}^T \boldsymbol\phi$ and $\frac{\partial \mathbf{f}}{\partial \boldsymbol\phi} = \mathbf{X}^T$, so:

This is a linear ODE with constant coefficients (since $\mathbf{X}\mathbf{X}^T$ does not depend on $\boldsymbol\phi$).

Evolution of the Residual

Define the residual vector $\mathbf{r}(t) = \mathbf{f}[\mathbf{X}, \boldsymbol\phi_t] - \mathbf{y}$. We seek an ODE for $\mathbf{r}(t)$.

Since $\mathbf{y}$ is constant:

Substituting (7b):

Definition (Tangent Kernel). The matrix $\mathbf{K} = \frac{\partial \mathbf{f}}{\partial \boldsymbol\phi} \cdot \frac{\partial \mathbf{f}}{\partial \boldsymbol\phi}^T \in \mathbb{R}^{I \times I}$ is the tangent kernel (or empirical NTK for neural networks). For the linear model:

The residual ODE is therefore:

This is a matrix exponential decay equation.

Closed-Form Solution for Function Evolution

When $\mathbf{K}$ is constant (which holds exactly for the linear model, since $\frac{\partial \mathbf{f}}{\partial \boldsymbol\phi} = \mathbf{X}^T$ is independent of $\boldsymbol\phi$), equation (8d) has the solution:

Substituting back $\mathbf{r}(t) = \mathbf{f}[\mathbf{X}, \boldsymbol\phi_t] - \mathbf{y}$ and $\mathbf{r}(0) = \mathbf{X}^T\boldsymbol\phi_0 - \mathbf{y}$:

Verification. At $t = 0$: $\mathbf{f} = \mathbf{y} + (\mathbf{X}^T\boldsymbol\phi_0 - \mathbf{y}) = \mathbf{X}^T\boldsymbol\phi_0$. Correct.

Interpretation. The model predictions start at their initial values $\mathbf{X}^T\boldsymbol\phi_0$ and exponentially decay toward the targets $\mathbf{y}$. The rate of decay is governed by the eigenvalues of $\mathbf{K} = \mathbf{X}^T\mathbf{X}$.

Trainability and Convergence Analysis

Spectral decomposition. Let $\mathbf{K} = \mathbf{X}^T\mathbf{X} = \sum_{j=1}^I \lambda_j \mathbf{u}_j \mathbf{u}_j^T$ where $\lambda_1 \geq \lambda_2 \geq \cdots \geq \lambda_I \geq 0$. Then:

Theorem 1 (Trainability). The linear model can fit the training data exactly (i.e., $\mathbf{r}(t) \to \mathbf{0}$ as $t \to \infty$) if and only if $\mathbf{K} = \mathbf{X}^T\mathbf{X}$ is full rank.

Proof. The residual at time $t$ is $\mathbf{r}(t) = \sum_j e^{-\lambda_j t} (\mathbf{u}_j^T \mathbf{r}(0)) \mathbf{u}_j$. Each component decays as $e^{-\lambda_j t}$. If $\lambda_j > 0$ for all $j$, every component decays to zero. If $\lambda_j = 0$ for some $j$, the corresponding component $\mathbf{u}_j^T \mathbf{r}(0)$ remains unchanged. $\blacksquare$

Corollary. The matrix $\mathbf{X} \in \mathbb{R}^{D \times I}$ has $\text{rank}(\mathbf{X}) \leq \min(D, I)$. The kernel $\mathbf{K} \in \mathbb{R}^{I \times I}$ is full rank when $I \leq D$ (more parameters than data). In the overparameterized regime, the model can interpolate training data.

Theorem 2 (Convergence Speed). The convergence rate is dominated by the smallest nonzero eigenvalue $\lambda_{\min}$ of $\mathbf{K}$:

The direction associated with $\lambda_{\max}$ converges fastest; the direction associated with $\lambda_{\min}$ converges slowest.

Implication. The condition number $\kappa = \lambda_{\max} / \lambda_{\min}$ of $\mathbf{K}$ determines how "anisotropic" convergence is. High condition number means some directions converge much faster than others — the loss drops quickly along principal axes but slowly along the minor axis. This motivates preconditioning and adaptive methods (see Gradient Descent).

Evolution of Parameters

We can also derive a closed-form for how the parameters themselves evolve. Start from the Taylor linearization (exact for linear models):

Equating with (10) and solving for $\boldsymbol\phi_t$:

Applying the pseudoinverse $(\mathbf{X}^T)^+$ to both sides:

This gives the exact parameter trajectory from initialization $\boldsymbol\phi_0$ as a function of continuous time $t$.

Comparison to the Closed-Form OLS Solution

Recall from Linear Regression that the OLS solution is:

Proposition (Asymptotic Agreement). As $t \to \infty$, the gradient flow solution (15) converges to the OLS solution (16).

Proof. Rearranging (15):

The first term is $\hat{\boldsymbol\phi}_{\text{OLS}}$. In the second term, $\mathbf{X}\exp[-\mathbf{X}^T\mathbf{X} \cdot t] \to \mathbf{0}$ as $t \to \infty$ because any component in the null space of $\mathbf{X}^T\mathbf{X}$ is annihilated by the left-multiplication by $\mathbf{X}$. Therefore $\boldsymbol\phi_\infty = \hat{\boldsymbol\phi}_{\text{OLS}}$. $\blacksquare$

Remark. In the underdetermined case ($I < D$), there are infinitely many solutions with zero training loss. Gradient flow converges to the minimum-norm solution $\boldsymbol\phi_\infty = \mathbf{X}(\mathbf{X}^T\mathbf{X})^{-1}\mathbf{y}$, which is the solution closest to the initialization $\boldsymbol\phi_0$ in the Euclidean sense. This implicit regularization is a key insight for understanding neural networks.

Evolution of Model Predictions

For a new test point $\mathbf{x}^*$, the prediction at time $t$ is $y^*_t = \mathbf{x}^{*T}\boldsymbol\phi_t$. Substituting (15):

As $t \to \infty$:

Observation. The final prediction (19) depends on the training data only through $\mathbf{X}\mathbf{y}$ (a sufficient statistic for the linear model). The prediction can be written as $\hat{y}^* = \mathbf{k}(\mathbf{x}^*)^T \mathbf{K}^{-1} \mathbf{y}$ where $\mathbf{k}(\mathbf{x}^*) = \mathbf{X}^T\mathbf{x}^*$. This has the same form as a kernel regression prediction — foreshadowing the NTK.

Distribution Over Predictions

Setup. Define a Gaussian prior over initial parameters:

Since $y^*_t$ is an affine function of $\boldsymbol\phi_0$ (equation 18), and affine transformations of Gaussians are Gaussian (see Multivariate Gaussian), we can compute:

where $a(t)$ is a scalar offset (from the $\mathbf{y}$ terms) and $\mathbf{b}(t)$ is a vector depending on $\mathbf{x}^*, \mathbf{X}, t$.

Result. The prediction follows:

Interpretation:

• At $t = 0$: high variance (we haven't learned anything yet), mean prediction is zero.
• As $t$ increases: variance decreases as the model commits to fitting training data; mean approaches the OLS prediction.
• Adding observation noise $\sigma^2$ to the predictive variance accounts for irreducible noise.

This gives a training-time-dependent uncertainty quantification — the model becomes more certain as training progresses.

Common Pitfalls and Misconceptions

Pitfall 1: Gradient flow $\neq$ gradient descent. Gradient flow is the $\alpha \to 0$ limit. With finite learning rate, trajectories can diverge, oscillate, or take different paths (especially near saddle points). Results here hold approximately for small $\alpha$.

Pitfall 2: The kernel $\mathbf{K}$ is constant only for linear models. For neural networks, $\frac{\partial \mathbf{f}}{\partial \boldsymbol\phi}$ changes as $\boldsymbol\phi$ evolves. The NTK theory shows that in the infinite-width limit, this change vanishes and the kernel becomes approximately constant — but this is an asymptotic result, not exact.

Pitfall 3: Confusing $\mathbf{X}^T\mathbf{X}$ with $\mathbf{X}\mathbf{X}^T$. The kernel governing function-space evolution is $\mathbf{K} = \mathbf{X}^T\mathbf{X} \in \mathbb{R}^{I \times I}$ (Gram matrix). The matrix governing parameter-space dynamics is $\mathbf{X}\mathbf{X}^T \in \mathbb{R}^{D \times D}$. Their nonzero eigenvalues are the same, but dimensions differ.

Pitfall 4: Minimum-norm is not always desirable. Gradient flow finds the minimum-norm interpolant in the underdetermined case. For neural networks, this implicit bias can help generalization (flat minima), but can also overfit depending on the parameterization.

Research Perspective

Historical context. The ODE perspective on optimization dates to Cauchy (1847). Its application to neural network training was formalized by the Neural Tangent Kernel (Jacot et al., 2018), which showed that sufficiently wide networks behave like linear models in their tangent-space linearization. The gradient flow framework provides the mathematical scaffolding for NTK theory.

Modern relevance:

• NTK theory uses exactly the framework developed here, replacing the linear model with a first-order Taylor expansion of a neural network around its initialization.
• Lazy training regime refers to networks where parameters stay close to initialization and the tangent kernel remains approximately constant — the regime where gradient flow analysis is accurate.
• Feature learning in finite-width networks departs from the NTK regime precisely because $\mathbf{K}$ evolves during training.
• Implicit regularization of gradient descent (converging to minimum-norm solutions) is studied extensively through the gradient flow lens.
• Scaling laws can be partially understood through the spectral structure of the tangent kernel — eigenvalues determine which "features" are learned at what rate.

Connection to subsequent parts:

• Part II replaces the linear model with a neural network → NTK derivation.
• Part III: Bayesian approach → posterior over parameters → Bayesian Linear Regression.
• Part IV: Function-space view → Gaussian Processes.

Summary of Takeaways

• Gradient flow is gradient descent with infinitesimal step size: $d\boldsymbol\phi/dt = -\nabla L$.
• Residual evolution follows $d\mathbf{r}/dt = -\mathbf{K}\mathbf{r}$ where $\mathbf{K} = \mathbf{X}^T\mathbf{X}$ is the tangent kernel.
• Closed-form solution: $\mathbf{f}(t) = \mathbf{y} + \exp[-\mathbf{K}t](\mathbf{f}(0) - \mathbf{y})$.
• Trainability requires $\mathbf{K}$ full rank (more parameters than data).
• Convergence speed is governed by $\lambda_{\min}(\mathbf{K})$.
• Parameter evolution converges to OLS as $t \to \infty$; in the underdetermined case, to the minimum-norm solution.
• Predictions at test points can be computed in closed form and have kernel regression structure.
• Gaussian prior over initialization → Gaussian distribution over predictions at every time $t$.
• This is the linear-model prototype of the Neural Tangent Kernel framework.

Exercises

Exercise 1. Verify directly that $L[\boldsymbol\phi_t]$ is monotonically decreasing under gradient flow by computing $dL/dt$.

Exercise 2. For a $2 \times 2$ matrix $\mathbf{K}$ with eigenvalues $\lambda_1 = 4, \lambda_2 = 1$, compute $\exp[-\mathbf{K}t]$ explicitly and sketch the residual components over time.

Exercise 3. Suppose $I = 2$ data points in $D = 3$ dimensions. What is the rank of $\mathbf{K} = \mathbf{X}^T\mathbf{X}$? Can the model fit both points exactly? How many solutions exist?

Exercise 4. Derive the gradient flow ODE for the regularized loss $L_\lambda = \frac{1}{2}\lVert \mathbf{X}^T\boldsymbol\phi - \mathbf{y}\rVert^2 + \frac{\lambda}{2}\lVert\boldsymbol\phi\rVert^2$. What is the equilibrium ($d\boldsymbol\phi/dt = 0$)?

Exercise 5. Show that in the underdetermined case ($I < D$), gradient flow starting from $\boldsymbol\phi_0 = \mathbf{0}$ converges to $\boldsymbol\phi_\infty = \mathbf{X}(\mathbf{X}^T\mathbf{X})^{-1}\mathbf{y}$, which is the minimum-norm interpolant.

Exercise 6. For the kernel $\mathbf{K} = \mathbf{X}^T\mathbf{X}$, express the convergence time $T_\epsilon$ (time to reduce residual below $\epsilon$) in terms of $\lambda_{\min}(\mathbf{K})$ and $\lVert\mathbf{r}(0)\rVert$.

Exercise 7. Starting from equation (18), compute the mean and variance of the prediction $y^*_t$ under the prior $\boldsymbol\phi_0 \sim \mathcal{N}(\mathbf{0}, \sigma_p^2\mathbf{I})$. Express your answer in terms of $\mathbf{x}^*, \mathbf{X}, \mathbf{y}, \sigma_p^2, t$.

Exercise 8. Consider a neural network with one hidden layer: $f(\mathbf{x}) = \mathbf{w}^T\sigma(\mathbf{V}\mathbf{x})$. Write the tangent kernel $\mathbf{K}$ for this model and explain why it depends on the current parameters $\boldsymbol\phi = (\mathbf{w}, \mathbf{V})$. Under what conditions does it become approximately constant?

Exercise 9. Prove that $\lVert\boldsymbol\phi_t - \boldsymbol\phi_0\rVert$ is minimized among all solutions with zero training loss when $\boldsymbol\phi_0 = \mathbf{0}$ and the system is underdetermined.