PhaseSAE: Time-Division Multiplexed Dictionary Learning for Memory-Efficient Mechanistic Interpretability

Abstract

Sparse Autoencoders (SAEs) resolve polysemanticity in Large Language Models (LLMs) by projecting dense residual activations $x \in \mathbb{R}^{d_{in}}$ into high-dimensional sparse latent spaces $\mathbb{R}^{d_{sae}}$ ($d_{sae} \gg d_{in}$). However, all prevailing SAE architectures (Top-$K$, Gated, JumpReLU) treat dictionary learning as an instantaneous spatial compression problem: evaluating the entire dictionary in a single forward pass.

During autoregressive generation ($B=1$), this imposes an intractable memory-bandwidth wall: dragging gigabytes of dictionary weights across the memory bus at an arithmetic intensity of only 0.5 FLOP/byte, leaving GPU Tensor Cores >99% idle.

PhaseSAE introduces Time-Division Multiplexing (TDM) to dictionary learning. By partitioning the dictionary tensor into $P$ phase tiles and evaluating cascaded micro-ticks against progressive residual errors, PhaseSAE strictly reduces peak parameter memory bandwidth to $1/P$ while eliminating greedy feature competition. In head-to-head evaluations on TinyStories-1M (blocks.0.hook_resid_post), PhaseSAE achieves 40.51% Cross-Entropy loss recovery vs 24.42% for baseline Top-$K$ at identical sparsity ($L_0 = 16.0$)—a +65.9% relative gain—with Explained Variance ($R^2$) improving by +65.7% ($0.1758 \to 0.2913$).

Executive Summary & Performance Benchmark

Figure 1: Time-Division Multiplexed dictionary partitioning into P phase tiles with cascaded residual micro-ticks and strict 1/P peak memory footprint.

1. The Memory-Bandwidth Wall in Monolithic SAEs

In autoregressive token generation ($B=1$), the arithmetic intensity (FLOPs per byte transferred) of an SAE forward pass is fundamentally bounded:

On an NVIDIA H100 GPU (3,350 GB/s HBM3 bandwidth, 1,979 TFLOP/s FP16 Tensor Core throughput), the roofline saturation threshold requires an arithmetic intensity of $\approx 590 \text{ FLOP/byte}$. Consequently, compute cores sit >99% idle, completely throttled by memory bus latency.

Figure 2: Memory-Bandwidth Wall & Roofline Saturation curves for Monolithic SAE vs. Fused PhaseSAE.

2. Mathematical Formalization of PhaseSAE

PhaseSAE partitions the latent dimension $d_{sae}$ into $P$ distinct phase tiles ($m = d_{sae} / P$). At micro-tick $p$, the network processes only the parameter slice $W_{enc}^{(p)} \in \mathbb{R}^{d_{in} \times m}$ and $W_{dec}^{(p)} \in \mathbb{R}^{m \times d_{in}}$ against the progressive residual error vector $r^{(p)}$:

Orthogonal Feature Complement Alignment

In standard Top-$K$ SAEs, selecting all $K$ latents simultaneously causes greedy feature competition: latents with overlapping decoder directions compete destructively, producing overshoot or mutual suppression.

In PhaseSAE, Phase 0 isolates broad high-variance features; subsequent phases only observe residual error $r^{(p)}$, forcing latents to span the orthogonal complement of previously explained variance. Empirical cross-phase decoder cosine similarity confirms near-orthogonal alignment:

3. Empirical Gauntlet on TinyStories-1M

We evaluated PhaseSAE head-to-head against standard Top-$K$ SAEs directly on activations from TinyStories-1M (blocks.0.hook_resid_post, $d_{in}=64$, $d_{sae}=256$, target $L_0=16.0$, 16,000 tokens of c4-10k) using the official sae_lens.evals gauntlet:

Figure 3: Downstream Cross-Entropy Recovery vs. Active Parameter Footprint across Sparsity Budgets.

4. Anytime Early-Exit Inference

Because residual errors contract monotonically at each micro-tick (41.5% → 71.0% → 84.0% → 90.5%), PhaseSAE natively enables Adaptive Early-Exit Inference: terminating evaluation once $\|r^{(p)}\|_2 / \|x - b_{dec}\|_2 < \tau_{\text{exit}}$:

5. Citation & Reference

@article{kupferberg2026phasesae,
  title={PhaseSAE: Time-Division Multiplexed Dictionary Learning for Memory-Efficient Mechanistic Interpretability},
  author={Kupferberg, Conner},
  year={2026},
  url={https://thalamocorticalsystems.com/papers/phase-sae.html},
  note={Canonical Implementation: github.com/connerkup/phase-sae; Upstream Fork: github.com/connerkup/SAELens}
}