← Writing

Why we proved convergence on the simplest network there is

The linear layer is the one component every neural network shares. Its loss is still nonconvex, and training it is still NP-complete. That is why it is the right place to look for a guarantee.

The linear layer is the core component every neural network shares. We proved why training it converges, and that analysis now informs how foundation models scale.

Why it matters

Training an AI model means repeatedly nudging it toward better answers. The math that guarantees those nudges arrive at the desired output is still at an early stage.

Gradient descent and learning rate tuning are foundational to AI engineering. Every fine-tuning run and every pre-training job depends on them. Yet the proof that gradient descent converges, for the loss functions it actually optimizes, remains incomplete. The constants controlling the learning rate can’t be computed beforehand, so they’re chosen by trial and error.

My PhD research studied this on linear neural networks: the same linear transformations that, composed with activations, build every transformer and MLP. Despite their simplicity, the optimization stays hard. These networks have been shown to share key training behavior with the nonlinear architectures used in production.

The simplest model predicted the behavior of the complex one.

How it works

A linear network seems like it should be easy to train. Even with identity activations, its loss stays nonconvex, and training it is NP-complete. The architecture is simple. The optimization stays complex.

Prior gradient flow analysis showed these networks preserve a balancedness invariant critical for convergence. We expected it to carry over to gradient descent. It didn’t. Discretization destroyed exact balancedness, and the approach that worked for prior results stopped working.

We developed an induction argument controlling weight deviation from balancedness under practical stepsizes. That led to convergence for almost all initializations, with stepsizes that do not decay exponentially with depth and whose constants are computable from the network and data.

For SGD, we proved one of the first iterate bounds for a concrete neural network loss and established almost sure convergence.

We then tested whether the findings extend: tanh networks at depths 2, 3, and 5 under the same conditions. Same convergence pattern. Same depth ordering.