[emerardd]

Description

Residual decomposition with mlp_input=True assumes that the current layer's attention output feeds its MLP. That is correct for sequential blocks, but parallel blocks feed both branches from resid_pre. On a tiny random NeoX with nonzero attention output biases, the existing component decomposition differs from the actual ln2 input by about 0.4; normalized attribution differs by about 23. Full head/neuron decomposition also includes the nonexistent current-layer attention contribution.

This change declares the MLP-input topology on BlockBridge / ParallelBlockBridge and uses the target block's topology in the analysis paths:

• Component and full residual decompositions include current-layer attention/head writes only for sequential blocks.
• accumulated_bias() excludes the current attention output bias for parallel MLP inputs, while retaining all preceding contributions.
• Accumulated residuals omit nonexistent parallel mid streams.
• Neuron remainders target parallel resid_pre rather than an absent resid_mid, including the fused normalized projection path.

Model forward computation, hook naming, neuron selection and ordinary attention-input decomposition are unchanged. The existing method docstrings now describe parallel MLP-input semantics.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Validation

The initial real-config NeoX regression had 16 failed / 16 passed before the fix: the parallel cases failed and the sequential controls passed. These models are built from config and do not download weights.

Expanded coverage checks layers 0/1, raw/normalized sums, batched/batchless caches, expanded/unexpanded neurons, vector/matrix projections, bias inclusion and remainder targets. Cross-architecture checks compare decomposition sums to captured normalization inputs, including shared-ln1 and separate-ln2 layouts.

• NeoX regression, mixed-norm decomposition and two existing activation-cache test files: 362 passed.
• Final parallel/sequential architecture identities and expanded NeoX regression: 108 passed (overlaps the NeoX cases above).
• Full mypy .: 389 source files passed.
• Pinned pycln/isort/Black checks on all five changed files and git diff --check: passed.

The complete unit suite, pretrained checkpoints and GPU/MPS hardware were not tested.

Checklist

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (API docstrings)
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes (affected tests run; full suite not run)
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[jlarson4]

This is a bug that existed before this issue, but it is related, would you be willing to fix this here as well? Most parallel blocks have no ln2, because ln1 feeds both branches. So apply_ln=True with mlp_input=True still raises a KeyError for ln2.hook_scale here, in decompose_resid, and in get_full_resid_decomposition. Can we adjust the cached-scale lookup to resolve to the norm that feeds the MLP, so the normalized paths work on those blocks too?

[emerardd]

Thanks for catching this! I'll address it in this PR by resolving the cached scale from the norm that actually feeds the MLP, keeping independent ln2 norms where present and using ln1 for shared-norm parallel blocks. I'll also add coverage for the normalized decomposition paths on those shared-norm layouts.

[emerardd]

Fixed in 558597d. The cached-scale lookup now uses ln1 for shared-norm parallel blocks while preserving ln2 for blocks with an independent MLP norm. Missing cached ln2 scales still raise rather than falling back to ln1. I also added normalized decomposition coverage for the shared-norm layouts and a two-norm Falcon case.

[jlarson4]

Good solution here @emerardd, just one comment about a related bug that would be nice to fix along side it. If you don't have time to address that as well right now, let me know and I can merge this + file that as a separate issue for someone else to take on.

[jlarson4]

Looks good! I'll merge once CI passes

[jlarson4]

Update: there is a bug with the Qwen3.5 testing model we use for CI that is preventing Full Code Coverage from passing, I will patch dev and push the patch to your branch when it is verified.