Scaling Interpretable Transformers with Parity Bottleneck Layers

Mechanistic interpretability often relies on Sparse Autoencoders (SAEs) to recover features from a model's residual stream post-hoc. However, it remains unclear whether these recovered features are truly the ones the model uses during computation or artifacts of the training process. The authors introduce the ParityTransformer, which aims to solve this by enforcing interpretability through a Deep Parity Bottleneck (DPB) integrated directly into the model's forward pass. Instead of storing a massive, learned dictionary of features in memory, the DPB uses a lightweight, deterministic hash function to compute feature directions on-chip. This allows the model to maintain a wide, sparse basis without the memory bandwidth bottlenecks that typically prevent per-layer interpretable bottlenecks at scale.

The DPB architecture uses a hierarchical mixture-of-experts approach to select active features. By using a parity-based hash function, the model derives feature directions from their integer indices using a small seed matrix that fits entirely in GPU registers. This eliminates the need to load large dictionary matrices from High Bandwidth Memory (HBM). The authors demonstrate that this approach is computationally feasible, though it introduces an 'interpretability tax.' While the ParityTransformer matches the performance of dense baselines on standard benchmarks like HellaSwag and LAMBADA, it requires significantly more training tokens and wall-clock time to reach equivalent loss levels.

Empirical results show that ParityTransformers perform as well as, or better than, post-hoc SAEs on sparse probing tasks. Because the features are native to the model's forward pass, they provide a more direct window into the model's actual computational process. The authors show that these features are effective for steering and fine-grained causal interventions, suggesting that designing models with interpretable-by-design bottlenecks is a viable alternative to post-hoc recovery methods.