PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Necessity of Attention-FFN Split in Vision Transformers

Junhyeok Kim, Jinyeong Kim, Jae Wan Park, Seong Jae Hwang

arXiv:2610.10303v1cs.CV

TL;DR

The necessity of the Transformer’s Attention–FFN alternation remains under-examined. The paper introduces AttenFeed and uViT to compare architectures that preserve or relax this separation, finding the dichotomy most restrictive at smaller model scales.

  • Problem

    The inductive bias of alternating Attention and FFN layers remains largely unexplored, leaving when this structure is necessary or redundant an open question.

    Source: §1 · p. 1

  • Method

    The authors introduce AttenFeed, a module that generalizes Attention and FFNs, and build uViT from sequences of AttenFeed modules to compare against ViT.

    Source: §1 · p. 2

  • Results

    uViT significantly outperforms ViT at smaller model scales, while their performance gap becomes marginal as model capacity increases.

    Source: §1 · p. 2

  • Takeaways & Limitations

    uViT provides a way to investigate when the conventional Attention–FFN structure facilitates learning and when it may be redundant.

    Source: §1 · pp. 1–2

    Pretraining datasets contain geographic, taxonomic, and demographic imbalances that downstream use of the resulting representations inherits.

    Source: Impact Statement · p. 9

Abstract

from arXiv · show

The standard Transformer architecture relies on a rigid pattern that alternates Attention and Feed-Forward Network (FFN) layers. Despite its widespread adoption, the inductive bias imposed by this strict separation has not been systematically examined. In this work, we investigate the necessity of the Attention-FFN dichotomy in Vision Transformers (ViTs). To facilitate this analysis, we introduce the AttenFeed module, a unified component that integrates the functional properties of both Attention and FFN. Based on this module, we devise the unified Vision Transformer (uViT), which replaces the conventional alternating Attention-FFN structure with a sequence of AttenFeed modules. We then use uViT as a control group that relaxes the Attention-FFN dichotomy of the standard ViT and systematically compare the two models across multiple datasets and model scales. Our experiments reveal that the Attention-FFN dichotomy can hinder performance at smaller model scales due to the rigid parameter allocation of ViTs. The AttenFeed module and uViT serve as new analytical tools for understanding the Attention-FFN structure and offer theoretical insights into the heuristically designed architecture of conventional ViTs.

1. Introduction

The paper examines whether ViT’s rigid Attention–FFN alternation is necessary, an architectural inductive bias that remains largely unexplored. It introduces AttenFeed and uViT to study that separation and reports that the standard dichotomy hinders performance at smaller model scales.

1. Introduction

  • The inductive bias imposed by the Transformer’s alternating Attention–FFN structure remains largely unexplored despite prior work on other architectural biases.Prior studies cited here found biases toward symmetric functions and sparse attention, while MoE work suggests consecutive Attention or FFN blocks can sometimes be more effective than alternation.Source: §1 · p. 1
  • AttenFeed is a unified component that generalizes the functional properties of both Attention and FFN, using a nonlinear activation to give attention FFN-like behavior.The paper describes this generalization as both mathematical and empirically validated.Source: §1 · p. 2
  • uViT consists of a sequence of AttenFeed modules, relaxing the strict structural separation between Attention and FFN in standard ViTs.The authors use uViT as a control group for investigating the Attention–FFN dichotomy.Source: §1 · p. 2
  • At smaller model scales, uViT significantly outperforms ViT, while their performance gap narrows and becomes marginal as model capacity increases.The authors attribute the smaller-scale disadvantage to the rigid parameter allocation of ViT.Source: §1 · p. 2
  • The authors present AttenFeed and uViT as tools for theoretical and empirical investigation of the heuristically established alternating Attention–FFN structure.They aim to support more systematic study of architectural inductive biases in ViT.Source: §1 · p. 2

2. Related Work

Related work characterizes several Transformer inductive biases and studies architectural changes that manipulate them. The paper identifies the necessity of rigid Attention–FFN separation as an unaddressed question.

2. Related Work

  • Prior work characterizes Transformer biases including preferences for symmetric functions, sparse attention, relational reasoning, and copy tasks.These findings examine biases within Transformers or compare them with other model families.Source: §2 · p. 2
  • Architectural modifications have manipulated inductive biases through locality-focused designs, reduced token-mixing bias, or training-time bias injection.Examples include Swin Transformer, MLP-Mixer, ResMLP, and methods for injecting biases during training.Source: §2 · p. 3
  • The paper identifies the necessity of rigid Attention–FFN separation as unexplored; GAU merges GLU and Attention for a different purpose, while this work uses a unified module to study the dichotomy.The cited work describes GAU as addressing limitations of linear attention rather than minimizing inductive bias.Source: §2 · p. 3

3. Method

The method introduces AttenFeed to unify Attention and FFN behavior, then builds uViT from these modules to relax their explicit separation. The comparison is designed to investigate when the conventional dichotomy benefits learning.

