PaperlayerView exact revision on arXiv

Source-linked AI summary

Copy the Same, Distill the Difference: Initializing Linear Vision Transformers

Huaiyuan Qin, Muli Yang, Gabriel James Goenawan, Shiqi Huang, Min Kass Chong, Wahyu Wiratama, Peng Hu, Chen Gong, Wu Liu, Xi Peng, Chun Jian Ho, Hongyuan Zhu

arXiv:2609.35745v1cs.CVcs.AIcs.LG

TL;DR

Linear ViTs offer linear-complexity token routing but typically require from-scratch pre-training and underperform Softmax ViTs, leaving effective initialization unclear. This study tests reuse of pre-trained Softmax weights and finds that copying MLPs while distilling attention outputs can close the gap and even surpass Softmax models.

  • Problem

    Linear ViTs typically require from-scratch pre-training and underperform Softmax ViTs, while effective reuse of pre-trained Softmax weights for initialization remains unclear.

    Source: §1 · p. 1

  • Method

    The study tests component-level Softmax-to-linear transfer, copying MLP weights and distilling attention outputs rather than copying attention weights.

    Source: §1 · p. 2

  • Results

    Copying MLPs and distilling attention outputs together close the gap to Softmax ViTs and can even let linear ViTs surpass them.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    Across the studied settings, copy operator-agnostic MLPs and distill attention behavior to reuse pre-trained weights across the Softmax-to-linear boundary.

    Source: §1 · p. 3

    The approach relies on public pre-trained Softmax weights, some trained on large curated or web-scale datasets that may inherit dataset biases.

    Source: ETHICS STATEMENT · p. 10

Abstract

from arXiv · show

Linear Vision Transformers (ViTs) are designed to replace the attention in Softmax ViTs with the linear-complexity attention operator for more efficient token routing, but they require from-scratch pre-training and typically underperform the original Softmax version. How to initialize linear ViTs both efficiently and effectively still remains unclear. In this work, we explicitly ask: given that most foundation ViTs are built on the mainstream Softmax attention, can linear ViTs benefit from their pre-trained weights? Recent works on Attention Transfer show that attention is the effective transferable component between Softmax ViTs, suggesting attention alone suffices for such reuse. However, we find the opposite for Softmax-to-linear transfer. The attention weights are operator-specific: copying them barely helps, and is sometimes even worse than random initialization. Instead, the attention's token routing behavior can be recovered through distillation with a proper loss design, letting linear ViTs reduce the gap and even match Softmax ones. In contrast, the MLP weights, which carry the learned representation, are operator-agnostic: they can be transferred by simple direct copying, which already carries most of the benefit of the pre-trained weights. Thus, copying MLPs can serve as an effective foundation for Softmax-to-linear transfer: paired with the distilled attention, linear ViTs eventually close the remaining gap and even surpass Softmax ones. These findings hold consistently across various linear ViT variants, different model sizes, and diverse datasets. We hope this study deepens the understanding of reusing pre-trained weights across attention operators: copy what stays the same and distill what differs, to recover the benefit across the Softmax-to-linear boundary.

1 INTRODUCTION

The paper asks whether linear ViTs can reuse pre-trained Softmax weights, then finds that transfer depends on the component: distill attention outputs and copy MLPs. Together, these strategies close the Softmax performance gap and can surpass Softmax ViTs.

1 INTRODUCTION

  • Linear ViTs offer linear-complexity token routing, but typically require from-scratch pre-training and underperform Softmax ViTs, motivating weight reuse.The paper asks whether pre-trained Softmax weights can initialize linear ViTs effectively.Source: §1 · p. 1
  • Copying pre-trained Softmax attention weights barely helps linear ViTs and can perform worse than random initialization.The failure is consistent across the linear ViT variants studied.Source: §1 · p. 2
  • Matching attention outputs through a properly designed distillation loss recovers token routing behavior and can close most of the gap to Softmax ViTs.The proposed loss matches outputs rather than operator-specific attention maps.Source: §1 · pp. 2–3
  • Copying MLPs and distilling attention outputs together close the Softmax gap and can let linear ViTs surpass Softmax ones.The authors report this pattern across diverse settings.Source: §1 · p. 2
  • MLP weights are operator-agnostic, carry learned representations, and provide most of the benefit of pre-trained weights when copied.Paired with distilled attention, copied MLPs can close the remaining gap or surpass Softmax ViTs.Source: §1 · pp. 2–3

