[Proposal] Contextual Decomposition: Nonlinear Source Attribution and Dataset-Wide Ranking on GPT-2

Read the current local single-source implementation and its 218 focused CD tests, then investigate per-example reconstruction budgets and independently test the allocation policy. The remaining work includes maintainer agreement, resolving integration validation, and implementing reference-mean fitting, supported joint-source attribution, and held-out dataset-wide ranking; the issue is done only when both outcomes are delivered and validated.

---

Proposal

Add maintained, Bridge-native contextual decomposition for GPT-2, with explicit
nonlinear allocation semantics and independently tested constituent shares.

This issue has two usable outcomes:

1. Single-source attribution from a head/head-position or residual anchor to
   downstream internal activations and logits/logit differences.
2. Reference-mean fitting, supported joint-source attribution, and dataset-wide
   independent source ranking with held-out characterization.

The first outcome is implemented locally but is not pushed or merged. The second
is not implemented. Keep this issue open until both outcomes are delivered; do
not close it after preparatory refactors or single-source delivery alone.

---

Motivation

TL has activation access, patching, gradient attribution, and vocabulary/weight
readouts, but no first-class nonlinear contextual decomposition workflow. CD-T
assigns beta/gamma constituents through attention, normalization, MLPs, and
readout without differentiating a task metric.

This complements intervention tools; it does not replace causal validation.
Reconstruction, policy-correct allocation, and causal usefulness are separate
claims. Independent constituent tests are necessary because defining gamma as
total minus beta can reconstruct a wrong allocation perfectly.

---

Current implementation and mathematical conventions

The local single-source API provides ContextualDecomposition, CDSource,
CDTarget, CDBaseline, CDTensor, CDResult, and reconstruction diagnostics.

• Raw fp32/eval GPT-2 TransformerBridge with local transformers driver and eager
  full-sequence attention, on one device.
• A selected head's residual write excludes shared output bias. Reference-centered
  beta is injected at the attention merge; gamma retains all other native terms.
• Residual pre/mid/post anchors and validated absolute source positions.
• Downstream internal targets and per-example logit/logit-difference selectors.
• Owned detached references/results and explicit intermediate retention.
• Affine bias allocation by projected magnitude; undefined jointly zero ratio
  completed by assigning bias to gamma.
• Ordered GELU differences, combined-input LayerNorm denominator, and printed
  nonlinear QK/softmax/value allocation. This is not a Shapley or LRP method.
• No silent cancellation redistribution, gamma correction, or frozen-pattern
  substitution.

The implementation policy is named magnitude-affine-ordered-gelu-v1. Printed
paper equations and inspected reference scripts differ in places; complete
reference compatibility is not claimed. No authors' kernels are vendored or
required at runtime. Repository-wide licensing of those kernels was not established.

---

Outstanding blockers identified by review

1. Batch-dependent error budget: native reconstruction currently uses a maximum
   scale over the whole batch. A large example can hide an incorrect smaller
   example that would fail alone. Fix validation per example and propagate
   appropriate target budgets before requesting merge. Do not enlarge the global
   tolerance or correct constituents to hide the problem.
2. Repository integration: full tests fail at Granite's MPS eager-scan parity
   assertion. Investigation locates the discrepancy in HF grouped MoE experts,
   not GPT-2 CD. Track the runtime issue separately; full validation remains blocked.
3. Maintainer agreement: no CD-T policy/API/scope approval has been found.
   This proposal requests it; implementation existence is not approval.

---

Validation observed

On the local feature head:

• 218 focused CD tests passed with 97% statement coverage.
• All five sanitized notebook code cells passed.
• Mypy, format/diff checks, docs build, source distribution, and wheel build passed.
• Standard unit: 7,452 passed; docstrings: 13 passed; acceptance: 163 passed.
• Standard integration: 1,541 passed and one Granite MPS failure.

These results include existing skips/xfails and docs warnings; they are not a
claim of warning-free or fully green validation. Trained GPT-2 characterization
pins revision 607a30d783dfa663caf39e06633721c8d4cfcd7e and uses explicit token IDs,
references, and CPU/fp32/eager settings. It is numerical characterization, not
IOI circuit-recovery evidence.

---

Second outcome and exclusions

The second PR should fit masked/aligned reference means from separate reference
data, implement explicitly supported disjoint joint-source sets, score sources
independently across examples, retain signed/ratiometric diagnostics, and provide
a held-out IOI ranking workflow. It must not synthesize joint attribution by
summing marginal scores or leak test examples into reference fitting.

Automatic circuit discovery/pruning, mixed head/MLP circuits, direct-edge CD,
modern decoder/encoder/seq2seq architectures, cached generation, and resumable
experiments are separate proposals. No SAE training or artifact dependency is needed.

---

Maintainer decisions requested

• Is a single-source GPT-2 CD tool plus a separate ranking/joint-source PR the right scope?
• Are the explicit allocation policy and zero-ratio completion acceptable?
• Is the current source/target/reference API appropriate for tools.analysis?
• Which per-example numerical acceptance contract should be required?
• Is independent equation/oracle validation sufficient without an externally
  licensed reference package, with reference comparisons deferred until permission?

---

References

• Hsu et al., Efficient Automated Circuit Discovery in Transformers using
  Contextual Decomposition: https://arxiv.org/abs/2407.00886.
• Murdoch et al., Beyond Word Importance: Contextual Decomposition to Extract
  Interactions from LSTMs: https://arxiv.org/abs/1801.05453.
• Authors' research scripts: https://github.com/adelaidehsu/CD_Circuit.

---

Checklist

• Fresh repository issue/PR searches found no matching contextual-decomposition proposal.
• Maintainers approve policy/API/scope.
• Batch-dependent reconstruction acceptance fixed and independently tested.
• Single-source PR merged with required validation resolved.
• Reference/joint-source/dataset-ranking outcome implemented and validated.
• Both outcomes complete before issue closure

---