Gradient family · Integrated Gradients averaged over baselines
Learn the method — a beginner's guide
New to explainable AI? This block explains the method as it is used on neural
networks, then how we run the very same method on an Atari 2600, where — unlike for a
network — we know the right answer. Every implementation step links to the code. See also the
XAI primer for the shared vocabulary.
The idea in plain words
Integrated Gradients (IG) needs a baseline, an input that means “nothing”. As
the IG guide showed, the answer depends heavily on that choice: a black image, a blurred image and
random noise each give a different map. There is rarely a single right “nothing”.
Expected Gradients (EG) sidesteps the choice. Instead of one hand-picked baseline it uses
a whole reference distribution: a pool of real inputs, for example images from the training
set. It runs IG from many randomly drawn references and averages the results. The attribution then
answers “how does this input differ from a typical input, and which features carry the
difference?” This links EG to Shapley values from game theory, which also average a
feature's contribution over many reference situations. EG is stable and keeps completeness in
expectation. But stable does not mean faithful, and averaging differences can cancel them.
How it works on a neural network
With \(y = f(x)\) and a reference distribution \(D\), EG assigns input \(i\) the
expected value (average) of the IG integrand over references \(x'\) and path positions
\(\alpha\):
\[ \mathrm{EG}_i(x) \;=\; \mathbb{E}_{x' \sim D,\;\alpha \sim U(0,1)}
\Big[(x_i - x'_i)\,\frac{\partial f}{\partial x_i}\big(x' + \alpha(x - x')\big)\Big]. \]
Erion et al. estimate it by Monte Carlo integration: draw \(M\) pairs
\((x'^{(m)}, \alpha_m)\), one gradient each, and average. There is no inner Riemann sum; the random
\(\alpha\) does the path integral. Summing over \(i\) gives completeness in expectation:
\(\sum_i \mathrm{EG}_i = f(x) - \mathbb{E}_{x'}[f(x')]\).
Toy example: averaging over baselines. \(f(x_1, x_2) = x_1 x_2\) at \(x = (2, 3)\), with
\(D = \{(0,0),\,(2,0)\}\), each drawn half the time. From the IG guide: IG from \((0,0)\) is
\((3, 3)\); IG from \((2,0)\) is \((0, 6)\). EG is the average, \((1.5,\ 4.5)\). Check:
\(1.5 + 4.5 = 6 = f(x) - \tfrac12\big(f(0,0) + f(2,0)\big) = 6 - 0\). The baseline choice is no
longer a single arbitrary pick, but the result now depends on the pool \(D\).
Toy example: cancellation. Let \(f(x) = x\) at \(x = 3\), with references
\(\{1, 5\}\). The two IG terms are \(3 - 1 = 2\) and \(3 - 5 = -2\). EG \(= 0\). Yet \(f\) depends on \(x\)
one-for-one. EG measures how far \(x\) is from the average reference. An input sitting at
the typical value gets no credit, however much the output depends on it.
Recipe
- Build a reference pool \(D\) of real inputs (on a network: training examples).
- Draw \(M\) references \(x'^{(m)}\) from \(D\) (with replacement) and one \(\alpha_m \sim U(0,1)\) for each.
- For each pair, form \(x' + \alpha (x - x')\), run one backward pass, and multiply the gradient by \((x - x')\).
- Average the \(M\) products: that is the EG attribution.
- Optionally check completeness against \(f(x) - \overline{f(x')}\), and repeat with other seeds to measure stability.
From a neural network to the Atari
| The network \(f\) | The Atari 2600 VCS running the ROM, emulated bit-exactly by jutari; its differentiable twin is JuTari.Diff. |
|---|
| The input \(x\) | The 128 RAM bytes at the intervention frame \(f^*\) (frame 90 of seeded random play). |
|---|
| The reference distribution \(D\) | The 90 genuine RAM states at frames 1–90 of the same seeded gameplay run: every row is a reachable machine state. |
|---|
| Monte Carlo samples | \(M = 64\) (reference, \(\alpha\)) draws per run, five independent seeds for the stability estimate. |
|---|
| The output \(y\) | A content output (one RAM byte read one-hot) and a position output (screen pixels changed versus the un-intervened frame, made differentiable by the bilinear sampler). |
|---|
| “Is the attribution right?” | The intervention oracle: poke each cause, re-run the real ROM, measure \(|\Delta y|\). |
|---|
How we apply it to the VCS
The runner implements Erion's single-\(\alpha\)-per-reference estimator over the RAM vector: 64 draws
per run, each a uniformly sampled row of \(D\) and a uniform \(\alpha\). It repeats this with five seeds,
reports the mean pairwise cosine similarity between the seeds' maps (stability), and also scores
single-baseline IG with the zeros baseline on the same output for a head-to-head comparison.
- The reference pool. \(D\) is built by replaying the shared seeded gameplay stream and
recording the RAM after each of the first 90 frames. An earlier version used a pool from a NOOP
(no-input) replay. There, many candidate bytes never changed, so \(x - x' = 0\) for every
reference and EG was exactly zero. The gameplay pool was introduced so the causal bytes actually
vary across references.
- Interpolated RAM. As in IG, the blended states \(x' + \alpha(x - x')\) have fractional bytes. They are
only fed to the differentiable read, never executed on the machine.
- Content output. The one-hot read has gradient \(e_s\) everywhere, so
\(\mathrm{EG}_s = \frac1M\sum_m (x_s - x'^{(m)}_s)\): the current byte minus the average of the
sampled references, and 0 for every other byte. This is the cancellation toy for real. A score byte
that only grows keeps a clear signal. A byte that sits near its own trajectory average gets almost
none, though the oracle may show it matters. The runner records the byte's value, its range over
\(D\), and whether EG vanished across all seeds.
- Position output. The naive gradient is zero for every reference. With the sampler, each draw
contributes \((x - x')\) times the kernel's slope at one random point of its path, a noisy one-sample
estimate of the occupancy difference between the sprite at the reference's column and at the real one.
- Scoring is the shared contract: correlation and precision@3 against the oracle,
deletion/insertion curves on the real ROM, and an oracle-vs-itself control.
Read the code
XAI the interpretability method
substrate the bit-exact emulator it runs on (jutari, Julia).
Reading the result
The stability score is high: different seeds give nearly the same map. That is EG's promise, and
it keeps it. But look at agreement with the oracle. On the content output the map is still a
single spike at the read byte, now sized by its distance from the average reference, and that
spike can shrink to almost nothing. On the position output agreement stays low. Overall EG is
among the least faithful attribution methods in the study while producing tidy, plausible-looking
maps. That combination, high plausibility and low faithfulness, is exactly what the intervention
oracle is designed to expose.
Pitfalls
- Stable is not faithful. Averaging removes randomness from the baseline choice, not error. A map can be perfectly reproducible and wrong.
- Cancellation. Signed differences from references above and below the input cancel. Inputs at their typical value get little credit even when the output depends on them.
- The pool is the new baseline. Results depend on \(D\). A pool where an input never varies gives it zero credit, as the NOOP pool did on the VCS.
- Monte Carlo noise. One gradient per draw is cheap but noisy; check stability across seeds, as the runner does.
- Still a gradient method. Where the gradient vanishes (sprite timing), no reference distribution can bring it back.
Try it yourself
julia --project=jutari tools/xai_study/phaseB_attribution/expected_gradients.jl --games core
Writes expected_gradients_<game>_{content,ball_pixel}.json/.npz and expected_gradients_core_summary.json into tools/xai_study/phaseB_attribution/out/. Defaults: --eg-samples 64 --eg-stability-seeds 5; the IG head-to-head uses --ig-steps 64.
Original papers & further reading
- Erion, Janizek, Sturmfels, Lundberg & Lee (2021). Improving performance of deep learning models with axiomatic attribution priors and expected gradients. Nature Machine Intelligence 3, 620–631. — doi.org/10.1038/s42256-021-00343-w
- Lundberg & Lee (2017). A Unified Approach to Interpreting Model Predictions (SHAP; Expected Gradients is its gradient-based relative). NeurIPS. — arxiv.org/abs/1705.07874
- Sturmfels, Lundberg & Lee (2020). Visualizing the Impact of Feature Attribution Baselines. Distill. — doi.org/10.23915/distill.00022
Read next: Integrated Gradients · KernelSHAP / Shapley · Grad×Input / DeepLIFT · Occlusion
