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

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-LLMs 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.

• The paper introduces a novel integrated gradient and norm-based approach to visualize feed-forward block dynamics in Transformers.
• It finds that feed-forward blocks amplify specific token compositions and often cancel effects from residual and normalization layers.
• The study highlights architectural differences between masked and causal models, suggesting pathways to optimize Transformer efficiency.

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

In the paper titled "Analyzing Feed-Forward Blocks in Transformers through the Lens of Attention Map," the authors tackle the intricate task of deciphering the internal dynamics of Transformer models, focusing particularly on the often-overlooked feed-forward (FF) blocks. The study aims to render the input contextualization effects of these FF blocks in a human-friendly visualization format using refined attention maps. This paper embarks on this analysis across both masked- and causal-LLMs, contributing a novel perspective to the existing body of research on Transformer interpretability.

Summary of Contributions

The analysis presented in the paper extends the norm-based approach to interpret the entire Transformer layer, incorporating FF blocks alongside attention mechanisms, residual connections, and normalization processes. The authors establish that FF blocks significantly modify the input contextualization patterns, with emphasis on specific token compositions. Furthermore, the research uncovers an intriguing interplay where FF blocks and surrounding components often negate each other's effects, hinting at possible redundancies within Transformer layers.

Methodology

The authors employ an integrated gradient (IG) method to overcome the inherent challenges posed by the non-linear activation functions within FF blocks. By coupling IG with a norm-based analysis, they achieve a component-wise breakdown of the Transformer layer, allowing them to track contextualization changes at a granular level. This methodological advancement provides a refined attention map that can capture subtler transformations induced by the FF networks.

Key Findings

1. Contextualization by FF Blocks: The study reveals that FF blocks emphasize certain linguistic compositions, such as subwords-to-word and words-to-multi-word-expression constructions. The occurrence of amplified relationships is evident across various layers, particularly in the mid-to-late stages of the network.
2. Redundant Processing: A counterintuitive but significant observation is that the FF and its adjacent layers appear to cancel each other's effects. For instance, residual connections (RES) carry original signals, which can dominate over the modifications imposed by the FF, thus diminishing the FF’s impact. Similarly, layer normalization (LN), through its weighting parameters, often suppresses the unique dimensions introduced by FF transformations.
3. Variable Contextualization Patterns Across Architectures: The paper notes differences in contextualization changes between masked and causal models, pointing towards architecture-specific dynamics. For example, causal models tend to show significant changes in earlier layers compared to their masked counterparts.

Implications and Future Directions

The findings suggest potential pathways for optimizing Transformer architectures by mitigating redundancy. This could involve pruning strategies for weight parameters in FF blocks or devising mechanisms to better exploit their representational capacities without unnecessary overlap with other components.

Looking forward, applying this refined attention map analysis to different model architectures, such as the LLaMA and OPT series, could yield further insights. Additionally, adapting this interpretative framework to analyze novel Transformer variants that integrate FF-centric enhancements, such as adapters, may provide valuable guidance on their design and implementation.

The paper serves as a pivotal step towards a deeper understanding of Transformer layers, offering an analytical toolkit that balances precision and interpretability. This work opens avenues for enhancing model efficiency and unlocking new architectures, aligning with the overarching goal of refined Transformer interpretability.