LIG: Layer-wise Integrated Gradients for Within-Layer Flow Analysis in Transformers

2026-06-19

AI Summary

This work addresses the limited fine-grained interpretability of information flow within Transformer models. It proposes Layer-wise Integrated Gradients (LIG), a novel method that extends Integrated Gradients from scalar-to-scalar mappings to set-to-set mappings for the first time, enabling attribution tracing at the boundaries between attention and MLP modules across layers with respect to output tokens. By employing L2 norm-based scalarization, LIG achieves set-level attribution while incorporating a Layer-wise Relevance Propagation–style composition rule to ensure attribution conservation at module boundaries. The method requires no model retraining and is readily applicable to standard architectures such as BERT-base. Evaluations on the PTB dataset demonstrate consistent intra-layer attributions under well-chosen baselines, effectively uncovering the internal mechanisms governing information propagation in Transformers.

Abstract

Transformers achieve strong performance, but their internal computations remain opaque. We view each Transformer layer as a dynamic graph whose nodes are token representations and per-head attention outputs, with Multi-Head Attention (ATT) and MLP as module boundaries. On this graph we use LIG (Layer-wise Integrated Gradients), which applies set-to-set Integrated Gradients (IG) at nonlinear module boundaries. Set-to-set IG applies IG to a map from a set of input token representations to a set of output representations, evaluating token-to-token contributions, which is not standard in prior IG applications. This extends IG from the usual scalar-objective setting to set-to-set maps via an L2 scalarization, and composes within-layer contributions in the spirit of Layer-wise Relevance Propagation (LRP), with IG completeness playing the role of LRP-style conservation at each boundary. We use LIG to analyze (i) the agreement between module-wise composition and layer-whole attribution under an L2 criterion, and (ii) within-layer information flow by tracing separated ATT and MLP contributions. On BERT-base and PTB, configurations that best preserved within-layer consistency used the target token's embedding as the ATT baseline and either the ATT output at a=0 or Zero as the MLP baseline. We therefore present LIG as a diagnostic XAI tool at module-boundary granularity, without model-specific retraining or per-operation interpreter design. Code is available at https://github.com/eightsuzuki/layer-wise-integrated-gradients.

Problem

• Transformers
• Integrated Gradients
• Explainable AI
• Information Flow
• Module-boundary Analysis

Innovation

• Layer-wise Integrated Gradients
• set-to-set attribution
• Transformer interpretability
• module-boundary analysis
• Integrated Gradients