Last-Layer Linearization in Neural Networks

• Last-layer linearization is the process of freezing a network’s nonlinear feature extractor and using a linear output head to simplify inference and retraining.
• It underpins methods such as Bayesian GLMs, posterior-aware conformal regression, and closed-form head optimization, offering significant efficiency and effective uncertainty modeling.
• While effective for many applications, its limitations in capturing extractor uncertainty spur the development of multi-layer fusion techniques for enhanced robustness.

Last-layer linearization is the treatment of a deep network as a fixed nonlinear feature map followed by a linear output head, so that inference, optimization, or retraining is performed only with respect to that head. In a standard feed-forward regression model, the backbone produces a feature map $\phi(x)\in\mathbb{R}^d$ from the last hidden layer and the predictor is $f_w(x)=\phi(x)^\top w$; more generally, linearizing a DNN $f_\theta(x)$ around a selected parameter subset $\theta_s$ gives $$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$, and choosing $\theta_s$ to be the final layer yields the last-layer form (Kim et al., 1 Dec 2025, Wilson et al., 29 May 2026). Recent work uses this construction in several distinct but related ways: as a Bayesian generalized linear model for epistemic uncertainty quantification, as the basis of posterior-aware conformal regression, as a device for closed-form or proximal optimization of the head, and as a retraining mechanism for robustness to spurious correlations and worst-group failure modes (Fiedler et al., 2023, Galashov et al., 6 Oct 2025, Welfert et al., 2024).

1. Formalization and model classes

In the Bayesian GLM view, both full-network and last-layer approximations are obtained by linearizing a DNN around selected parameters and placing a Gaussian prior on those linearized parameters. If $\theta_s\sim\mathcal N(0,\lambda^{-1}I)$, the resulting linearized model is a Bayesian generalized linear model whose predictive distribution matches that of a Gaussian process with kernel $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$. When the subset is all parameters, the induced kernel is the Neural Tangent Kernel $$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$; when the subset is only the final layer parameters $\theta_L$, $f_w(x)=\phi(x)^\top w$0 and the induced kernel is the Conjugate Kernel $f_w(x)=\phi(x)^\top w$1. The corresponding terminologies in the literature are DNN-GLM for full-network linearization and LL-GLM for last-layer linearization (Wilson et al., 29 May 2026).

A closely related formulation appears in Bayesian last-layer neural networks. There, the hidden layers are deterministic and define a learned feature vector $f_w(x)=\phi(x)^\top w$2, while the output layer is linear and probabilistic: $f_w(x)=\phi(x)^\top w$3, with Gaussian observation noise and a Gaussian prior on $f_w(x)=\phi(x)^\top w$4. In this setting, the model is exactly Bayesian linear regression on learned features, rather than a first-order approximation over all network weights (Fiedler et al., 2023).

This dual usage explains why the term “last-layer linearization” covers more than one research line. In some papers it means a first-order linearization with respect to final-layer parameters; in others it means freezing a representation and exploiting the exact linear structure of the output head. The common object is the same: a nonlinear backbone whose role is feature extraction, and a final affine map whose linearity is used analytically or algorithmically.

2. Bayesian last-layer models and uncertainty quantification

Several leading methods for uncertainty quantification linearize DNNs to form Bayesian GLMs, with the posterior predictive distribution as the key object for epistemic uncertainty. Under the Gaussian-process equivalence, the predictive posterior is Gaussian, and the posterior covariance quantifies uncertainty. In Bayesian last-layer models, this becomes especially explicit: the predictive distribution is $f_w(x)=\phi(x)^\top w$5, where $f_w(x)=\phi(x)^\top w$6 and $f_w(x)=\phi(x)^\top w$7, with posterior precision $f_w(x)=\phi(x)^\top w$8 (Fiedler et al., 2023).

A central comparison in this literature is whether one gains materially by linearizing the entire network rather than only the final layer. Using Bayes Free Energy, $f_w(x)=\phi(x)^\top w$9, as the comparison metric, random matrix theory in a double-asymptotic regime yields the conclusion that there is no meaningful theoretical advantage of full-network linearization over last-layer linearization in UQ; the only regime where NTK can outperform CK is when $f_\theta(x)$0, and even there the paper suggests the advantage may be an initialization artifact. The same study reports that LL-GLM is comparable to, or sometimes better than, DNN-GLM on toy 1D regression, UCI regression benchmarks, image classification tasks from MNIST through ImageNet, and GPT-2 fine-tuned on IMDB sentiment classification. On GPT-2/IMDB, LL-GLM showed about $f_\theta(x)$1 faster wall-time and about $f_\theta(x)$2 lower memory use, while having no statistically meaningful difference in VARROC-ID (Wilson et al., 29 May 2026).

