Vision Transformer Finetuning Benefits from Non-Smooth Components

AI summary

Overview

Research area: Transfer learning and parameter-efficient finetuning of vision transformers, combined with the theory of smoothness and Lipschitz continuity in neural networks.

Technical level: Intermediate to Advanced. The empirical sections are accessible; Sections 4 and Appendix B require comfort with Lipschitz constants, spectral norms, and Frobenius-norm bounds.

Scope: The paper argues, theoretically and through roughly 1,000 finetuning runs on large-scale vision transformers, that the non-smooth ("plastic") components of a ViT—attention modules and feedforward layers—are the ones worth updating when adapting to downstream data.

What This Paper Is About

Smoothness in transformers has mostly been studied for generalization, training stability, and adversarial robustness, but its role in transfer learning is poorly understood. The authors introduce a measure they call plasticity—the average rate of change of a module's output with respect to its input—and use it to ask which transformer components should be prioritized during finetuning and why. Their answer runs against the common assumption that smoothness is desirable: high plasticity, meaning low smoothness, is what consistently helps finetuning.

Key Contributions

1. An intuitive, formal measure. They define the plasticity of a transformer component as its average rate of change over pairs of token sequences, $\mathscr{P}(f)=\mathbb{E}{(x,y)\sim\nu}\left[|f(x)-f(y)|{\mathrm{F}}/|x-y|_{\mathrm{F}}\right]$, which is always bounded above by the Lipschitz constant $\mathrm{Lip}(f)$.

2. A theoretical analysis. They derive upper bounds on plasticity for LayerNorms, feedforward linear layers, and multi-head self-attention, including a tighter attention bound based on the total energy of input images, and use these bounds to rank the components.

3. A validated plasticity ranking. They verify the predicted ordering on pretrained ViTs of three sizes (86M, 307M, and 632M parameters) across 11 classification benchmarks, and show the same pattern on DINOv3 and GPT2, giving the ranking MHA → FC1 → FC2 → LN2 → LN1.

4. Exhaustive finetuning evidence. They run component-wise finetuning of large-scale ViTs under both SGD and Adam, showing that high-plasticity components give higher and more stable accuracy across initializations and learning rates.

Main Findings

• Theory ranks attention as most plastic. For a LayerNorm, $\mathscr{P}(f) \leq \frac{1}{\sigma}|\gamma|{\infty}$; for a feedforward linear layer, $\mathscr{P}(f) \leq |W|{2}$; and for multi-head self-attention, $\mathscr{P}(f) \leq \sum_{h=1}^{H}|O^{h}|{2}|V^{h}|{2}\sqrt{3n+(12n+3)r^{4}|A^{h}|_{2}^{2}}$. The theoretical takeaway is MHA → FC1 ≈ FC2 → LN2 ≈ LN1.

• A tighter attention bound exists. Assuming input images have bounded total energy $\mathcal{E}$, the bound becomes $\sum_{h=1}^{H}|O^{h}|{2}|V^{h}|{2}\left(\sqrt{n}+\alpha^{2}\mathcal{E}|A^{h}|_{2}\right)$, improving the growth rate in sequence length.

• Empirical plasticity refines the ranking. Measured on pretrained ViTs, the ordering is MHA → FC1 → FC2 → LN2 → LN1, with attention and feedforward layers exceeding a plasticity of 1 (amplifying input differences) and LayerNorms falling below 1 (contracting them).

• Longer sequences increase attention plasticity. ViT-Huge, with sequence length $n = 257$, shows higher attention plasticity than ViT-Base, consistent with the sequence-length dependence in the attention upper bound.

• The pattern generalizes beyond supervised ViTs. DINOv3, a 7B-parameter self-supervised vision transformer, and GPT2, a 124M-parameter decoder-only language model, show the same component ordering.

• Plastic components finetune better. Averaged over the 11 benchmarks with SGD, top-1 accuracy is 90.8 for MHA, 90.7 for FC1, 90.3 for FC2, 89.9 for LN2, and 89.8 for LN1, with the MHA advantage statistically significant for all modules except FC1 according to a Wilcoxon signed-rank test at a 5% confidence level.

• Plastic components are more stable. Across a grid of learning rates and 3 seeds, the attention module shows the smallest variability, and the benefits are most pronounced on harder datasets such as Cifar100, Clipart, and Sketch.

• Plasticity aligns with gradient norms. During finetuning on Sketch at learning rate $\eta = 1\mathrm{e}{-}2$, the ordering of gradient norms matches the plasticity ranking, and the validation loss descends more steeply for high-plasticity components.

