Technical Summary: Closed-Form Last Layer Optimization

Problem Statement

Deep neural networks are typically optimized using Stochastic Gradient Descent (SGD) or its variants, treating all layers uniformly. However, under a squared loss function, the optimal weights for the linear last layer of a neural network admit a closed-form solution (specifically, a ridge regression solution) given the features produced by the backbone.

Standard SGD treats the last layer as a set of parameters updated via gradient steps, ignoring this analytical optimality. This creates a mismatch: while the backbone parameters ($\theta$) require iterative optimization, the last layer weights ($W$) could theoretically be computed directly. Furthermore, in stochastic settings with mini-batches, naively computing the closed-form solution for each batch leads to severe overfitting to that specific batch, causing instability. Existing approaches that attempt to leverage this (e.g., in causal inference) often require backpropagation through the closed-form inverse, which is computationally expensive, or decouple the last layer from the backbone, leading to optimization instabilities.

Methodology

The authors propose a training procedure that treats the last layer as an implicit, deterministic function of the backbone parameters, optimizing solely for the backbone while ensuring the last layer remains optimal at every step.

1. Theoretical Foundation (Deterministic Case)

For a fixed backbone parameter $\theta$, the optimal last layer $W^\star(\theta)$ minimizes the regularized squared loss:
$$W^\star(\theta) = \arg\min_W \sum \|y_i - W\phi_\theta(x_i)\|^2 + \beta\|W\|_F^2$$
The authors leverage the Envelope Theorem to show that the gradient of the composite loss $L^\star(\theta) = L(W^\star(\theta), \theta)$ with respect to $\theta$ is identical to the gradient of the original loss where $W$ is fixed at $W^\star(\theta)$:
$$\nabla_\theta L^\star(\theta) = \nabla_\theta L(W^\star(\theta), \theta)$$
This critical result implies that one does not need to backpropagate through the matrix inversion required to compute $W^\star(\theta)$. Instead, the algorithm alternates between:

1. Computing $W_{t+1} = W^\star(\theta_t)$ (closed-form update).
2. Updating $\theta_{t+1} = \theta_t - \alpha \nabla_\theta L(W_{t+1}, \theta_t)$ (standard backprop through the backbone).

2. Practical Adaptation (Stochastic Mini-Batch)

In the stochastic setting, computing $W^\star$ on a single mini-batch causes overfitting. To address this, the authors introduce a proximal regularizer that penalizes the distance between the current closed-form solution and the previous estimate $W_t$. The objective becomes:
$$L_{B_t, W_t}^{\text{prox}}(W, \theta) = \sum_{(x,y) \in B_t} \|y - W\phi_\theta(x)\|^2 + \lambda \|W - W_t\|_F^2$$
The closed-form solution for this proximal loss is:
$$W^\star_{B_t, W_t}(\theta) = (Y_t \phi_\theta(X_t)^\top + \lambda W_t)(\phi_\theta(X_t)\phi_\theta(X_t)^\top + \lambda I)^{-1}$$
This formulation effectively implements an approximate Kalman filter on the last layer, accumulating information from previous batches while adapting to the current batch. The algorithm updates the backbone via SGD and then re-fits the last layer using this proximal closed-form solution.

3. Classification Extension

While primarily designed for regression, the method is extended to classification by treating one-hot labels as targets for a squared loss. The output is interpreted via an $\arg\max$ operation. The authors note that while this heuristic works for small-to-moderate class counts, it does not guarantee probability normalization.

Key Contributions

1. Optimization Framework: A practical algorithm that integrates closed-form last layer updates into standard SGD pipelines without requiring backpropagation through the inverse matrix operation, leveraging the envelope theorem.
2. Stochastic Regularization: A proximal formulation that prevents overfitting in mini-batch settings by regularizing the last layer towards its previous estimate, interpreted as approximate Kalman filtering.
3. Theoretical Analysis:
   • NTK Regime: Proven convergence to a global minimum in the Neural Tangent Kernel (NTK) infinite-width limit, despite the non-convexity of the loss in function space.
   • One-Step Analysis: A theoretical bound quantifying the superiority of the closed-form step over a gradient descent step, showing that the gain is related to the condition number (variance of singular values) of the feature matrix.
4. Empirical Validation: Demonstrated improvements over standard SGD and naive closed-form ridge regression across diverse tasks.

Experimental Results

The method ("$\ell_2$ c.f. proximal ($\lambda$)") was evaluated against standard squared loss SGD/Adam and a naive closed-form ridge baseline ("$\ell_2$ c.f. ridge ($\beta$)").

• Quantum Chemistry (QM9): The proximal method consistently outperformed standard SGD, particularly at small batch sizes where the naive ridge baseline failed to converge. It showed robustness to the choice of backbone optimizer (SGD vs. Adam).
• Fourier Neural Operators (FNO): In operator learning for the viscous Burgers equation, the proximal method remained effective across batch sizes, whereas the naive ridge baseline failed at small batches. The method effectively handled the high-dimensional spatial degrees of freedom.
• Deep Feature Instrumental Variable (DFIV): In causal inference, the method eliminated the need for costly re-estimation of last layers on the full dataset after training, achieving performance comparable to re-estimation strategies while being computationally efficient.
• Classification (CIFAR-100, ImageNet): On CIFAR-100, the method achieved accuracy comparable to or better than Cross Entropy for small-to-moderate batch sizes. However, on ImageNet (large number of classes), it underperformed Cross Entropy, highlighting a limitation of squared loss for high-cardinality classification.

Significance and Claims

The paper claims that leveraging the closed-form solution for the last layer is a viable and effective strategy to accelerate optimization and improve stability, particularly in regression tasks and small-batch regimes.

• Theoretical Significance: The work provides the first analysis of the critical points of the modified loss in function space and proves convergence to a global minimum in the NTK regime, addressing the non-convexity inherent in the problem.
• Practical Significance: The proximal formulation solves the overfitting issue inherent in mini-batch closed-form updates, making the method robust and compatible with standard training pipelines.
• Limitations: The authors explicitly state that the method is currently limited to squared loss. While it performs well on classification with few classes, it underperforms Cross Entropy on large-scale classification (ImageNet), suggesting that adapting the approach to cross-entropy loss is a necessary direction for future work. The method also incurs a computational overhead (matrix inversion) that scales with the last layer dimension, making it slower than standard SGD for very large last layers in small-batch regimes, though this cost diminishes as the backbone size increases.