2 RELATED WORK

Prior work explores efficient linear attention, ViT initialization, attention transfer, and conversion from Softmax Transformers. This paper distinguishes its focus as reuse across the attention-operator boundary, with copying for MLPs and distillation for attention outputs.

Efficient Linear ViTs

  • Linear-attention designs span vision, language, and state-space models, but typically rely on dedicated pre-training from scratch.Examples include InLine (Han et al., 2024), MHLA (Zhang et al., 2026), and TTT.Source: §2 · p. 3

Initialization for ViTs

  • Analytical initialization uses structural priors, while other methods reuse weights without crossing attention operators; this work studies transfer across that boundary.Mimetic (Trockman & Kolter, 2023) shapes query-key products, while structured initialization (Zheng et al., 2025) imposes convolution-like patterns.Source: §2 · p. 3

Attention Transfer for ViTs

  • Attention Transfer recovers pre-trained benefits between Softmax ViTs, but its effectiveness across attention operators had not been studied.Li et al. (2024) showed attention patterns can suffice within Softmax ViTs; Qin et al. (2026b) studied architectural mismatch within that family.Source: §2 · p. 3

Linearizing Transformers

3 PRELIMINARIES

Softmax attention routes tokens using an attention map, while linear attention reorders computation to reduce its cost. The experimental setup compares five linear ViT variants initialized from pre-trained Softmax models across multiple datasets and against random and Softmax baselines.

3.1 FROM SOFTMAX TO LINEAR ATTENTION

  • Softmax attention routes queries, keys, and values through an attention map with quadratic O(N^2) cost in token count.Linear attention instead uses a kernel function to reorder computation, reducing cost to O(N).Source: §3.1 · p. 4

3.2 EXPERIMENTAL SETUP

  • The study tests five linear ViT variants at Tiny, Small, and Base sizes, using corresponding pre-trained DeiT weights as transfer sources.Unmatched parameters in variants with additional components retain their default initialization.Source: §3.2 · p. 4
  • Tiny and Small models are evaluated on six transfer datasets, while all three sizes are also evaluated on ImageNet-1K.Experiments use 300 or 600 training epochs depending on dataset scale, and reported results average three random seeds.Source: §3.2 · p. 4
  • The random baseline represents training from scratch, while the same-size, same-task pre-trained Softmax ViT defines the recovery target.A transfer is considered effective only if it improves over the random baseline.Source: §3.2 · p. 4

4 SOFTMAX ATTENTION WEIGHTS ARE OPERATOR-SPECIFIC

The section tests whether Softmax attention weights transfer by copying into linear ViTs and finds no reliable advantage over random initialization. This failure persists across model sizes, pre-training sources, and distribution shifts.

4.2 MAIN RESULTS

  • Across five linear ViTs and seven datasets, attention-only Copy never stands out and can underperform Random, unlike the expected transfer advantage.The results identify copied Softmax attention weights as operator-specific: their benefit does not reliably survive the change of attention operator.Source: Figure 2 · p. 5, §4.2 · p. 6

4.1 ATTENTION-ONLY TRANSFER

  • Copy loads pre-trained Softmax Q, K, and V projections into each linear ViT block while leaving other components random, testing whether attention alone preserves performance.Mimetic and Impulse provide analytical, pre-training-free comparison baselines.Source: §4.1 · pp. 5–6

4.3 ROBUSTNESS OF THE FAILURE

  • Across DINO, MoCov3, iBOT, and MAE pre-trained weights, copying shifts performance only marginally and never reliably beats Random.The pattern rules out the supervised DeiT weights as the sole explanation for the failure.Source: Table 1 · p. 6, §4.3 · pp. 6–7
  • On four distribution-shift benchmarks, Copy gains no additional robustness over Random and remains equally far from the Softmax reference.The out-of-distribution result is consistent with the in-distribution comparison.Source: Table 2 · p. 6, §4.3 · p. 7
  • At every tested model size, attention-only Copy remains at or near Random, while analytical initializations retain their advantage as size grows.This supports the conclusion that the failure is not an artifact of limited model capacity.Source: Figure 2 · p. 5, §4.3 · p. 6

5 DISTILL THE ROUTING, COPY THE REPRESENTATION

