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

Conner Kupferberg · September 2026 Classification: Deep Learning Architecture · Mechanistic Interpretability · Hardware Systems Optimization Canonical Reference Codebase: connerkup/phase-sae Upstream Distribution Fork: connerkup/SAELens (feat/phase-multiplexed-sae) Canonical Paper & Interactive Press Suite: thalamocorticalsystems.com/papers/phase-sae

> 🏛️ Official Publication Notice: The canonical research paper, Triton GPU rooflines, interactive figures, and downloadable benchmark datasets have been migrated to Conner Kupferberg’s Frontier AI lab: [Thalamocortical Systems](https://www.thalamocorticalsystems.com/papers/phase-sae).

Executive Summary

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 architecture hits 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 achieves three foundational results:

1. The Memory-Bandwidth Wall

In autoregressive 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), saturating compute requires an arithmetic intensity of $\approx 590\text{ FLOP/byte}$. At 0.5 FLOP/byte, execution is >99% memory-bound.

2. PhaseSAE: Time-Division Multiplexed Formulation

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)}$:

In standard Top-$K$ SAEs, selecting all $K$ latents simultaneously causes greedy feature competition: latents with overlapping decoder directions compete destructively. 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 ($\mathbb{E}[|\cos (W_{dec}^{(p)},W_{dec}^{(q)})|]\approx 0.11$).

3. Hardware Roofline & Fused Triton Kernel

To avoid the latency penalty of sequential PyTorch kernel launches, PhaseSAE implements a custom Fused TDM Triton GPU Kernel. Residual vectors $r^{(p)}$ are pinned directly in GPU registers, eliminating intermediate DRAM roundtrips across all $P$ micro-ticks:

At $B=16$, the Fused TDM Triton kernel matches monolithic wall-clock latency (41.7 ms vs 40.3 ms) while maintaining a strict 4× reduction in peak parameter bandwidth and achieving +65.9% higher downstream cross-entropy recovery.

4. Downstream LM Benchmark: 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 evaluation gauntlet:

4.1 Pareto Frontier Sweep

To confirm that PhaseSAE's dominance is structural across sparsity regimes, we evaluated both architectures across varying sparsity budgets $K\in \{8,16,24\}$:

Pareto Invariant: PhaseSAE strictly dominates monolithic Top-$K$ across every operating regime, achieving higher downstream fidelity at $K=16$ (44.52% CE recovery) than baseline Top-$K$ achieves even at $K=24$ (38.61%), with 4× lower peak parameter bandwidth.

4.2 Optimization Dynamics & Residual Contraction

Gradient backpropagation and loss contraction across 200 optimization steps:

Hierarchical residual contraction across progressive micro-ticks on 256 unseen validation vectors:

5. The Anytime SAE: Dynamic Token Early-Exit

Because residual errors contract monotonically at each micro-tick, PhaseSAE natively enables Adaptive Early-Exit Inference: terminating the forward pass once $\Vert r^{(p)}{\Vert }_2/\Vert x-b_{dec}{\Vert }_2<{\tau }_{\text{exit}}$:

Empirical adaptive exit performance across varying thresholds $\tau$:

6. Implementation & Framework Integration

PhaseSAE is implemented as both a standalone PyTorch reference library (connerkup/phase-sae) and a fully registered upstream architecture within sae_lens (feat/phase-multiplexed-sae):

from phase_sae import PhaseMultiplexedSAE, PhaseSAEConfig, AnytimeSAE

# 1. Initialize PhaseSAE (P=4 phases, k=4 per phase -> total k=16)
cfg = PhaseSAEConfig(d_in=768, d_sae=12288, num_phases=4, k_per_phase=4)
sae = PhaseMultiplexedSAE(cfg)

# 2. Forward pass with in-register residual streaming
reconstruction, feature_acts = sae(x)

# 3. Stream micro-ticks for 1/P memory consumption
for phase_idx, acts_p, recon_p, residual in sae.stream_phase_ticks(x):
    print(f"Phase {phase_idx}: residual norm = {residual.norm(dim=-1).mean():.4f}")

# 4. Adaptive early-exit under compute budget
anytime = AnytimeSAE(sae)
out = anytime.run_under_budget(x, max_phases=2, exit_threshold=0.2)
print(f"Exited after {out.phases_executed} phases (early_exit={out.early_exited})")

Test Matrix & CI Status: • tests/saes/test_phase_multiplexed_sae.py: 8/8 tests passed in 0.19s. • tests/test_llm_sae_training_runner.py: 2/2 tests passed in 8.89s. • Standalone Triton & Anytime test suite: 6/6 tests passed. • Regression suite across Gated and JumpReLU SAEs: 51/51 passed with zero regressions.

7. Research Brief & Press Package

For journalists, researchers, and engineers reviewing or covering PhaseSAE, a complete media and data package is available for download:

Citation & Attribution

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