PaperlayerView exact revision on arXiv

Source-linked AI summary

ViT-AdaLA: Adapting Vision Transformers with Linear Attention

Yifan Li, Seunghyun Yoon, Viet Dac Lai, Franck Dernoncourt, Jason Kuen, Yu Kong, Trung Bui

arXiv:2603.16063v1cs.CV

TL;DR

Vision foundation-model ViTs provide strong representations but face quadratic attention costs, while existing linear-attention training and LLM-oriented linearization approaches are costly or transfer poorly to vision. ViT-AdaLA progressively aligns linear attention and features with a pretrained softmax teacher before supervised fine-tuning, achieving strong classification and segmentation results with efficiency and resolution scalability.

  • Problem

    ViTs suffer quadratic complexity for long visual-token sequences, while existing ViT linear-attention methods require costly training from scratch and LLM linearization transfers poorly to vision models.

  • Method

    ViT-AdaLA uses attention alignment, feature alignment with a frozen softmax-ViT teacher, and supervised fine-tuning to transfer pretrained VFM priors to linearized ViTs.

  • Results

    ViT-AdaLA performs effectively across classification and segmentation tasks, maintaining classification accuracy within 1% of the original softmax backbone while improving efficiency and resolution scalability.

  • Takeaways & Limitations

    The framework transfers pretrained ViT priors to linear-attention models without large-scale pretraining and supports efficient scaling to higher resolutions.

  • Takeaways & Limitations

    Softmax attention can remain more efficient than linear attention when sequence length N is smaller than head dimension D, such as N = 256 and D = 1024.

Abstract

from arXiv · show

Vision Transformers (ViTs) based vision foundation models (VFMs) have achieved remarkable performance across diverse vision tasks, but suffer from quadratic complexity that limits scalability to long sequences. Existing linear attention approaches for ViTs are typically trained from scratch, requiring substantial computational resources, while linearization-based methods developed for large language model decoders do not transfer well to ViTs. To address these challenges, we propose ViT-AdaLA, a novel framework for effectively adapting and transferring prior knowledge from VFMs to linear attention ViTs. ViT-AdaLA consists of three stages: attention alignment, feature alignment, and supervised fine-tuning. In the attention alignment stage, we align vanilla linear attention with the original softmax-based attention in each block to approximate the behavior of softmax attention. However, residual approximation errors inevitably accumulate across layers. We mitigate this by fine-tuning the linearized ViT to align its final-layer features with a frozen softmax VFM teacher. Finally, the adapted prior knowledge is transferred to downstream tasks through supervised fine-tuning. Extensive experiments on classification and segmentation tasks demonstrate the effectiveness and generality of ViT-AdaLA over various state-of-the-art linear attention counterpart.

1. Introduction

ViT-AdaLA addresses quadratic attention and weak transfer from decoder-focused linearization by adapting pretrained ViT priors through progressive alignment. The framework avoids expensive training from scratch and is evaluated across classification and segmentation.

  • Motivation: Training-from-scratch linear attention methods require large-scale pretraining, whereas linearization leverages an off-the-shelf pretrained ViT.The latter paradigm substantially reduces the need for extensive pretraining.
  • Motivation: Softmax self-attention scales quadratically with visual-token sequence length, increasing computational and memory overhead for long-sequence vision applications.Linear attention reduces the cited complexity from O(N^2D) to O(ND^2).
  • Motivation: LLM linearization methods transfer poorly to ViTs because decoder-only LLMs generate targets, while ViTs primarily extract features for separate task-specific heads.The paper attributes the performance drop to errors accumulating temporally in LLMs but spatially and hierarchically in ViTs.
  • ViT-AdaLA: ViT-AdaLA adapts pretrained softmax-ViT knowledge through attention alignment, feature alignment, and supervised fine-tuning.Attention modules are first aligned block-wise; the full linearized ViT then aligns final-layer features with a frozen teacher.
  • Contribution: The proposed paradigm shifts linear-attention ViT development from designing approximations alone toward adapting prior knowledge from pretrained ViTs.This enables linearized ViTs to inherit existing VFM priors without expensive training from scratch.
  • Evaluation: Experiments across classification and segmentation tasks validate ViT-AdaLA’s effectiveness, efficiency, and resolution scalability across different VFMs and downstream tasks.The evaluation compares against a wide range of state-of-the-art linear-attention baselines.

2. Related Work