The paper separates transferable behavior from transferable weights: it distills attention outputs to recover token routing and directly copies MLPs to transfer representation. Combining these transfers closes the remaining gap to the Softmax target and can match or surpass it.

5.1 TOKEN ROUTING CAN BE RECOVERED THROUGH DISTILLATION

  • For Softmax-to-linear transfer, distillation matches attention outputs rather than operator-specific maps, recovering most of the gap across operators and sometimes reaching the Softmax target.On Flowers, attention distillation alone leaves a larger gap, and for MHLA it falls below Copy; the text attributes this to weak task signal for relearning representation.Source: §5.1 · p. 8

5.2 MLP WEIGHTS ARE OPERATOR-AGNOSTIC

  • Copying MLP blocks raises accuracy above Random and attention-only Copy across datasets, carrying most of the benefit of pre-trained weights.The authors characterize MLP weights as operator-agnostic and the transferred representation as residing primarily in those blocks.Source: Figure 4 · p. 8, §5.2 · p. 9

5.3 THE RECIPE

  • Combining copied MLPs with distilled attention largely closes remaining gaps across operators and datasets, with MLPCopy and FullCopy plus distillation reaching the Softmax level.MLPCopy w/ Distill lands close to FullCopy w/ Distill, so copied attention weights add no consistent advantage once routing is distilled.Source: §5.3 · p. 9

5.4 DISCUSSION AND ANALYSIS

  • Across ReLU and ELU, map matching and output matching have comparable peak recovery, but for signed-map InLine, map matching consistently falls behind.The analysis supports output matching as a more general design; the output-matching peak is stable across a broad λ plateau, motivating λ = 50 as the default.Source: Figure 5 · p. 9, §5.4 · p. 10
  • The transfer patterns also hold across linear ViT variants, model sizes, and dense prediction tasks under the same experimental protocols.The main section results otherwise use ReLU-based ViT-Small on classification.Source: §5.4 · p. 10

6 CONCLUSION

The conclusion finds that copying Softmax attention weights alone can fail across the Softmax-to-linear boundary, whereas attention distillation and MLP copying transfer useful behavior and representation. Together, they can match or surpass the Softmax target.

6 CONCLUSION

  • Copied Softmax attention weights are operator-specific and can perform below random initialization when transferred alone.The conclusion contrasts this with recovering token routing through distillation using attention-output matching.Source: §6 · p. 10
  • MLP weights are operator-agnostic, and directly copying them carries most of the benefit of pre-trained weights.The conclusion identifies the MLP representation as the component that transfers effectively by copying.Source: §6 · p. 10
  • Pairing copied MLPs with distilled attention closes the remaining gap to the Softmax target and can even surpass it.The paper summarizes the recipe as copying what stays the same and distilling what differs.Source: §6 · p. 10

AI USE STATEMENT

The authors disclose AI assistance in manuscript preparation and experiment code, state that AI-assisted content was reviewed, and accept responsibility for the final content.

AI USE STATEMENT

  • AI assisted with writing, polishing, figure plotting, and code implementations for the experiments.The authors state that all AI-assisted content shown in the manuscript was reviewed and verified.Source: AI USE STATEMENT · p. 10
  • The authors say they reviewed and verified all AI-assisted content shown in the manuscript.The statement covers the displayed manuscript content.Source: AI USE STATEMENT · p. 10
  • The authors take responsibility for the final content, including text, claims, and artifacts produced with AI assistance.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

The paper studies Softmax-to-linear initialization as a methodological question and identifies both a potential benefit and a scope for caution. Its recipe does not remove biases or safety concerns inherited from the pre-trained teacher weights.

ETHICS STATEMENT

  • The work’s direct positive impact is to encourage more efficient and effective linear ViT initialization and sharpen understanding of weight reuse across attention operators.The paper frames this as a methodological study of initialization using pre-trained Softmax weights.Source: ETHICS STATEMENT · p. 10
  • The findings should not be interpreted as showing that Attention Transfer fails across attention operators; the authors present a recipe for that transfer instead.The stated principle is to copy what stays the same and distill what differs.Source: ETHICS STATEMENT · p. 10
  • The recipe can make transfer more effective but does not remove biases or safety concerns inherited from public pre-trained Softmax weights.Some source weights are trained on large-scale curated or web-scale datasets.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

The authors commit to releasing code and provide experimental-protocol details to support replication. They report repeating experiments over three random seeds and point to a statistical-significance analysis.

