PRISM: Deriving the Transformer as a Signal-Denoising Operator via Maximum Coding Rate Reduction

TL;DR

PRISM derives a white-box Transformer by treating attention as a gradient ascent on the maximum coding rate reduction objective $ΔR(Z)$ from the $MCR^2$ framework. It introduces two geometric biases, an overcomplete dictionary and $π$-RoPE based spectral separation, to enforce non-resonant signal-noise separation and enable de-noising of representations. Experiments on TinyStories with a 50M-parameter Prism-mini show unsupervised functional disentanglement, with low-frequency signal heads capturing long-range dependencies and high-frequency noise heads modeling local syntax, while maintaining competitive performance. The work suggests interpretability and performance can be unified through principled geometric constraints, offering a scalable white-box pathway for Transformer design.

Abstract

Deep learning models, particularly Transformers, are often criticized as "black boxes" and lack interpretability. We propose Prism, a white-box attention-based architecture derived from the principles of Maximizing Coding Rate Reduction ($\text{MCR}^2$). By modeling the attention mechanism as a gradient ascent process on a distinct signal-noise manifold, we introduce two physical constraints: an overcomplete dictionary to expand the representational phase space, and an irrational frequency separation ($π$-RoPE) to enforce incoherence between signal and noise subspaces. We demonstrate that these geometric inductive biases can be viewed as a physical constraint and they are sufficient to induce unsupervised functional disentanglement alone. Using TinyStories as a controlled testbed for verifying spectral dynamics, we observe that Prism spontaneously specializes its attention heads into spectrally distinct regimes: low-frequency heads capturing long-range causal dependencies (signal) and high-frequency heads handling local syntactic constraints (noise). Our results suggest that interpretability and performance are not a trade-off, but can be unified through principled geometric construction.

Figures (3)

• Figure 1: The PRISM White-Box Architecture.(Top) The input latent state $\mathbf{Z}^\ell$ is projected into an overcomplete feature space via dictionary $\mathbf{U}=[\mathbf{U}_s,\mathbf{U}_n]$. The data flow in prism bifurcates into Signal Stream (Blue) and Noise Stream (Orange). $\pi$-RoPE: the Signal stream applies low-frequency rotary embeddings scaled by $\pi$ ($e^{i\pi\theta}$) to represent long-range semantic structures, while the Noise stream applies high-frequency embeddings ($e^{i\pi^{-1}\theta}$) to capture short-range syntactic artifacts. The layer output is computed via a differential operator, $\mathbf{Z}^{\ell+1} \leftarrow S - \lambda N$, where $\lambda$ is an annealing coefficient that dynamically suppresses the noise subspace, encouraging the model to learn long-term correlation (e.g. causal logic) during training.
• Figure 2: Training Dynamics Emergent Functional Specialization(a) Training and validation loss curves on the TinyStories dataset. (b) Evolution of the L2 Gradient Norm. (c)-(d) Comparative visualization of attention maps in Layer 6, Head 2. (c) Signal Head: Driven by low-frequency $\pi$-RoPE, it captures sparse, long-range semantic dependencies (e.g., attending to specific entities). (d) Noise Head: Driven by high-frequency $\pi$-RoPE, it captures local syntactic artifacts and background noise.
• Figure 3: Overcomplete Signal Subspace Visualization:(a) Subject-centric: Maintains a persistent vertical attention band on the word "frog", tracking the narrative protagonist. (b) Object-centric (Causal): When predicting the target token "door" (last row), this head explicitly attends to the distal causal antecedent "key", bridging the logical gap between the tool and the target. (c) Structural: Attends to the determiner "a", likely tracking the introduction of new entities. (d) Action-centric: Anchors the sequence to the primary predicate "found" and "looked", tracing the action.