The Bayesian last-layer literature also studies how predictive covariance relates to extrapolation. One paper defines an affine cost $f_\theta(x)$3 through a constrained optimization over the training features and proves that, when $f_\theta(x)$4, $f_\theta(x)$5. This identifies the Bayesian last-layer covariance with a scaled extrapolation cost in feature space and motivates tuning $f_\theta(x)$6 after marginal-likelihood training. In the reported simulation study, the model with tuned $f_\theta(x)$7 achieved the best test log-predictive density among the compared methods: $f_\theta(x)$8 for NN with BLL, versus $f_\theta(x)$9 for BLR with NN features and $\theta_s$0 for a BNN trained with variational inference (Fiedler et al., 2023).

Standard Gaussian BLL models have also been generalized. A flexible Bayesian last-layer model, DVI-IBLL, replaces the Gaussian prior with an implicit prior $\theta_s$1 and uses diffusion posterior sampling rather than Gaussian or mean-field posterior approximation. The method is presented as addressing non-Gaussian, outlier-rich, or high-dimensional settings while keeping the last-layer-only Bayesian structure. Reported runtime overhead remains modest relative to deterministic training, with 0.343 s/batch for DVI-IBLL versus 0.321 s/batch for the DNN baseline, and the paper reports CIFAR-10 accuracy $\theta_s$2, ECE $\theta_s$3, and NLL $\theta_s$4 (Xu et al., 2024).

3. Posterior-aware conformal regression

A recent development is the combination of last-layer Bayesianization with conformal prediction. CLAPS defines a conformal regression method that uses a Last-Layer Laplace Approximation together with split-conformal calibration. The backbone is kept fixed, and only the scalar linear head $\theta_s$5 is Bayesianized. With a Gaussian prior $\theta_s$6 and homoscedastic Gaussian observation noise $\theta_s$7, the last-layer MAP estimate is the ridge-regression solution

$\theta_s$8

and the Laplace approximation gives

$\theta_s$9

The induced posterior predictive is

$$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$0

with $$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$1 and $$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$2, decomposing into aleatoric and epistemic terms (Kim et al., 1 Dec 2025).

CLAPS uses the full posterior predictive CDF rather than a residual score. If $$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$3 denotes the CDF of the Gaussian predictive, the conformity score is

$$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$4

For the Gaussian case, with $$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$5, this becomes $$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$6. Split conformal calibration is then applied with the standard rank/order-statistic rule, and because $$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$7 is monotone in $$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$8 under the Gaussian predictive, the prediction set is a central interval symmetric around $$f_{\hat\theta}(\theta_s,x)=f_{\hat\theta}(x)+J_s(\theta_s,x)(\theta_s-\hat\theta_s)$$9. The finite-sample marginal coverage guarantee is unchanged: under exchangeability of calibration and test points conditional on the trained backbone/posterior, $\theta_s$0 (Kim et al., 1 Dec 2025).

The method is accompanied by a diagnostic suite. The predictive variance decomposition $\theta_s$1 motivates the epistemic share

$\theta_s$2

together with $\theta_s$3 as a global posterior-uncertainty summary and the heteroscedasticity signal $\theta_s$4. The reported empirical pattern is that CLAPS is most effective when $\theta_s$5 is nontrivial and heteroscedasticity is weak. On Airfoil and kin8nm, it achieved nominal coverage around $\theta_s$6 with the narrowest intervals among the compared methods, with reported widths about $\theta_s$7 on Airfoil and $\theta_s$8 on kin8nm. On CASP it was competitive but not always the narrowest, and on YearPredictionMSD the advantage largely disappeared as posterior variance collapsed and Normalized-CP or CQR became tighter (Kim et al., 1 Dec 2025).

4. Exact last-layer optimization and stochastic training

Last-layer linear structure is also exploited for optimization under squared loss. For a model $\theta_s$9, with backbone feature map $\theta_s\sim\mathcal N(0,\lambda^{-1}I)$0 and final linear layer $\theta_s\sim\mathcal N(0,\lambda^{-1}I)$1, the regularized squared-loss objective

$\theta_s\sim\mathcal N(0,\lambda^{-1}I)$2