Prior efficient-Transformer research includes memory optimization, token reduction, local attention, state-space modeling, and two linearized-Transformer paradigms. Linearization-based methods adapt pretrained softmax Transformers, but the cited approaches include methods developed for decoder-based LLMs.

  • Efficient Attention: Efficient-attention research reduces Transformer memory and computation overhead through methods including FlashAttention, token selection or merging, shifted windows, and Mamba-based state-space modeling.These approaches target the quadratic cost of standard attention through different computational or token-processing mechanisms.
  • Linear Attention: Linearized Transformers divide into training-from-scratch approaches that design attention approximations and linearization-based approaches that adapt existing softmax Transformers.The first category trains models to acquire prior knowledge, while the second category modifies pretrained models.
  • Linear Attention: Training-from-scratch methods improve attention approximations through alternative activation functions and related linear-attention designs.Their paradigm targets accurate attention approximation while obtaining prior knowledge through training from scratch.
  • Linear Attention: Linearization-based methods include Hedgehog, LoLCATS, and Lizard, which approximate or combine attention mechanisms for pretrained Transformer adaptation.LoLCATS uses attention transfer and low-rank linearization based on LoRA for decoder-based LLMs, while Lizard adds local attention to global GLA attention.

3. Method

ViT-AdaLA adapts pretrained softmax-based ViTs into linear-attention models through attention alignment, feature alignment, and supervised fine-tuning. The method uses kernel-based linear attention for linear complexity while preserving attention behavior and final-layer features.

  • Preliminaries: Softmax attention has quadratic complexity O(N^2D), whereas kernel-based linear attention reorders query-key computation to achieve O(ND^2).The linear formulation uses ϕ(·)=ELU(·)+1 to decompose the exponential kernel.
  • Attention Alignment: Stage 1 tunes the added linear-attention module against the corresponding softmax module while keeping the original model features unchanged.The alignment loss uses mean squared error over attention outputs from each block, and only the linear module's query, key, and value projections are updated.
  • Attention Alignment: Unlike LoLCATS, ViT-AdaLA directly tunes query, key, and value projections in vanilla linear attention, yielding higher efficiency and improved approximation quality.LoLCATS tunes two additional query- and key-mapping modules based on Hedgehog linear attention.
  • Feature Alignment: Stage 2 replaces all softmax modules with Stage 1 linear modules and aligns the resulting ViT's final representations with those of the frozen original ViT.The feature-alignment loss uses MSE between the original and linearized models' final features, with λ controlling output-loss scale.
  • Supervised Fine-tuning: Stage 3 transfers the adapted linearized ViT to downstream tasks by appending a task-specific head and jointly fine-tuning the backbone and head.The three-stage pipeline is attention alignment, feature alignment, and supervised fine-tuning.

4. Experiment

Experiments evaluate ViT-AdaLA across classification, semantic segmentation, resolution scaling, and training-stage ablations. The method preserves strong performance while improving efficiency and transferring prior knowledge from pretrained vision foundation models.

  • Classification: ViT-AdaLA achieves the highest top-1 accuracy among evaluated vision foundation models, remaining within 1% of the original softmax backbone while preserving efficiency.The comparison uses ImageNet-1K with a single linear-layer classification head.
  • Classification: Aligning final feature representations with the original VFM teacher is essential for effective knowledge transfer in vision models.Attention alignment alone leaves accumulated approximation errors across layers insufficiently corrected.
  • Semantic Segmentation: ViT-AdaLA demonstrates strong performance across VFMs on ADE20K semantic segmentation, rivaling supervised IN1K-pretrained ViT baselines.Segmentation evaluation uses a Mask2Former head across methods.
  • Training Ablations: Stage 1 accelerates Stage 2 convergence and enhances prior-knowledge extraction from the VFM teacher.Stage 2 aligns final-layer features and later adopts early stopping after performance saturation.
  • Training Ablations: Vanilla linear attention provides a superior softmax-attention approximation to other evaluated linear attention variants while retaining high computational efficiency.The Stage 1 comparison tunes query, key, and value weights in every layer for each baseline.
  • Training Ablations: ViT-AdaLA is independent of the specific linear attention architecture and supports direct adaptation of downstream task models.Task-model adaptation slightly outperforms VFM-only adaptation by inheriting more task-relevant priors.

5. Conclusion

ViT-AdaLA adapts vision foundation model priors into linearized ViTs without large-scale pretraining through a three-stage alignment process. The framework distills priors, accelerates convergence, and supports resolution scalability.

  • ViT-AdaLA adapts vision foundation model priors into linearized ViTs without large-scale pretraining.
  • Its three-stage alignment process distills priors, accelerates convergence, and supports robust resolution scalability.

Impact Statement