3. Method

  • The method aims to investigate when the inductive bias of Attention–FFN structure becomes beneficial by introducing AttenFeed and uViT.The section proceeds from a ViT block review to the unified module and then the overall architecture.Source: §3 · p. 3

3.1. Preliminary

  • A conventional ViT receives one [CLS] token and T image-patch tokens, then alternates multi-head Attention and position-wise FFN layers.The input token sequence at each layer is denoted Xℓ.Source: §3.1 · p. 3

3.2. AttenFeed Module

  • AttenFeed adds a nonlinear activation to the attention value projection, allowing its OV circuit to be interpreted as an FFN while retaining standard attention properties.The paper presents the module as a mathematical generalization of both Attention and FFN, with empirical validation.Source: §3.2 · p. 3

3.3. Unified Vision Transformer

  • uViT replaces the coupled Attention–FFN structure with one AttenFeed module per layer; its parameter count is matched to ViT by increasing layers and model dimension.The authors describe uViT as having approximately three times fewer parameters before this adjustment.Source: §3.3 · p. 4, Table 1 · p. 4

4. AttenFeed as a Generalization of Attention and FFN

AttenFeed is presented as a unified module that generalizes Attention and FFN, with uViT testing the properties of this design. The section combines mathematical results with empirical evidence of FFN-like, attention-like, and intermediate-parameter behavior.

4.1. Mathematical Description

  • Under self-focused attention, AttenFeed is equivalent to an FFN and can approximate standard attention, including under looser conditions.The approximation is also established under the small pre-activation condition in Proposition 4.2.Source: §4.1 · p. 4

4.2. FFN-like Behavior of the AttenFeed Module

  • uViT consistently has higher output rank than xViT, and the findings suggest GELU gives AttenFeed FFN-like rank-preserving properties.The comparison uses identical configurations; higher rank correlates with improved test accuracy.Source: §4.2 · pp. 4–5, Table 2 · p. 5

4.3. Attention-like Behavior of AttenFeed Module

4.4. Comparing Parametric Spaces of uViT and ViT

  • AttenFeed's WV and WO values occupy intermediate ranges between the corresponding standard attention and FFN parameters.This supports the claim that its parameter space does not deviate significantly from the regimes of standard attention and FFN.Source: §4.4 · p. 6

4. AttenFeed as a Generalization of Attention and FFN

  • AttenFeed mathematically generalizes standard Attention and FFN and empirically exhibits properties of both.Source: §4 · p. 4, §4.4 · p. 6

5. Impact of Inductive Bias in ViT

Comparisons of uViT and ViT indicate that relaxing the Attention–FFN split is associated with lower inductive bias and stronger performance at smaller scales. The performance gap narrows as model capacity increases.

5.1. Assessing Inductive Bias via Transfer Learning

  • uViT outperforms ViT in every evaluated linear-probing scenario and in most full fine-tuning cases across four downstream datasets.The authors take these transfer results as evidence that uViT has less inductive bias and more flexible output representations and internal parameters.Source: §5.1 · p. 7, Table 4 · p. 8

5.2. The Role of Inductive Bias at Scale

  • uViT significantly outperforms ViT at the Small scale, while their performance gap diminishes as models scale from Small to Base.The experiments used identical training conditions; the authors interpret the diminishing gap as evidence that the split becomes counterproductive at smaller scales.Source: §5.2 · p. 8, Table 5 · p. 8
  • The authors attribute uViT’s small-scale advantage to ViT’s rigid parameter allocation, which fixes parameters in Attention or FFN roles.They argue this can prevent small models from allocating more capacity to attention when token relationships are important for a dataset.Source: §5.2 · p. 8

6. Conclusion

The paper introduces AttenFeed and uViT to examine the Attention–FFN split in Vision Transformers, concluding that the split restricts performance at smaller model scales.

6. Conclusion

  • AttenFeed mathematically and empirically generalizes standard attention and FFNs, and uViT replaces their alternating structure with AttenFeed modules.The authors report that this unified design relaxes the structural inductive bias imposed by explicit functional separation.Source: §6 · p. 8
  • The authors conclude that the conventional Attention–FFN dichotomy acts as a restrictive inductive bias at smaller model scales.Source: §6 · p. 8

Impact Statement

The paper notes that pretraining data imbalances carry into downstream representations. Its mathematical analysis also shows that self-focused AttenFeed behavior reduces to a position-wise FFN.

Impact Statement

  • Downstream uses of representations from the studied pretraining datasets inherit their geographic, taxonomic, and demographic imbalances.The paper says these uses should follow established fairness and robustness practices for pretrained vision backbones.Source: Impact Statement · p. 9

A. Mathematical Descriptions

A.1. FFN-like Behavior.

  • When an AttenFeed head attends only to its own token, spatial mixing vanishes and its output approximates a position-wise FFN.The resulting expression is described as a standard token-wise GELU FFN.Source: §A.1 · p. 13

A.2. Attention-like Behavior.

The appendix supports AttenFeed’s attention-like behavior through mathematical equivalence arguments and describes a further behavior: spatially weighted mixing of FFN outputs.