is a ridge-regression problem in $\theta_s\sim\mathcal N(0,\lambda^{-1}I)$3 for fixed $\theta_s\sim\mathcal N(0,\lambda^{-1}I)$4. The optimal final layer is available in closed form: $\theta_s\sim\mathcal N(0,\lambda^{-1}I)$5 This induces a reduced objective $\theta_s\sim\mathcal N(0,\lambda^{-1}I)$6 in which the last layer is treated as an implicit function of the backbone parameters (Galashov et al., 6 Oct 2025).

The key theorem is a gradient-equivalence result: for fixed $\theta_s\sim\mathcal N(0,\lambda^{-1}I)$7, $\theta_s\sim\mathcal N(0,\lambda^{-1}I)$8. Because $\theta_s\sim\mathcal N(0,\lambda^{-1}I)$9 at the optimum, one does not need to backpropagate through the matrix inverse defining $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$0. Training can therefore alternate an exact solve for the head with an ordinary gradient step on the backbone: $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$1 A stochastic minibatch extension adds a proximal term $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$2, yielding a closed-form proximal update that regularizes the batch solve toward the previous iterate (Galashov et al., 6 Oct 2025).

The same work shows that the reduced loss is not convex and may have non-global critical points, but proves an NTK-regime convergence result: under sequential infinite-width limits, positive definite NTK and NGPK, and $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$3, gradient flow on the reduced objective converges almost surely to a global minimizer. Empirically, the proximal closed-form last-layer method improves over squared-loss SGD on 1D Burgers with a Fourier Neural Operator, improves minibatch training in deep feature instrumental variable regression, generally outperforms plain squared-loss SGD on CIFAR-10 and CIFAR-100 with ResNet-18, and improves over squared-loss SGD on ImageNet with NF-Nets, although cross-entropy still remains stronger at ImageNet scale (Galashov et al., 6 Oct 2025).

5. Retraining, spurious correlations, and worst-group accuracy

Another major use of last-layer linearization is last-layer retraining on frozen representations. In the group-robustness setting, the ambient input $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$4 is mapped to a latent representation $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$5, and only a linear classifier $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$6 is retrained. Under Gaussian group-conditional latent distributions $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$7, recent theory analyzes downsampling (DS), upweighting (UW), and intra-class domain mixup (MU). The paper proves that DS and UW are distribution-free equivalent for any loss, derives closed-form optimal linear classifiers for DS/UW, SRM/ERM, and MU under the Gaussian model, and shows that under the orthogonality condition $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$8 and $k(x,x')=J_s(\theta_s,x)J_s(\theta_s,x')^\top$9,

$$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$0

The common optimal worst-group error for the balanced methods is $$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$1. On CMNIST, CelebA, and Waterbirds, the reported worst-group error means $$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$2 std over 10 runs support the theory, with DS, UW, and MU all substantially better than ERM (Welfert et al., 2024).

Data selection for last-layer retraining has also been automated. One framework uses OWL-ViT to score the presence of the target object, ranks images within each class by that detector score, and retrains only the final linear classifier head on class-balanced subsets such as top $$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$3, mid $$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$4, bot $$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$5, and rnd $$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$6. On ImageNet-1k with ResNet-50 and DeiT-small, the reported trends are that top 100 models perform better on low-ranked test examples, top 100 models are consistently more robust than bot 100 under foreground and background noise, and bot $$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$7 models achieve higher accuracy on ImageNet-A (Azizmalayeri et al., 2023).

The same line of work also exposes a limitation of last-layer-only retraining. H2T-DFR combines DFR with Head2Toe feature selection from all layers, using group-lasso regularization on a balanced validation set to select useful features beyond the last layer. With a ResNet-50 backbone, the reported mean worst-group accuracy over 5 seeds improved from $$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$8 for DFR to $$K_{\mathrm{NTK}}(x_i,x_j)=J(\theta,x_i)J(\theta,x_j)^\top$$9 for H2T-DFR on CelebA, and from $\theta_L$0 to $\theta_L$1 on HAM10000, while Waterbirds did not improve and remained close to DFR (Hameed et al., 2024).

6. Critiques of last-layer sufficiency and generalized last-layer heads

