PaperlayerView exact revision on arXiv

Source-linked AI summary

Vision Transformer Finetuning Benefits from Non-Smooth Components

Ambroise Odonnat, Laetitia Chapel, Romain Tavenard, Ievgen Redko

arXiv:2602.06883v3cs.LGcs.CVstat.ML

TL;DR

Transfer learning in vision transformers lacks a clear understanding of how component smoothness affects adaptation. The paper formalizes plasticity, ranks components theoretically and empirically, and finds that finetuning high-plasticity attention and feedforward modules gives better, more stable performance.

  • Problem

    The role of transformer smoothness in transfer learning remains poorly understood despite distribution shifts requiring weight updates for downstream adaptation.

  • Method

    The paper defines plasticity as components’ average rate of change, derives theoretical plasticity rankings, and validates them through extensive vision-transformer finetuning experiments.

  • Results

    Finetuning high-plasticity attention modules and feedforward layers consistently produces higher and more stable performance across initialization and learning rates.

  • Takeaways & Limitations

    Prioritizing non-smooth attention and feedforward components offers practical guidance for adapting vision transformers.

  • Takeaways & Limitations

    The study does not examine how tailored optimization, such as adaptive learning rates or schedulers, affects module-specific finetuning performance.

Abstract

from arXiv · show

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.

1. Introduction

The paper studies which vision-transformer components should be prioritized during finetuning by defining plasticity as an average rate of change. Theory and experiments show that higher-plasticity attention and feedforward modules yield better, more stable adaptation performance.

  • Our approach: Plasticity formalizes a module’s response to input variation as its average rate of change, with higher values indicating lower smoothness.The measure captures how component outputs adapt when inputs change.
  • Our contributions: The authors derive upper bounds that establish a theoretical ranking of transformer components by plasticity.This ranking provides a principled basis for prioritizing components during adaptation.
  • Our contributions: Attention has the highest plasticity, followed by the first and second feedforward layers, feedforward-preceding LayerNorm, and attention-preceding LayerNorm.The ordering is observed in pretrained ViTs and also holds for DINOv3 and GPT2.
  • Our contributions: Adapting high-plasticity attention and feedforward modules produces higher and more stable performance across initialization and learning rates.The study uses exhaustive finetuning runs on diverse classification benchmarks with both SGD and Adam optimizers.
  • Our contributions: The work links plasticity to gradient norms and presents this relationship as guidance for designing more efficient adaptation methods.Its findings offer a novel perspective on smoothness in finetuning.

2. Background

This section defines neural-network smoothness through concepts including Lipschitz continuity and introduces vision transformers as patch-based image models built from transformer encoders. It also identifies LayerNorm, multi-head self-attention, and feedforward linear layers as key transformer components.

  • Neural network smoothness: Smoothness encompasses differentiability, Lipschitz continuity, and robustness to input perturbations.The section presents Lipschitz continuity as a common way to quantify smoothness.
  • Neural network smoothness: A function is Lipschitz continuous when output differences are bounded by a constant multiple of input differences.The smallest such constant is the function’s Lipschitz constant, denoted Lip(f).
  • Vision transformers: A vision transformer splits 2D images into patches, embeds them as token sequences, and processes them through successive transformer encoders.A classification token, CLS, is prepended for classification.
  • Transformer components: Transformer encoders contain LayerNorms, multi-head self-attention, and feedforward linear layers operating on token sequences.The section recalls the operations of each module, including row-wise softmax in attention and GeLU between the feedforward linear layers.

3. Vision transformer plasticity

This section defines transformer plasticity as sensitivity to input changes and uses it to argue that non-smooth components may adapt more effectively during finetuning. Plasticity complements worst-case Lipschitz analysis by measuring average rates of change over input sequences.

  • Expected benefits of plasticity: Although smoothness supports generalization, training stability, and adversarial robustness, excessive smoothness can constrain capacity and adaptability to new tasks.This motivates studying plasticity as a potentially useful property for transfer learning.
  • Plasticity measure: Plasticity is defined as a module’s ability to adapt its output in response to changes in inputs.The definition considers distinct pairs of token sequences under a uniform distribution.
  • Plasticity measure: The plasticity measure captures average component behavior over input sequences, while satisfying P(f) ≤ Lip(f).The Lipschitz constant provides a worst-case bound, whereas plasticity measures the overall rate of change.
  • Connection to finetuning: High plasticity implies a high Lipschitz constant and therefore low smoothness.This links sensitivity to input changes with the smoothness properties of transformer components.
  • Connection to finetuning: Input-output and weight-output smoothness are theoretically connected, making plasticity relevant to parameter adaptation during finetuning.Prior theory characterizes these notions as two sides of the same coin.
  • Expected benefits of plasticity: Highly plastic, non-smooth components are expected to permit larger gradient norms, enabling faster and better downstream adaptation.The paper explicitly does not expect a linear relationship between plasticity and downstream performance.