REPRODUCIBILITY STATEMENT

  • The authors commit to making source code and replication scripts publicly available upon publication.The planned code includes implementations and scripts to replicate the reported empirical results.Source: REPRODUCIBILITY STATEMENT · p. 11
  • The paper points to protocol details covering the model zoo, transfer and evaluation protocols, and reference-baseline selection.These details are provided in Sec. 3.2 and further in Sec. B.Source: REPRODUCIBILITY STATEMENT · p. 11
  • All experiments are repeated over three random seeds, with a statistical-significance analysis provided in Table F.Source: REPRODUCIBILITY STATEMENT · p. 11

A SCOPE AND RELATION TO SIMILAR METHODS

This is a controlled empirical study of Softmax-weight reuse for initializing linear ViTs, not a new linearization method. It distinguishes which components transfer by copying from those better transferred behaviorally, and sets its comparisons apart from related conversion and Attention Transfer settings.

A SCOPE AND RELATION TO SIMILAR METHODS

  • The study factorizes pre-trained Softmax weights by component and tests copying or behavioral transfer under fixed recipes across downstream tasks and five linear attention operators.The authors do not claim a new linearization method or improvement over the related methods on their own benchmarks.Source: A SCOPE AND RELATION TO SIMILAR METHODS · p. 15
  • The paper’s attention-output matching follows LoLCATs’ established objective, while its contribution is showing that the behavioral target works for ViTs across five operators, including those without attention maps.Source: A SCOPE AND RELATION TO SIMILAR METHODS · p. 15
  • Compared with ViT-AdaLA’s dedicated alignment and fine-tuning pipeline for vanilla linear attention, the paper offers an operator-general recipe without a dedicated adaptation stage.Source: A SCOPE AND RELATION TO SIMILAR METHODS · p. 15
  • The component-level analysis relates to LiT’s selective inheritance and hybrid distillation, and to T5’s parameter inheritance, by identifying MLPs as copyable and attention as better transferred by distillation.The paper places attention-output distillation at initialization rather than distilling task predictions, and contrasts its decomposition with T5’s FullCopy-like setting.Source: A SCOPE AND RELATION TO SIMILAR METHODS · p. 15
  • DiD studies post-hoc, label-free conversion under a frozen detector, whereas this work initializes a linear ViT for supervised training with trainable downstream heads.The authors describe the findings as complementary: DiD aligns behavior for a fixed detector, while this study examines component transfer at initialization.Source: A SCOPE AND RELATION TO SIMILAR METHODS · p. 15
  • Unlike prior Attention Transfer between same-architecture Softmax ViTs, this study tests whether Softmax attention weights can initialize a linear operator.Here copied projections are initialization only, and all parameters are jointly trained; the Copy failure therefore does not contradict within-Softmax Attention Transfer.Source: A SCOPE AND RELATION TO SIMILAR METHODS · pp. 15–16

B.1 DETAILS OF MODEL ZOO

  • The model zoo uses five linear ViT variants whose attention operators differ, including ELU/ReLU, InLine, MHLA, and TTT; projection shapes are not uniform across all variants.MHLA adds mixing and depthwise-convolution components, while TTT uses additional inner-model components and a wider output projection.Source: §B.1 · p. 16

B.2 DETAILS OF TRANSFER PROTOCOL

  • The transfer protocol uses same-size ImageNet-1K DeiT weights by default, compares Random, Copy, MLPCopy, and FullCopy settings, and jointly trains all parameters after initialization.For TTT, only standard QKV rows are loaded where the output projection shape differs; unmatched components use their default initializations.Source: §B.2 · pp. 16–17

B.3 DETAILS OF EVALUATION PROTOCOL

Across additional operators, model sizes, and evaluation settings, the reported pattern is consistent: copied attention alone brings no reliable advantage, while MLP copying carries most transfer benefit and attention distillation closes much of the remaining gap. The section also specifies the evaluation and distillation protocols underlying these comparisons.

B.3 DETAILS OF EVALUATION PROTOCOL

  • In Softmax-to-linear transfer experiments, results use standard train/test splits, report Top-1 Accuracy, average over 3 random seeds, and follow 300- or 600-epoch schedules by dataset scale.Source: §B.3 · p. 17
  • The authors report that copying patch embeddings, positional embeddings, and the class token does not yield significant gains, so the analysis focuses on attention and MLP initialization.Source: §B.3 · p. 17

