Saved in:
| Main Authors: | , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2510.04606 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866915990193831936 |
|---|---|
| author | Galashov, Alexandre Da Costa, Nathaël Xu, Liyuan Hennig, Philipp Gretton, Arthur |
| author_facet | Galashov, Alexandre Da Costa, Nathaël Xu, Liyuan Hennig, Philipp Gretton, Arthur |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2510_04606 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Closed-Form Last Layer Optimization Galashov, Alexandre Da Costa, Nathaël Xu, Liyuan Hennig, Philipp Gretton, Arthur Machine Learning 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. |
| title | Closed-Form Last Layer Optimization |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2510.04606 |