# Closed-Form Last Layer Optimization

arXiv: https://arxiv.org/abs/2510.04606 · HTML: https://mlsanity.com/paper/2510.04606/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 28 references that are in the map.
- First posted on arXiv 6 Oct 2025.

## Builds on (most cited first, 20 of 20)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Large Batch Training of Convolutional Networks](https://mlsanity.com/paper/1708.03888/) (2017)
- [High-Performance Large-Scale Image Recognition Without Normalization](https://mlsanity.com/paper/2102.06171/) (2021)
- [Evaluation of Neural Architectures Trained with Square Loss vs Cross-Entropy in Classification Tasks](https://mlsanity.com/paper/2006.07322/) (2020)
- [Variational Bayesian Last Layers](https://mlsanity.com/paper/2404.11599/) (2024)
- [Learning Deep Features in Instrumental Variable Regression](https://mlsanity.com/paper/2010.07154/) (2020)
- [Learning time-scales in two-layers neural networks](https://mlsanity.com/paper/2303.00055/) (2023)
- [An Introduction to Bi-level Optimization: Foundations and Applications in Signal Processing and Machine Learning](https://mlsanity.com/paper/2308.00788/) (2023)
- [Deep Proxy Causal Learning and its Application to Confounded Bandit Policy Evaluation](https://mlsanity.com/paper/2106.03907/) (2021)
- [On Instrumental Variable Regression for Deep Offline Policy Evaluation](https://mlsanity.com/paper/2105.10148/) (2021)
- [Functional Bilevel Optimization for Machine Learning](https://mlsanity.com/paper/2403.20233/) (2024)
- [Layer-Specific Adaptive Learning Rates for Deep Networks](https://mlsanity.com/paper/1510.04609/) (2015)
- [Leveraging the two timescale regime to demonstrate convergence of neural networks](https://mlsanity.com/paper/2304.09576/) (2023)
- [Kalman Filter for Online Classification of Non-Stationary Data](https://mlsanity.com/paper/2306.08448/) (2023)
- [Wide Neural Networks as Gaussian Processes: Lessons from Deep Equilibrium Models](https://mlsanity.com/paper/2310.10767/) (2023)
- [Bayesian Optimization via Continual Variational Last Layer Training](https://mlsanity.com/paper/2412.09477/) (2024)
- [Mean-Field Langevin Dynamics for Signed Measures via a Bilevel Approach](https://mlsanity.com/paper/2406.17054/) (2024)
- [Ultra-fast feature learning for the training of two-layer neural networks in the two-timescale regime](https://mlsanity.com/paper/2504.18208/) (2025)

## Built on by (most cited first, 5 of 5)

- [Rethinking Neural Network Learning Rates: A Stackelberg Perspective](https://mlsanity.com/paper/2605.15530/) (2026)
- [Doubly Robust Proxy Causal Learning with Neural Mean Embeddings](https://mlsanity.com/paper/2605.09514/) (2026)
- [PEIRA: Learning Predictive Encoders through Inter-View Regressor Alignment](https://mlsanity.com/paper/2605.17671/) (2026)
- [Training Under Challenge: Executable Certificates and Challenge-Closed Optimality for Neural Networks](https://mlsanity.com/paper/2608.12655/) (2026)
- [Teacher Geometry Shapes Learnability in Teacher-Student Networks](https://mlsanity.com/paper/2609.09595/) (2026)

## Abstract (arXiv)

Neural networks are typically optimized with variants of stochastic gradient descent. Under a squared loss, however, the optimal solution to the linear last layer weights is known in closed-form. We propose to leverage this during optimization, treating the last layer as a function of the backbone parameters, and optimizing solely for these parameters. We show this is equivalent to alternating between gradient descent steps on the backbone and closed-form updates on the last layer. We adapt the method for the setting of stochastic gradient descent, by trading off the loss on the current batch against the accumulated information from previous batches. We provide theoretical analyses showing convergence of the method to an optimal solution in the neural tangent kernel regime, as well as quantifying the gains compared to standard SGD in a one-step analysis. Finally, we demonstrate the effectiveness of our approach compared with SGD and Adam on a squared loss in several regression tasks, including neural operators and causal inference.