Closed-Form Last Layer Optimization

• URL: http://arxiv.org/abs/2510.04606v1
• Date: Mon, 06 Oct 2025 09:14:39 GMT
• Title: Closed-Form Last Layer Optimization
• Authors: Alexandre Galashov, Nathaël Da Costa, Liyuan Xu, Philipp Hennig, Arthur Gretton,
• Abstract summary: Under a squared loss, the optimal solution to the linear last layer weights is known in closed-form.<n>We show this is equivalent to alternating between gradient descent steps on the backbone and closed-form updates on the last layer.
• Score: 72.49151473937319
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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. Further, we prove that, in the Neural Tangent Kernel regime, convergence of this method to an optimal solution is guaranteed. Finally, we demonstrate the effectiveness of our approach compared with standard SGD on a squared loss in several supervised tasks -- both regression and classification -- including Fourier Neural Operators and Instrumental Variable Regression.

Related papers

• Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning [17.65459083031186]
  We train dynamics of gradient descent in a softmax self-attention layer trained to perform linear regression.<n>We show that a simple first-order gradient descent can converge to the globally optimal self-attention parameters.
  arXiv  Detail & Related papers  (2026-03-02T06:44:54Z)
• Safeguarded Stochastic Polyak Step Sizes for Non-smooth Optimization: Robust Performance Without Small (Sub)Gradients [16.39606116102731]
  The vanishing Polyak delivering adaptive neural network has proven to be a promising choice for gradient descent (SGD)<n> Comprehensive experiments on deep networks corroborate tight convex network theory.<n>In this work, we provide rigorous convergence guarantees for non-smooth optimization with no need for strong assumptions.
  arXiv  Detail & Related papers  (2025-12-02T02:24:32Z)
• Flattened one-bit stochastic gradient descent: compressed distributed optimization with controlled variance [55.01966743652196]
  We propose a novel algorithm for distributed gradient descent (SGD) with compressed gradient communication in the parameter-server framework. Our gradient compression technique, named flattened one-bit gradient descent (FO-SGD), relies on two simple algorithmic ideas.
  arXiv  Detail & Related papers  (2024-05-17T21:17:27Z)
• Directional Smoothness and Gradient Methods: Convergence and Adaptivity [16.779513676120096]
  We develop new sub-optimality bounds for gradient descent that depend on the conditioning of the objective along the path of optimization.<n>Key to our proofs is directional smoothness, a measure of gradient variation that we use to develop upper-bounds on the objective.<n>We prove that the Polyak step-size and normalized GD obtain fast, path-dependent rates despite using no knowledge of the directional smoothness.
  arXiv  Detail & Related papers  (2024-03-06T22:24:05Z)
• Stochastic Gradient Descent for Gaussian Processes Done Right [86.83678041846971]
  We show that when emphdone right -- by which we mean using specific insights from optimisation and kernel communities -- gradient descent is highly effective. We introduce a emphstochastic dual descent algorithm, explain its design in an intuitive manner and illustrate the design choices. Our method places Gaussian process regression on par with state-of-the-art graph neural networks for molecular binding affinity prediction.
  arXiv  Detail & Related papers  (2023-10-31T16:15:13Z)
• Sampling from Gaussian Process Posteriors using Stochastic Gradient Descent [43.097493761380186]
  gradient algorithms are an efficient method of approximately solving linear systems. We show that gradient descent produces accurate predictions, even in cases where it does not converge quickly to the optimum. Experimentally, gradient descent achieves state-of-the-art performance on sufficiently large-scale or ill-conditioned regression tasks.
  arXiv  Detail & Related papers  (2023-06-20T15:07:37Z)
• Improved Overparametrization Bounds for Global Convergence of Stochastic Gradient Descent for Shallow Neural Networks [1.14219428942199]
  We study the overparametrization bounds required for the global convergence of gradient descent algorithm for a class of one hidden layer feed-forward neural networks.
  arXiv  Detail & Related papers  (2022-01-28T11:30:06Z)
• STORM+: Fully Adaptive SGD with Momentum for Nonconvex Optimization [74.1615979057429]
  We investigate non-batch optimization problems where the objective is an expectation over smooth loss functions. Our work builds on the STORM algorithm, in conjunction with a novel approach to adaptively set the learning rate and momentum parameters.
  arXiv  Detail & Related papers  (2021-11-01T15:43:36Z)
• Differentiable Annealed Importance Sampling and the Perils of Gradient Noise [68.44523807580438]
  Annealed importance sampling (AIS) and related algorithms are highly effective tools for marginal likelihood estimation. Differentiability is a desirable property as it would admit the possibility of optimizing marginal likelihood as an objective. We propose a differentiable algorithm by abandoning Metropolis-Hastings steps, which further unlocks mini-batch computation.
  arXiv  Detail & Related papers  (2021-07-21T17:10:14Z)
• Stochastic Optimization with Heavy-Tailed Noise via Accelerated Gradient Clipping [69.9674326582747]
  We propose a new accelerated first-order method called clipped-SSTM for smooth convex optimization with heavy-tailed distributed noise in gradients. We prove new complexity that outperform state-of-the-art results in this case. We derive the first non-trivial high-probability complexity bounds for SGD with clipping without light-tails assumption on the noise.
  arXiv  Detail & Related papers  (2020-05-21T17:05:27Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.