Analyzing Feed-Forward Blocks in Transformers through the Lens of Attention Maps

TL;DR

This work investigates how feed-forward blocks (FFBs) in Transformer layers influence input contextualization. It extends a norm-based attention-map framework with Integrated Gradients to decompose FF, RES, and LN effects across four scopes, enabling visualization of FF-induced context shifts. Across masked and causal LMs, FFs amplify particular linguistic compositions and these effects are largely canceled by surrounding components, revealing redundancy in Transformer computation. The approach yields interpretable, component- and layer-level insights and offers a path toward more robust mechanistic interpretability.

Abstract

Transformers are ubiquitous in wide tasks. Interpreting their internals is a pivotal goal. Nevertheless, their particular components, feed-forward (FF) blocks, have typically been less analyzed despite their substantial parameter amounts. We analyze the input contextualization effects of FF blocks by rendering them in the attention maps as a human-friendly visualization scheme. Our experiments with both masked- and causal-language models reveal that FF networks modify the input contextualization to emphasize specific types of linguistic compositions. In addition, FF and its surrounding components tend to cancel out each other's effects, suggesting potential redundancy in the processing of the Transformer layer.

Figures (113)

• Figure 1: Overview of the Transformer layer for Post-LN and Pre-LN architectures, annotated with analysis scopes, e.g., AtbFfResLn. The right part of this figure (token-to-token attention map) illustrates the component-by-component changes of the attention maps. See Appendix \ref{['appendix:sec:samples_visualization']} for concrete examples of attention maps.
• Figure 2: An illustration of possible contextualization effects by FFB. The FFB does not have the function of mixing input tokens together; however, its input already contain mixed information from multiple tokens, and the FFB is capable of altering these weights. Here, output $\bm y_3$ is computed based on $[\bm x_1, \cdots, \bm x_5]$. Based on the vector norm, the most influential input seems to be $\bm x_3$ before FF; however, after the FF's transformation, $\bm x_5$ becomes the most influential input.
• Figure 3: +Ff (Atb$\leftrightarrow$AtbFf)
• Figure 4: +Res (AtbFf$\leftrightarrow$AtbFfRes)
• Figure 5: +Ln (AtbFfRes$\leftrightarrow$AtbFfResLn)