The Dual-Stream Transformer: Channelized Architecture for Interpretable Language Modeling

TL;DR

The paper introduces the Dual-Stream Transformer, a novel architecture that separates token and context processing to enhance interpretability, while maintaining competitive performance on language modeling tasks.

Contribution

It proposes a dual-stream residual architecture with tunable head mixing strategies, balancing interpretability and performance in language models.

Findings

Fully independent head mixing increases validation loss by 8%.

Kronecker mixing costs only 2.5% performance loss.

Models remain robust under attention amplification up to 16x.

Abstract

Standard transformers entangle all computation in a single residual stream, obscuring which components perform which functions. We introduce the Dual-Stream Transformer, which decomposes the residual stream into two functionally distinct components: a token stream updated by attention and a context stream updated by feed-forward networks. Information flow between attention heads is controlled through a hierarchy of mixing strategies, from fully independent (maximum interpretability) to dense (standard transformer behavior). This design exposes a tunable tradeoff between interpretability and performance. We measure this tradeoff on language modeling tasks at 29M parameters. Fully independent head mixing increases validation loss by 8\% relative to dense baselines. The recommended Kronecker mixing strategy, which permits scalar communication between heads while preserving within-head structure, costs only 2.5\%. All configurations maintain functional generation under attention amplification (scaling logits by factors up to 16 at inference time), with degradation ranging from 16\% to 27\%. This robustness suggests the architectures learn discrete algorithms that operate independently of soft probabilistic mixing. The architecture provides a foundation for interpretable language models where internal structure is exposed by design. \footnote{This work was partially supported by DARPA Contract HR001125C0302.}