• Adam shows the same effect. With Adam, for example, Cifar100 gives MHA 91.0 ± 0.2, FC1 91.3 ± 0.6, FC2 90.6 ± 1.4, LN2 89.4 ± 2.7, and LN1 88.4 ± 3.2; Sketch gives MHA 69.0 ± 0.5 and LN2 63.1 ± 2.9.

• Single-component finetuning can rival LoRA. Finetuning only the LayerNorms with Adam reaches 89.4% on Cifar100 with 28K trainable parameters, versus 88.1% with 400K parameters for LoRA—15× fewer trainable parameters.

• Practical recommendation. The authors conclude that the components to prioritize are the attention module and the first feedforward linear layer.

Methodology in Plain English

The authors isolate one component type at a time rather than testing all combinations, which would be computationally prohibitive. They first write down a measure of how strongly a module's output changes when its input changes, then prove upper bounds on that measure for each module type using standard linear-algebra tools like spectral norms, plus assumptions that hold in practice (tokens have consistent per-position statistics; token sequences lie in a bounded ball; input images have bounded total energy). Next they estimate the measure empirically: they embed 12,800 pretraining images from ImageNet1k as the sequences $x$ and downstream images as the sequences $y$, then record the distribution of rate-of-change values for every module in the network. Finally, they finetune ViTs of 86M, 307M, and 632M parameters—pretrained on ImageNet22k—one component type at a time, over 11 benchmarks, 4 well-spaced learning rates, and 3 seeds, using both SGD and Adam, tracking accuracy, gradient norms, and validation loss.

Why This Matters

Impact on research. The paper challenges the prevailing assumption that promoting smoothness helps learning. Prior work linked Lipschitz regularization to better generalization, stability, and robustness, and showed that reducing the Lipschitz constant hurts a 145M-parameter Lipschitz-constrained transformer on FineWeb (where matching a NanoGPT baseline required a Lipschitz constant up to $10^{264}$). This paper turns that into an actionable rule for selecting which parameters to adapt, and offers a bridge between theoretical smoothness analysis and practical parameter-efficient finetuning.

Real-world applications (derived from the benchmark suite; the paper does not enumerate applications itself).

• Adapting vision models to shifted visual domains such as Clipart and Sketch imagery.
• Image classification under corruptions and conditions including Gaussian noise, motion blur, snow, contrast changes, and speckle noise.
• Fine-grained recognition tasks such as Flowers102 and Pets.
• Low-cost adaptation of large pretrained backbones when only a fraction of parameters can be updated.

Industry relevance. Selective and parameter-efficient finetuning is already standard practice, and the paper shows that the choice of which component to update matters more than raw parameter count. Its result that this choice is robust across learning rates and seeds is directly useful for teams that want fewer hyperparameter sweeps, and the finding that LayerNorm-only tuning can beat LoRA on parameter budget suggests cheaper alternatives to established PEFT methods.

Future Directions

• Why is FC1 more plastic than FC2? The empirical ranking separates the two feedforward layers, while the theory only places them as approximately equal; the cause is left open.
• What about combinations? The study finetunes each module type in isolation, and the authors explicitly note that considering all possible combinations is computationally prohibitive. Whether the ranking still holds when several components are updated together is untested.
• How does distillation or alignment shift the picture? Remark 5.2 notes that plasticity depends heavily on model weights, so distilled or instruct-version models may show different values.
• Beyond the tested settings. The analysis is validated on ViT-Base, ViT-Huge, ViT-Large (307M), a 632M ViT, DINOv3, and GPT2; other architectures, modalities, and larger scales remain unexamined.

Target Audience

Researchers and practitioners working on transfer learning, parameter-efficient finetuning, and the theoretical properties of transformers. It is especially relevant to engineers choosing which parameters to update when adapting large vision or language backbones, and to theorists interested in how smoothness, Lipschitz continuity, and gradient norms interact during adaptation.

Authors’ abstract

The smoothness of the transformer architecture has been extensively studied in the context of generalization, training stability, and adversarial robustness. However, its role in transfer learning remains poorly understood. In this paper, we analyze the ability of vision transformer components to adapt their outputs to changes in inputs, or, in other words, their \emph{plasticity}. Defined as an average rate of change, it captures the sensitivity to input perturbation; in particular, a high plasticity implies a low smoothness. Our theoretical analysis and extensive experiments -- over $1,000$ finetuning runs on large-scale vision transformers -- showcase that this perspective provides principled guidance in choosing the components to prioritize during adaptation. A key takeaway for practitioners is that the high plasticity of the attention modules and feedforward layers consistently leads to better finetuning performance. Our findings depart from the prevailing assumption that smoothness is desirable, offering a novel perspective on transformers' functional properties. The code is available at https://github.com/ambroiseodt/vit-plasticity.