Layer-wise Integrated Gradients (LIG)

Set-to-set Integrated Gradients for within-layer flow analysis in Transformers.

LIG applies Integrated Gradients (IG) at ATT and MLP module boundaries inside each Transformer layer, composes within-layer token-to-token contributions, and compares them with layer-whole attribution under an $L_2$ diagnostic — without model-specific retraining or per-operation interpreter design.

• Website & interactive demo: eightsuzuki.github.io/layer-wise-integrated-gradients
• Paper & arXiv: Coming soon

This release implements L2-scalarized set-to-set IG at module boundaries (z→u, u→z, z→z), with baselines ITB (self_input_token), ITB-zeroRatio (itb_zero_ratio), ITB-mapRatio (itb_map_ratio), Zero, and ATTITBa=0 (att_itb_a0) for MLP.

Quick start

Install PyTorch first (CUDA/CPU wheel index differs), then LIG:

# GPU (CUDA 12.1 example)
pip install torch --index-url https://download.pytorch.org/whl/cu121

# From PyPI
pip install layer-wise-integrated-gradients

# Or from GitHub
pip install "layer-wise-integrated-gradients @ git+https://github.com/eightsuzuki/layer-wise-integrated-gradients.git"

# Development clone
git clone https://github.com/eightsuzuki/layer-wise-integrated-gradients.git
cd layer-wise-integrated-gradients
pip install torch --index-url https://download.pytorch.org/whl/cu121
pip install -e ".[dev]"

Explain one sentence → JSON

from lig import explain

result = explain(
    "The cat sat on the mat.",
    model="bert-base-uncased",
    num_steps=32,
    granularity="all",           # att (z→u), mlp (u→z), layer (z→z)
    layers=[0, 11],
)

lig explain "The cat sat on the mat." --steps 32 --granularity all -o attributions.json

Supported models

GPT-2 example:

explain(
    "Hello world",
    model="gpt2",
    granularity="all",
    baseline_att="self_input_token",   # or itb_zero_ratio / itb_map_ratio
    baseline_mlp="att_itb_a0",         # or zero
    layers=[0],
)

ATT ITB self-contribution variants (see method section in the upcoming paper):

explain("Hello world", baseline_att="itb_zero_ratio", granularity="att", layers=[0])
explain("Hello world", baseline_att="itb_map_ratio", granularity="att", layers=[0])
explain("Hello world", baseline_layer="itb_zero_ratio", granularity="layer", layers=[0])

Other decoders (Llama, …): docs/DECODER_DESIGN.md.

Inspect boundaries (no IG run)

from lig import describe_boundaries

# Config-only (no weight download)
describe_boundaries("gpt2", load_weights=False)

# Full module introspection
describe_boundaries("bert-base-uncased", load_weights=True, device="cpu")

python examples/detect_model_boundaries.py bert-base-uncased gpt2 distilbert-base-uncased

Docker

docker compose -f docker-compose.cpu.yml build
docker compose -f docker-compose.cpu.yml run --rm lig explain "Hello world" \
  --steps 4 --granularity layer --layers 0 --target-tokens 1 -o /output/out.json

See docs/DOCKER.md.

Dev setup (uv)

bash scripts/ops/setup_uv_env.sh          # GPU torch + pip install -e ".[dev]"
bash scripts/ops/setup_uv_env.sh --cpu    # CPU torch
source .venv/bin/activate

Regression tests

bash scripts/run_regression_tests.sh

Paper demo (no PTB download)

Interactive z2z map for the paper example sentence (PTB dev sample 410, Zero baseline):

• GitHub Pages: demo — switch samples and routes (direct z→z / composed ATT×MLP; Zero / ITB / ITB-zeroRatio)
• Regenerate (PTB cache required): python scripts/build_paper_demos.py && python scripts/build_demo_html.py
• Streamlit: pip install -e ".[demo]" then streamlit run lig/demo/streamlit_app.py
• Precomputed JSON: examples/paper_demo/json/ (see DATA_NOTICE)

Paper reproduction (Experiment A)

Requires Treebank-3 (LDC99T42). See docs/REPRODUCTION.md.

Features

• One-call API — explain(text) → JSON with z→u, u→z, and z→z attributions
• CLI — lig explain "..." -o out.json
• Model-agnostic block boundaries — auto-detect z (residual stream) and u (ATT output / MLP input) from module layout; BERT-family + GPT-2 full granularity; block-level z→z for MPNet, DistilBERT, Switch, Mamba
• Boundary introspection — describe_boundaries("bert-base-uncased") or python examples/detect_model_boundaries.py
• Configurable — integration steps, baselines, granularity, layers, target tokens

Project layout

.
├── lig/                  # Public API + encoder adapters + viz demo
├── examples/paper_demo/  # One-sentence PTB excerpt + LIG JSON
├── docs/                 # DOCKER, PyPI, REPRODUCTION, Pages demo
├── utils/calculations/ig/  # Core LIG implementations
├── scripts/
│   ├── reproduce/        # Experiment A (PTB-gated)
│   ├── build_demo_html.py
│   └── run_regression_tests.sh
└── test/

Citation

@inproceedings{suzuki2026lig,
  title     = {LIG: Layer-wise Integrated Gradients for Within-Layer Flow Analysis in Transformers},
  author    = {Suzuki, Eight and Hino, Hideitsu and Murata, Noboru},
  year      = {2026},
  note      = {Code: https://github.com/eightsuzuki/layer-wise-integrated-gradients}
}

License

MIT License for this software.

PyPI

Package metadata and publish steps: docs/PYPI.md.

Maintainers: sync core IG code from the parent monorepo with bash scripts/sync_from_monorepo.sh when needed. See MANIFEST.md.