4. Theoretical analysis

The theoretical analysis derives plasticity upper bounds for LayerNorm, feedforward, and attention components, yielding a ranking in which attention is most plastic and normalization least plastic. The results also extend to causal attention in decoder-only language models.

  • Upper bounds: LayerNorm has plasticity upper bound P(f) ≤ 1 under the stated common-token-statistics and positive-minimum-standard-deviation assumptions.The assumptions are justified by ImageNet1k preprocessing and token embedding properties.
  • Upper bounds: Feedforward linear layers have plasticity upper bound P(f) ≤ ∥W∥2, while attention receives looser control because its bounds depend on head spectral norms and sequence length.The attention analysis restricts token sequences to bounded sets because self-attention is not globally Lipschitz continuous.
  • Upper bounds: The tighter attention bound improves the sequence-length growth rate to √n under a bounded-total-energy image assumption.The bounded-token setting is considered practical; ViTs typically have average token norm 20 and sequence length around 200.
  • Limitations: Mean-field bounds independent of sequence length are mathematically interesting but unsuitable for vision transformers, whose sequence length is usually below 10^3.The limitation arises because practical ViT sequence lengths remain finite and comparatively small.
  • Theoretical ranking: The proposed plasticity ranking is MHA → FC1 ≈ FC2 → LN2 ≈ LN1.This ranking follows from comparing the relative tightness of the component-wise upper bounds.
  • Extension to large language models: The LayerNorm, feedforward, and attention results extend to decoder-only language models because causal masking preserves the relevant attention bound.Decoder blocks share the encoder-block global structure, differing at the attention module through causality.

5. Experiments

Experiments validate the theoretical plasticity ranking of ViT components and show that prioritizing high-plasticity attention and feedforward layers improves finetuning performance, stability, and optimization.

  • Empirical ranking: The empirical plasticity ranking across benchmarks is MHA → FC1 → FC2 → LN2 → LN1, supporting the theoretical ordering.Attention and feedforward linear layers have P(f) > 1, while LayerNorms are comparatively rigid.
  • Better performance: High-plasticity attention and feedforward components achieve better finetuning performance than LayerNorms on most benchmarks, especially Cifar100, Clipart, and Sketch.The attention module’s improvement is statistically significant relative to other modules except FC1.
  • Better performance: Single-component finetuning can compete with parameter-efficient adaptation: LayerNorms reach 89.4% with 28K parameters on Cifar100, versus LoRA’s 88.1% with 400K parameters.The LayerNorm result uses Adam, while the LoRA comparison is from the cited prior work.
  • Interplay between plasticity and optimization: Higher plasticity facilitates optimization, with larger gradient norms and steeper loss descent for attention and feedforward layers across benchmarks, learning rates, and seeds.These results support prioritizing the attention module and first feedforward linear layer during finetuning.
  • Visual Task Adaptation Benchmark: On VTAB with 1,000 training samples and 2,500 finetuning steps, prioritizing high-plasticity components, especially MHA, outperforms low-plasticity LN1 and LN2 overall.The protocol follows the original ViT paper setup.

6. Related work

Prior work has studied smoothness across several deep-learning settings and has established that estimating neural-network Lipschitz constants is difficult, especially for self-attention. Parameter-efficient finetuning methods broadly use selective parameter updates or additive adapters, positioning this work among selective approaches.

  • Smoothness: Smoothness has been studied in deep learning for generalization, training stability, generative modeling, adversarial robustness, and differential privacy.
  • Lipschitz constant estimation: Estimating neural-network Lipschitz constants is difficult because theoretical bounds are often loose, while self-attention makes transformer estimation more involved.Kim et al. (2021) showed that vanilla attention is not globally Lipschitz.
  • Parameter-efficient finetuning: Parameter-efficient finetuning includes selective approaches that update only a subset of parameters and additive methods that insert small adapters into models.The passage gives normalization layers as an example of inserted adapters.

7. Discussion · Appendix · A. Extended related work

