Research map / Architectures and efficiency / Attention and transformers
Analyzing Feed-Forward Blocks in Transformers through the Lens of Attention Maps
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Extends norm-based attention-map analysis to whole Transformer layers including feed-forward blocks by decomposing GELU with integrated gradients, and finds FF blocks amplify specific linguistic compositions while FF and LayerNorm/residual cancel each other's contextualization effects.
Problem
Attention-map interpretations of Transformers ignore feed-forward (FF) blocks, even though they hold about two-thirds of layer parameters and may change how tokens are contextualized; the nonlinear activation prevents straightforward norm-based decomposition.
Approach
Decompose the FF nonlinearity (GELU and similar activations) with integrated gradients using a zero baseline, which satisfies completeness, so that norms of per-token contributions can be computed across the full layer (ATTN, RES, LN, FF) without approximation, yielding refined attention maps with component-by-component scopes.
Key results
- FF in BERT-base layers 3 and 9-11 and LN in layers 10-11 produce large contextualization changes; by-layer average change from FF was 0.21 (BERT) and 0.15 (GPT-2), and from LN 0.27 (BERT) and 0.38 (GPT-2).
- Residual vectors (RES2) are more than twice the L2 norm of FF outputs in most layers, and LN2 cancels FF's contextualization effect; ignoring outlier dimensions reduced FF change from 0.21 to 0.09 in BERT and 0.15 to 0.02 in GPT-2.
- Top FF-amplified token pairs often are subword-to-word or multi-word expression pairs; FF amplification correlates weakly with PMI (Spearman 0.06-0.14).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply FF to the attention-block contextualization | BERT-base and GPT-2 · Wikipedia | contextualization change (1 - ρ): 0.21 and 0.15 | Helped |
| Apply LN after the attention block | BERT-base and GPT-2 · Wikipedia | contextualization change (1 - ρ): 0.27 and 0.38 | Helped |
| Add RES2 after FF | BERT and GPT-2 Transformer layers · Wikipedia | contextualization change relative to before FF | Hurt |
| Add LN after FF and RES | BERT and GPT-2 Transformer layers · Wikipedia | contextualization change relative to before FF | Hurt |
| Ignore bottom 1% of dimensions with the lowest LN γ when calculating FF contextualization (ablation) | BERT · Wikipedia | contextualization change (1 - ρ), layer average: 0.21 to 0.09 | Hurt |
| Ignore bottom 1% of dimensions with the lowest LN γ when calculating FF contextualization (ablation) | GPT-2 · Wikipedia | contextualization change (1 - ρ), layer average: 0.15 to 0.02 | Hurt |
| Use FF to amplify selected token-pair contextualization | BERT-base and GPT-2 · Wikipedia | FF-amp score | Helped |
| Explain FF amplification using word co-occurrence | BERT-base · Wikipedia | Spearman correlation between FF-amp scores and PMI: 0.06--0.14 | No clear effect |
| Use LN γ weights against mean absolute FF output by dimension | BERT and GPT-2 · Wikipedia | Pearson correlation: -0.45 to -0.78 in BERTs; -0.22 to -0.59 in GPT-2 | Mixed |
Surprising findings
- FF-induced contextualization changes are largely canceled by subsequent residual and layer normalization components.
- Simple word co-occurrence (PMI) does not explain which token pairs FF amplifies.
Limitations
- Decomposition of FF with IG is an approximation-free but the analysis is restricted to small-to-medium BERT variants and two causal LMs; generality to large models is untested.
- Linguistic category labels for amplified token pairs were manually annotated by one author.
- Shapley value approach was too slow to be used, limiting comparison of attribution methods.
Experimental setup
Datasets: Wikipedia excerpts (992 sequences) Stanford Sentiment Treebank v2 (872 sequences from validation set)
Benchmarks: Wikipedia excerpts Stanford Sentiment Treebank v2
Metrics: 1 - Spearman rank correlation between attention maps (contex FF-amp score Pearson correlation (LN gamma vs FF output) PMI correlation
Baselines: Atb (norm-based attention block analysis) Shapley value (SHAP) DecompX
Builds on: norm-based analysis (ATTN decomposition) ALTI / attention block norm analysis Integrated Gradients GlobEnc
Closest papers by meaning
- On the Emergence of Position Bias in Transformers (ICML 2025 · similarity 0.89)
- What are you sinking? A geometric approach on attention sink (NeurIPS 2025 · similarity 0.88)
- Circuit Component Reuse Across Tasks in Transformer Language Models (ICLR 2024 · similarity 0.87)
- Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens (NeurIPS 2024 · similarity 0.87)
- Causal Head Gating: A Framework for Interpreting Roles of Attention Heads in Transformers (NeurIPS 2025 · similarity 0.87)
- Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers (Student Abstract) (AAAI 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.