ViT-AdaLA targets efficient linear attention for vision transformers, reducing the computational burden of long sequences while adapting existing pretrained models. Its broader deployment may preserve or amplify risks associated with vision foundation models.

  • Linearizing pretrained ViTs can reduce computational and memory costs for long input sequences in resource-constrained settings.The paper identifies edge devices, robotics, and real-time perception systems as potential settings.
  • The techniques rely on publicly available datasets and do not introduce new data-collection or supervision modalities.
  • Efficient models may amplify existing vision-foundation-model risks, including pretrained-data bias and misuse in surveillance or privacy-sensitive applications.The paper calls for responsible use, dataset auditing, and appropriate governance.
  • Compared baselines: The compared linear-attention baselines include Hedgehog, LoLCATS, Linformer, Performer, Cosformer, Nyströmformer, and Monarch Attention.The experiments reproduce these methods using their publicly available codebases.
  • Activation choices: Kernel-based linear attention applies non-negative feature maps to queries and keys, with softmax, softplus, ReLU, and ELU+1 compared as variants.
  • Activation choices: ELU + 1 consistently achieves lower Stage 1 training loss and is selected as the default activation for numerical stability and convergence.

A.4. Training Time Comparison

Stage 1 is substantially faster than Stage 2, while ViT-AdaLA maintains a small performance gap from the original softmax model across different model sizes.

  • Stage 1 is much faster than Stage 2 across the evaluated vision foundation models.Stage 1 uses COCO, Stage 2 uses ImageNet-22K, and experiments run on 8 H100s.

A.6. Comparison with Training-from-Scratch-based Methods

Training linear-attention ViTs from scratch substantially lags ViT-AdaLA after the reported training schedules. The comparison suggests that much longer pretraining may be needed for competitive accuracy.

  • Training-from-scratch linear attention substantially lags ViT-AdaLA across ADE20K, Cityscapes, and ImageNet-1K.The vanilla baseline was pretrained on ImageNet-1K for 200 epochs, while ViT-AdaLA used 20 Stage 2 epochs on ImageNet-22K.
  • The reported gap suggests that much longer pretraining may be required for training-from-scratch linear attention to reach competitive accuracy.
  • Tuning only the QKV matrices decreases performance compared with tuning all components.The authors hypothesize that QKV projections alone lack sufficient capacity to absorb the teacher's high-dimensional distributions.

A.8. Experiments on Classification Tasks

ViT-AdaLA outperforms approximation-based linear attention baselines and remains competitive with softmax attention across classification and segmentation evaluations.

  • Approximation-based Monarch and Nyströmformer significantly outperform other linear attention baselines on CLIP-L classification benchmarks.The comparison indicates that approximating softmax attention is advantageous.
  • ViT-AdaLA and its Stage 2 variant surpass Monarch and Nyströmformer on the reported classification benchmarks.
  • ViT-AdaLA maintains competitive performance with softmax attention on ADE20K and Cityscapes segmentation tasks.Linear attention baselines show larger gaps from the softmax upper bound in segmentation than in classification.

B. More Visualization Results

Visualization and ablation results show that ViT-AdaLA retains more pretrained VFM knowledge than Monarch Attention, with Stage 2 providing the strongest preservation and Stage 1 further improving transfer.

  • More Visualization Results: Figure 14 ablates the Stage 1 and Stage 2 training procedures to compare their effects on PCA-projected final-layer features.
  • More Visualization Results: ViT-AdaLA features more closely match original softmax VFM features than Monarch Attention features in PCA visualizations.The comparison uses final-layer DINOv2-L features projected into three PCA channels.
  • More Visualization Results: Stage 2 preserves most original VFM features, while combining Stages 1 and 2 yields the strongest prior-knowledge retention.
  • More Visualization Results: Table 10 reports Top-1 fine-tuning accuracy on ImageNet-1K under CLIP-L, while Table 9 compares tuning strategies on CLIP-L and DINOv2-L at 512×512 resolution.

C. Limitations

The paper identifies scope, performance, efficiency, and resolution-dependent limitations for ViT-AdaLA, while suggesting extensions to other tasks and vision-language models.

  • Limitations: ViT-AdaLA has been validated on classification and segmentation, but its generalizability to object detection and image generation remains uninvestigated.
  • Limitations: A marginal performance gap remains between ViT-AdaLA and full softmax attention in segmentation tasks.
  • Limitations: Stage 2 training efficiency could be improved through advanced distillation strategies such as masked image modeling.
  • Limitations: ViT-AdaLA has higher computational overhead than softmax attention on low-resolution images, reflecting a challenge shared by many linear attention architectures.
  • Future Directions: Future work could extend linear attention to vision large language models processing long visual sequences, including high-resolution images and long videos.
Loading 2603.16063v1…