The discussion concludes that prioritizing high-plasticity attention and feedforward components improves finetuning, while identifying tailored optimization and decoder-only models as future directions. The appendices organize proofs, implementation details, code, carbon reporting, experiments, and extended related work on smoothness and parameter-efficient finetuning.

  • 7. Discussion: High-plasticity attention modules and feedforward layers consistently produce better and more stable finetuning performance.The paper uses plasticity, understood through average smoothness, to identify components to prioritize during adaptation.
  • 7. Discussion: Tailored optimization, including adaptive learning rates and schedulers, is proposed for studying module-specific finetuning performance.The authors also suggest extending their methodology and theoretical insights to decoder-only models and LLM adaptability.
  • Appendix: The appendix roadmap assigns additional related work to Appendix A, theoretical proofs to Appendix B, and code and carbon-footprint details to Appendix C.These appendices are followed by full implementation details in Appendix D and extensive additional experiments in Appendix E.
  • A. Extended related work: Prior work studies neural-network smoothness, commonly quantified by Lipschitz constants and spectral norms, in generalization, training stability, generative modeling, and adversarial robustness.The extended related-work section positions the paper within these established smoothness research contexts.
  • A. Extended related work: At large scale, smoothness research addresses training instabilities and loss spikes that can cause large models such as LLMs to diverge.The passage also notes companion work inspired by mechanistic interpretability that studies gradient descent, although the supplied text is truncated.
  • A. Extended related work: Linear and activation layers have known tight Lipschitz constants, whereas feedforward-network estimation is NP-hard and has loose theoretical upper bounds.The nonlinear nature of self-attention makes its Lipschitz-constant estimation challenging as well.
  • A. Extended related work: Parameter-efficient finetuning methods comprise five main families, including selective methods that update parameter subsets and additive methods that insert trainable adapters.The passage identifies selective methods as common in vision models and additive methods as adapter networks placed between model layers.

B. Proofs … C.1. Open-source code

The appendices derive plasticity bounds for LayerNorm, linear layers, and self-attention using matrix-norm properties, bounded token sequences, and image-energy assumptions. The accompanying code and findings are publicly available and designed for reuse and extension.

  • B. Proofs: The proofs rely on standard Frobenius and spectral norm properties, including submultiplicativity, operator-norm bounds, and invariance under transposition.These properties are introduced as tools for the theoretical derivations.
  • B.1. Proof of Proposition 1: LayerNorm plasticity is bounded by applying matrix-norm inequalities under equal token means and standard deviations across paired sequences.The standard deviations are assumed positive and dependent on the embedding layer rather than the input sequences.
  • B.2. Proof of Proposition 2: Linear-layer plasticity bounds follow directly from the Frobenius norm inequality for the first and second feedforward weight matrices.The two matrices have shapes d×4d and 4d×d, respectively.
  • B.3. Proof of Proposition 3: The self-attention proof applies a Lipschitz bound on bounded token sequences and converts supremum rate-of-change bounds into expected plasticity bounds.The argument applies the cited self-attention theorem to each attention head module.
  • B.4. Proof of Proposition 4: For image inputs, bounded image energy implies bounded token-sequence Frobenius norms after patch embedding, enabling the self-attention bound.The embedding spectral norm determines the resulting token-radius parameter.
  • B.4. Proof of Proposition 4: The image-based attention proof combines triangular inequalities, matrix-norm bounds, row-stochastic attention weights, and a row-wise softmax Lipschitz result.These steps control both transformed token differences and attention-score differences.
  • C.1. Open-source code: The code and findings are publicly available at github.com/ambroiseodt/vit-plasticity, with a library that supports adapting all or part of the implementation.The authors intend the library to facilitate knowledge transfer, reduce redundant training runs, and support extensions to large language models.

C.2. Carbon footprint · D. Implementation details · D.1. Vision transformers

