AI Insight
This paper proposes a neural network optimization method that replaces gradient-based updates for the final layer with closed-form solutions, while still using gradient descent for earlier layers. The approach is adapted for stochastic gradient descent by balancing current batch information with previously accumulated data. Theoretical analysis demonstrates convergence guarantees in the neural tangent kernel regime, and empirical results show performance improvements over standard SGD and Adam on regression tasks including neural operators and causal inference problems.
Why it matters
This method could make neural network training more efficient and stable for regression problems by eliminating the need to tune learning rates for the final layer. The approach is particularly relevant for applications in scientific computing and causal inference where squared loss functions are commonly used.
Understand the Science
⚠️ Preprint – Noch nicht peer-reviewed
Dieser Artikel wurde noch nicht von unabhängigen Experten begutachtet. Die Ergebnisse sind vorläufig und sollten mit Vorsicht interpretiert werden.
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. 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.