Closed-Form Last Layer Optimization

2025-10-06

AI Summary

Standard neural network training suffers from inefficient optimization of the final-layer linear weights. Method: Under squared loss, this work explicitly leverages the closed-form optimal solution of the last-layer weights with respect to backbone parameters, treating them as an analytic function of the backbone and performing gradient updates solely on the backbone. This approach strictly embeds the closed-form solution into training—equivalent to alternating optimization between backbone and last layer—and establishes convergence within the Neural Tangent Kernel (NTK) framework. Contribution/Results: By integrating the analytic solution with stochastic gradient descent (SGD), the method dynamically balances per-batch loss minimization against historical information. It significantly outperforms standard SGD on diverse tasks—including Fourier neural operators and instrumental variable regression—while remaining applicable to both regression and classification.

Technology Category

• Machine Learning: Optimization
• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data

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.

Problem

• Optimizing neural network last layer weights with closed-form solutions
• Replacing gradient descent with analytical updates for final layer
• Improving convergence and performance in supervised learning tasks

Innovation

• Closed-form optimization for last layer weights
• Treating last layer as backbone parameter function
• Convergence guaranteed in Neural Tangent Kernel regime