The paper reports a large experimental effort and specifies the ViT architectures and implementation conventions used in its finetuning experiments. The experiments involved approximately 1,000 runs, while the models were ImageNet22k-pretrained ViTs spanning three sizes and corresponding patch configurations.

  • C.2. Carbon footprint: Around 1,000 finetuning runs required the equivalent of 3,700 GPU hours.Using public cloud providers such as Azure or Amazon Web Services could cost up to $40,000.
  • C.2. Carbon footprint: Estimated emissions were roughly 259 kgCO2eq under a carbon efficiency of 0.1 kgCO2eq/kWh in France.The paper compares this estimate with a round-trip Paris–Madrid Boeing 737 flight and notes that France’s electricity grid makes the figure low.
  • D. Implementation details: The implementation-details section provides full implementation details.
  • D.1. Vision transformers: ViT inputs are 2D images split into square patches of size P, flattened and linearly embedded in dimension d.A classification token CLS is prepended before positional embeddings are added, and the resulting token sequence is processed by successive transformer encoders.
  • D.1. Vision transformers: The experiments use ViT models with 86M, 307M, and 632M parameters and patch sizes 16, 16, and 14, respectively.All models are pretrained on ImageNet22k.
  • D.1. Vision transformers: The code follows the original ViT implementation and uses a convolutional layer to embed images.The paper identifies this convolutional embedding as standard in the HuggingFace implementation.
  • D.1. Vision transformers: Table 4 lists ViT variants by patch size, sequence length, layers, attention heads, embedding dimension, FFN hidden dimension, and parameter count.
  • D.1. Vision transformers: Figure 6 presents the ViT-Base implementation.

D.2. Data preprocessing … E.1.1. THEORETICAL PLASTICITY RANKING

The paper evaluates plasticity and isolated-component finetuning across diverse vision benchmarks using standardized preprocessing and controlled optimization. Its theoretical analysis ranks transformer components by upper-bound plasticity, with attention and feedforward layers ranked highest.

  • D.2. Data preprocessing: Training images receive random cropping, 224×224 resizing, and random horizontal flips, while evaluation images are resized before center cropping and ImageNet1k-normalized.The stated normalization uses mean [0.485, 0.456, 0.406] and standard deviation [0.229, 0.224, 0.225].
  • D.3. Plasticity setup: Plasticity is computed from pretrained-model token sequences over pretraining and downstream distributions without assuming a known shift category.The implementation loops over N GPU batches of size b, stores outputs on the CPU, and uses N × b samples.
  • D.4. Finetuning setup: Finetuning isolates one trainable component at a time by freezing the remaining ViT-Base or ViT-Large weights and randomly initializing the classification head.The components are LN1, MHA, LN2, FC1, and FC2.
  • D.4. Finetuning setup: All configurations share the same inference cost, while their training memory differs through optimizer-state and gradient computation.The model parameters, optimizer states, gradients, and activations contribute to GPU memory load, but only optimizer and gradient computation vary between configurations in this setting.
  • D.4. Finetuning setup: Optimization uses SGD with momentum 0.9, no weight decay, cosine decay, 2000 warmup steps, batch size 512, and gradient clipping at norm 1.Each dataset–configuration pair is evaluated over 4 learning rates and 3 runs with different network-initialization and dataloader seeds.
  • D.4. Finetuning setup: Final performance is measured as test accuracy from the checkpoint achieving the best validation accuracy.Validation uses 20% of the training set.
  • E.1.1. THEORETICAL PLASTICITY RANKING: On ViT-Base, theoretical plasticity upper bounds rank the attention module highest, followed by FC1, FC2, LN2, and LN1.The computation uses sequence length n = 197, H = 12 attention heads, and reference radius r = 19.4 from Cifar10.

E.1.2. PLASTICITY EXPERIMENTS OF ALL BENCHMARKS

Across additional benchmark datasets, plasticity consistently follows the theoretical ranking: attention is highest, followed by the feedforward linear layers, while LayerNorms are most rigid. LayerNorm plasticity remains below 1, reinforcing the observed component hierarchy.

  • Benchmark-wide results: The experiments extend the Section 5.1 analysis to additional datasets and show findings aligned with the theoretical analysis in Section 4.Results are displayed in Figures 8 to 17.
  • Component ranking: Attention has the highest plasticity, followed by the first feedforward linear layer and then the second feedforward linear layer.This ranking is observed across the additional benchmark experiments.
  • Component ranking: LayerNorms are the most rigid components, with plasticity below 1.The ordering places the LayerNorm preceding the feedforward before the LayerNorm preceding attention.

E.1.3. PLASTICITY EXPERIMENTS ON DINOV3 AND GPT2 … E.2.2. EXTENSION TO VIT-LARGE (307M)