The assumption that the last layer is a sufficient representation is explicitly challenged in several domains. In sentence embedding, one study argues that sentence meaning is not concentrated only in the final transformer layer. It reports that on STS-B, BERT with $\theta_L$2 pooling can perform much better at an intermediate layer than at the last layer, with the 4th layer around $\theta_L$3 better than the last layer for STS-B. The proposed LayerAttPooler uses attention-based pooling over all layers, concatenates the last-layer $\theta_L$4 representation with the learned layer-pooled vector, and improves over last-layer baselines on STS and semantic search. For example, in an unsupervised BERT$\theta_L$5 reimplementation, the average score increases from $\theta_L$6 with last-layer $\theta_L$7 to $\theta_L$8 with LayerAttPooler (Oh et al., 2022).

A similar critique appears in visual tokenization. DRoRAE argues that last-layer features are a lossy bottleneck for image tokenization because low-level visual details survive in the last layer only as attenuated residuals after semantic abstraction. The method fuses all encoder layers through energy-constrained routing and incremental correction rather than using only the last encoder layer. On ImageNet-256, the reported reconstruction rFID decreases from $\theta_L$9 for the RAE baseline to $f_w(x)=\phi(x)^\top w$00 for the full three-phase DRoRAE model, and generation FID improves from $f_w(x)=\phi(x)^\top w$01 to $f_w(x)=\phi(x)^\top w$02 with AutoGuidance. The same work reports a log-linear scaling law between fusion capacity and reconstruction quality, $f_w(x)=\phi(x)^\top w$03 with $f_w(x)=\phi(x)^\top w$04 (Zhu et al., 11 May 2026).

In neural operators, the “last layer” has also been generalized away from linear regression. The diffusion last layer (DLL) is explicitly described as not being classical last-layer linearization in the Bayesian-NN sense, yet it preserves the same decomposition into frozen backbone and lightweight probabilistic head. The backbone produces an input-dependent rank-$f_w(x)=\phi(x)^\top w$05 basis $f_w(x)=\phi(x)^\top w$06, the output field is approximated as $f_w(x)=\phi(x)^\top w$07, and a conditional diffusion model is trained on the coefficient vector $f_w(x)=\phi(x)^\top w$08. The paper therefore characterizes DLL as a last-layer generative approximation or function-space probabilistic last layer rather than a parameter-space linearization (Park et al., 4 Feb 2026).

7. Geometry, recoverability, and recurring limitations

The linearity of the last layer has geometric consequences beyond training and uncertainty estimation. In a model-stealing setting where a transformer returns logits $f_w(x)=\phi(x)^\top w$09, the attainable logits lie in the linear subspace $f_w(x)=\phi(x)^\top w$10. With normalization before the last layer, the outputs satisfy linear constraints together with a quadratic equation $f_w(x)=\phi(x)^\top w$11, so the output set is an ellipsoid inside $f_w(x)=\phi(x)^\top w$12. Under a clean rank-$f_w(x)=\phi(x)^\top w$13 spectral gap, a nondegenerate quadric, and local stability, the last layer can be recovered exactly up to an orthogonal gauge $f_w(x)=\phi(x)^\top w$14, $f_w(x)=\phi(x)^\top w$15. The same analysis shows an identifiability wall below the last layer: deeper hidden networks cannot be fully reverse engineered from final outputs alone (Khilar, 5 Jun 2026).

The literature is equally explicit about the limitations of last-layer linearization. In CLAPS, only the last linear layer is Bayesianized, so epistemic uncertainty from the feature extractor itself is not captured, and the method loses its advantage when the last-layer posterior contracts toward zero or when heteroscedasticity is strong (Kim et al., 1 Dec 2025). In the LL-GLM versus DNN-GLM comparison, the theory is for randomly initialized networks rather than trained DNNs on realistic datasets, the sampling framework models only epistemic uncertainty, and near-rank-deficiency in trained last-layer Jacobians can create finite-precision issues (Wilson et al., 29 May 2026). In robustness and representation learning, several papers argue that the penultimate representation may already have discarded useful robust, low-level, or task-relevant signals, motivating all-layer feature selection or multi-layer fusion instead of strict last-layer dependence (Hameed et al., 2024, Oh et al., 2022, Zhu et al., 11 May 2026).

A central interpretation emerging from these results is that last-layer linearization is most effective when the decisive uncertainty or task signal lies in how learned features are used, rather than in uncertainty about the feature extractor itself. When that condition holds, the approach offers analytic tractability, exact or approximate Gaussian posteriors, conformal calibration with standard validity, efficient closed-form optimization, and strong computational savings. When it does not hold, the literature increasingly replaces last-layer sufficiency with multi-layer or function-space constructions that preserve the computational advantages of a lightweight head while relaxing the representational bottleneck (Wilson et al., 29 May 2026, Park et al., 4 Feb 2026).