Rewrite dataset attribution storage + 3 metrics

[ocg-goodfire]

Summary

• Dict-of-dicts storage: Replace flat-indexed source_to_component / source_to_out_residual matrices with attrs[target_layer][source_layer] = Tensor[target_d, source_d]
• Canonical naming: Storage uses canonical layer names (embed, output, 0.glu.up). Harvester stays concrete internally; translation at storage boundary via topology.target_to_canon()
• 3 attribution metrics: attr (signed mean), attr_abs (attribution to |target|, via backprop through .abs()), mean_squared_attr (pre-sqrt, mergeable across workers)
• Split entrypoints: Separate run_worker.py and run_merge.py instead of combined run.py
• Bug fixes: alive_targets bool→index, embed CI KeyError, scatter_add OOB (vocab_size vs num_embeddings), merge entrypoint missing config_json, attr_abs sign correctness (backprop through |y| instead of flipping by source sign)

Test plan

• Storage unit tests (17 tests: has_source, has_target, save/load roundtrip, merge weighted average)
• Type checks pass (basedpyright)
• End-to-end SLURM test on s-17805b61 (3 GPUs × 2 batches) — workers complete, merge pending with more memory
• Benchmark: fast (scatter_add_) vs slow (per-element loop) comparison in progress

🤖 Generated with Claude Code

[ocg-goodfire]

@claude please review this PR. specifically the exact correctness of the code. The interfaces and types are of course different but that's obviously intentional, I want to know if I've made mistakes implementing what should hopefully be somewhat obvious what the goal was

[claude]

Claude finished @ocg-goodfire's task —— View job

---

Code Review Complete ✅

Reviewing the dataset attribution storage rewrite for correctness. Todo list:

• Understand the new dict-of-dicts storage format vs old flat matrices
• Review attribution metric calculations (attr, attr_abs, mean_squared_attr)
• Check storage serialization/deserialization logic
• Review worker/merge split implementation
• Look for potential bugs and edge cases
• Verify type correctness and consistency

Review Summary

The code looks correct and well-implemented. The rewrite successfully addresses the stated goals while fixing several important bugs. Here's my detailed analysis:

✅ Storage Architecture

• Dict-of-dicts format: Much cleaner than flat matrices. attrs[target_layer][source_layer] = Tensor[target_d, source_d] is intuitive.
• Canonical naming: Smart translation boundary - concrete paths internally, canonical names (embed, output, 0.glu.up) for storage.
• Output residual storage: Elegant solution storing attributions to pre-unembed residual rather than full vocab, reducing storage from O(V×C×V) to O(V×C×d_model).

✅ Attribution Metrics

All three metrics are mathematically correct:

• attr: E[∂y/∂x · x] - standard signed attribution
• attr_abs: E[∂|y|/∂x · x] - correctly computed via backprop through .abs(), not sign flipping
• mean_squared_attr: E[(∂y/∂x · x)²] - stored pre-sqrt for proper merging across workers

✅ Bug Fixes Verified

• alive_targets bool→index: Line 269 correctly uses torch.where(alive_targets)[0].tolist() instead of boolean indexing
• scatter_add_ bounds: Uses vocab_size instead of num_embeddings preventing OOB errors
• attr_abs sign: Correctly computes via abs_grad * act * source_ci from backprop through .abs()

✅ Worker/Merge Implementation

• Clean separation: run_worker.py for computation, run_merge.py for combining
• Weighted averaging: Merge correctly weights by n_tokens, promotes to double precision to avoid numerical issues
• Robust handling: Proper validation of CI thresholds and layer consistency across workers

✅ Storage & Serialization

• Save/load roundtrip: Comprehensive unit tests verify correctness
• CPU tensor storage: Proper device handling
• Type safety: Good use of Pydantic and jaxtyping

✅ Performance Optimizations

• scatter_add_ for embedding tokens vs vectorized .add_() for components (mentioned 14x speedup)
• Summed gradients: Reduces backward passes from O(positions × components) to O(components)

Minor Observations

1. has_source/has_target implementation (lines 82-102 in storage.py): The logic correctly handles different layer types and bounds checking.

2. Gradient connectivity discovery remains unchanged and looks solid - correctly tests pairs and handles unused gradients.

3. Unit test coverage is excellent (17 tests covering storage, merge, bounds checking).

Verdict: LGTM 🚀

This is a solid rewrite that improves both the interface and fixes real bugs. The storage format is more intuitive, the attribution calculations are mathematically sound, and the implementation handles edge cases properly. The extensive unit tests give good confidence in correctness.