The paper extends plasticity analysis to DINOv3 and GPT2, finding that attention and feedforward components rank highest in plasticity. Across finetuning benchmarks and ViT-Large, higher-plasticity components consistently achieve better performance.

  • E.1.3. PLASTICITY EXPERIMENTS ON DINOV3 AND GPT2: The analysis extends to DINOv3-7B and GPT2-124M, using Cifar10 for DINOv3 and accounting for their respective pretraining sets.DINOv3 was self-supervised and pretrained on a set including ImageNet22k; GPT2 was pretrained on WebText.
  • E.1.3. PLASTICITY EXPERIMENTS ON DINOV3 AND GPT2: Across ViT-Base, DINOv3-7B, and GPT2, attention has the highest plasticity, followed by the first and second feedforward linear layers.The observed rate-of-change distributions follow the theoretical ranking described in Section 4.
  • E.2. Finetuning analysis: The supplementary finetuning analysis reports additional results, figures, and experiments related to Section 5.2.These materials include performance comparisons and an extension to a larger vision transformer.
  • E.2.1. PERFORMANCE COMPARISON: Finetuning performance is collected for trainable-component configurations across 11 image-classification benchmarks, with relative gain defined against linear probing.Table 8 reports best top-1 test accuracy over each dataset’s learning-rate grid, using three runs with different seeds.
  • E.2.1. PERFORMANCE COMPARISON: Across the reported comparisons, higher plasticity is associated with better finetuning performance.The same pattern appears in the relative-gain analysis and in the overall benchmark visualization.
  • E.2.2. EXTENSION TO VIT-LARGE (307M): The finetuning analysis extends to ViT-Large with 307M parameters and reaches similar conclusions.Results are reported on the Clipart benchmark in Table 9.
  • E.2.2. EXTENSION TO VIT-LARGE (307M): Table 9 evaluates ViT-Large configurations using average top-1 test accuracy on Clipart over the learning-rate grid.Entries average over learning rates, while non-smooth components are highlighted and the best component configuration is bolded.

E.2.3. ROBUSTNESS ANALYSIS

Across learning rates and random seeds on 11 benchmarks, plastic transformer components produce more stable finetuning performance, with attention consistently performing better and showing little variation.

  • Robustness analysis: The attention module consistently delivers better and more stable finetuning performance, with small performance variation.Components are ordered by decreasing plasticity in the robustness comparison.
  • Robustness analysis: Plastic transformer components result in more stable finetuning performance across learning rates and seeds on 11 benchmarks.The robustness comparison evaluates learning rates from Table 7 and three seeds relative to network initialization and dataloaders.

E.2.4. GRADIENT NORM ANALYSIS · E.2.5. EXTENSION TO ADAM

Gradient-norm dynamics follow the plasticity ranking: plastic components have higher gradient norms, faster validation-loss descent, and better convergence, especially on challenging datasets. The extension to Adam preserves the main finding that non-smooth, highly plastic components provide better and more stable finetuning performance, although Adam’s accuracy is slightly lower.

  • E.2.4. GRADIENT NORM ANALYSIS: High-plasticity components converge faster and better, with especially salient benefits on Cifar100, Clipart, and Sketch.The analysis links the scale of gradient norms and validation-loss descent to the plasticity ranking.
  • E.2.4. GRADIENT NORM ANALYSIS: Plastic components consistently have higher gradient norms, producing steeper validation-loss descent and better downstream performance.This pattern is observed across benchmarks, learning rates, and seeds.
  • E.2.4. GRADIENT NORM ANALYSIS: On Cifar10, repeated training dynamics show the same faster and better convergence pattern for components with high plasticity.The pattern is reported for seeds 0, 42, and 3407.
  • E.2.4. GRADIENT NORM ANALYSIS: On Cifar100, the training dynamics likewise show faster and better convergence for components with high plasticity.This is shown for seed 0 and reported consistently for seeds 42 and 3407.
  • E.2.4. GRADIENT NORM ANALYSIS: On Pet and Flowers102, all components have lower gradient norms because the pretrained model already achieves good linear-probing performance.These datasets therefore differ from more challenging benchmarks in gradient-norm scale.
  • E.2.5. EXTENSION TO ADAM: Adam preserves the central conclusion from SGD: non-smooth components yield better and more stable finetuning performance.The comparison uses Adam’s decoupled weight-decay variant and reports similar takeaways for both optimizers.
  • E.2.5. EXTENSION TO ADAM: Adam’s accuracy is slightly lower than SGD’s in the comparison summarized by Table 10.The table nevertheless states that the non-smooth-component advantage remains.
  • E.2.5. EXTENSION TO ADAM: With Adam, plastic components perform better and more stably across learning rates, while showing larger gradient norms and better generalization.Figure 54 reports similar trends for Adam and SGD on Sketch.
Loading 2602.06883v3…