B.4 DETAILS OF ATTENTION DISTILLATION

  • Attention distillation uses a frozen Softmax teacher and per-block MSE matching of student and teacher attention outputs during downstream fine-tuning, with objective L = Ltask + λ Ldistill.For MHLA and TTT, only patch tokens are matched; the default λ is 50 across operators, datasets, and initialization settings.Source: §B.4 · p. 18
  • The map-matching ablation uses recorded attention-module inputs, compares teacher Softmax maps with student kernel maps, and excludes MHLA and TTT because they do not materialize maps.Output-matching and map-matching loss weights are swept over separate ranges because their magnitudes differ.Source: §B.4 · p. 18

C.1 ADDITIONAL RESULTS ON MORE LINEAR OPERATORS

  • Across all five linear operators and datasets, Copy remains near Random, MLPCopy lifts performance, FullCopy adds gains, and MLPCopy with distilled attention approaches FullCopy with distillation and often matches or surpasses Softmax.Adding attention distillation lifts accuracy in nearly every setting; Copy with Distill exceeds Random with Distill for almost every operator and dataset.Source: §C.1 · p. 18, §C.2 · p. 19

C.2 ADDITIONAL RESULTS ON MORE MODEL SIZES

  • At smaller model sizes, the patterns persist with smaller margins: distillation lifts nearly every setting, and MLPCopy with Distill lands close to FullCopy with Distill and matches the Softmax target.Source: §C.2 · p. 19

C.3 ADDITIONAL RESULTS ON IMAGENET AND DENSE PREDICTION

  • On ImageNet-1K and dense prediction, Copy brings no reliable advantage over Random, while MLPCopy carries most pre-trained-weight benefit and distillation lets MLPCopy approach FullCopy with Distill and match or exceed Softmax.The reported dense tasks include object detection, instance segmentation, and semantic segmentation; the comparison uses ReLU and TTT.Source: §C.3 · p. 20

C.4 ADDITIONAL RESULTS ON COPYING THE REMAINING COMPONENTS

The additional results find that copying class, positional, and patch embeddings adds no significant accuracy benefit beyond copying attention and MLP weights. The accompanying interpretation explains why copied MLPs and distilled attention address different sources of mismatch, while noting the study’s empirical scope and training requirements.

C.4 ADDITIONAL RESULTS ON COPYING THE REMAINING COMPONENTS

  • Removing copied class, positional, and patch embeddings causes no significant accuracy drops when attention and MLP weights remain copied.Table G describes the ImageNet-1K, Base-size, ReLU ablation and its Top-1 Accuracy metric; Table D describes ImageNet-1K and dense prediction evaluation settings.Source: Table D · p. 21, §C.4 · p. 22

C.5 ADDITIONAL RESULTS ON STATISTICAL SIGNIFICANCE

  • The study reports standard deviations for ImageNet-1K Top-1 Accuracy with and without distillation across three random seeds.These significance analyses concern results reported in Tables D and E.Source: §C.5 · p. 22

D ADDITIONAL INTERPRETATION OF SOFTMAX-TO-LINEAR TRANSFER

D.1 PRESERVATION UNDER WEIGHT COPYING

  • Replacing Softmax with a kernel-based linear operator changes attention outputs even for fixed Q, K, and V, so copied projections do not preserve the pretrained attention function.Randomly initialized remaining components can also change the inputs received by copied attention projections.Source: §D.1 · p. 22
  • Because the MLP architecture is unchanged, copying its weights preserves its token-wise computation, although its input may differ between teacher and student.The input mismatch is affected by the attention output in the same transformer block.Source: §D.1 · p. 22

D.2 WHY COPYING MLPS AND DISTILLING ATTENTION ARE COMPLEMENTARY

  • MLPCopy removes the computation-mismatch term, while attention distillation targets the remaining input mismatch caused mainly by replacing the attention operator.The decomposition separates differences in MLP computation from differences in MLP inputs.Source: §D.2 · p. 23

E LIMITATIONS

  • The main transfer experiments use ImageNet-1K DeiT-pretrained weights, the recipe requires an auxiliary objective during fine-tuning, and the interpretation is not a formal derivation.The authors identify a complete formal theory of Softmax-to-linear or Attention Transfer as future work.Source: E LIMITATIONS · p. 23
Loading 2609.35745v1…