A.2. Attention-like Behavior.

  • With sufficiently small pre-activation inputs, an AttenFeed head approximates a standard attention head.The stated condition is that the relevant inputs to GELU are sufficiently small.Source: §A.2 · p. 13
  • The appendix also argues that activation sparsity makes AttenFeed structurally converge to standard attention through a linear approximation of GELU.It treats large positive activations as approximately linear and near-zero activations as negligible, while assuming large negative activations are effectively suppressed.Source: §A.2 · pp. 13–14
  • Under Theorem 4.1’s conditions, some AttenFeed-head parameterizations approximate a standard attention head.The argument uses the mathematical equivalence between a SiLU-activated FFN layer and a standard attention head.Source: §A.2 · pp. 13–14

A.3. Mixture of FFN-like Behavior.

  • AttenFeed also exhibits spatially weighted mixing of FFN outputs, with the attention QK circuit governing their integration into the residual stream.This is presented as an emergent behavior in addition to AttenFeed’s ability to function as attention or as a standard FFN.Source: §A.3 · p. 15

B. Training Details

The experiments compare ViT and uViT in toy-scale classification and large-scale pre-training, holding non-architectural hyperparameters fixed to isolate the architectural difference.

B. Training Details

  • The evaluation spans toy-scale classification on small datasets and large-scale pre-training on ImageNet-1k, Places365, and iNaturalist-2021.The toy-scale regime is used for rank analyses; the large-scale regime uses three pre-training datasets.Source: B. Training Details · p. 16
  • In both regimes, hyperparameters unrelated to the ViT–uViT architectural difference are held fixed.The passage frames the comparison as controlling the optimization recipe across the two architectures.Source: B. Training Details · p. 16
  • The controlled setup is intended to attribute observed differences to the architectural change rather than the optimization recipe.Source: B. Training Details · p. 16

B.1. Toy-Scale Training

The appendix documents the training protocols and rank measurement procedure, then reports that uViT preserves or improves pre-head rank and accuracy on additional datasets, alongside qualitative attention-map agreement with ViT.

B.1. Toy-Scale Training

  • Toy-scale training uses CIFAR-10, CIFAR-100, STL-10, and SVHN, with matched augmentation and optimization settings across ViT, uViT, and noFFN.The models train for 100 epochs with AdamW; data splits, random seeds, and other listed settings are matched within each dataset.Source: §B.1 · p. 16

B.2. Large-Scale Pre-training

  • Large-scale pre-training uses ImageNet-1k, Places365-standard, and iNaturalist-2021, with the same training and evaluation pipeline for ViT and uViT.The protocol specifies 224×224 inputs, patch size 16, and 300 epochs with AdamW; runs use matched global batch sizes.Source: §B.2 · p. 16

C.1. Rank Measurement Procedure

  • The rank procedure uses a 95% variance effective-rank threshold and reports normalized rank to compare representations with different feature dimensions.Pre-head rank probes the backbone representation before classification, while post-head rank probes the logits; the two measures are described as complementary.Source: §C.1 · p. 17

C.2. Rank Preservation of AttenFeed Module

  • On STL-10, uViT improves pre-head rank by up to 0.081 and test accuracy by up to 3.20 percentage points over xViT.Across configurations, uViT’s pre-head rank is equal to or higher than xViT’s, with higher rank consistently accompanied by higher test accuracy.Source: Table 6 · p. 18, §C.2 · p. 18

C.3. Attention Map Comparison of uViT and ViT

  • Across varied images, uViT-B attends to salient object regions identified by ViT-B, with only minor differences in spatial extent.The figure caption describes the comparison as qualitative and says uViT-B reliably highlights the same salient object regions.Source: §C.3 · p. 19, Figure 4 · p. 19

C.4. Parametric Space Comparison of uViT and ViT

At Small and Large scales, uViT’s parameter distributions remain between the attention and FFN regimes observed in ViT, supporting a scale-agnostic unified character for AttenFeed.

C.4. Parametric Space Comparison of uViT and ViT

  • At Small and Large scales, uViT parameter distributions are wider than ViT attention parameters but narrower than ViT FFN parameters.The same intermediate-regime pattern is reported for Base-scale models, indicating that the pattern is not specific to one capacity.Source: Figure 5 · p. 20, §C.4 · p. 20, Figure 6 · p. 21

D. Limitation

The empirical analysis is limited to computer vision, and the comparison uses a standard ViT training recipe that may not be optimal for uViT.

D. Limitation

  • Experiments evaluate ViTs on image datasets only, leaving the Attention–FFN dichotomy and AttenFeed unvalidated in language or audio domains.The authors identify testing those non-vision domains as future work.Source: D. Limitation · p. 22
  • All models use a standard ViT training recipe for fairness, although uViT may have a distinct optimal optimization landscape.A uViT-specific recipe could potentially improve performance and alter observed scaling behaviors.Source: D. Limitation · p. 22
Loading 2610.10303v1…