Temporal Sparse Autoencoders

Chapter 1: The Key Insight

Read this sentence: "Photosynthesis is the process by which plants convert sunlight into energy."

Now ask yourself two questions about each token:

1. What topic is this about? Plant biology. Scientific explanation. Every single token in this sentence belongs to the same topic. The topic is temporally consistent — it doesn't change from one token to the next.
2. What syntactic role does this token play? "Photosynthesis" is a capitalized proper noun. "is" is a copula verb. "the" is a determiner. "by" is a preposition. These roles change every single token. Syntax is temporally local.

This isn't a new idea in linguistics. Chomsky (1965) distinguished syntactic and semantic structure. Griffiths et al. (2004) formally argued that semantics exhibit long-range dependencies while syntax depends on short-range interactions. What's new is using this insight as a training objective for SAEs.

The Analogy: Radio Tuning

Think of an LLM's hidden representation at each token as a radio signal. It carries two overlapping broadcasts simultaneously:

• The FM station (semantics): A slow-changing carrier wave. The topic, the intent, the style. It changes when the passage shifts from biology to history, but within a passage it's roughly constant.
• The AM modulation (syntax): Rapid fluctuations riding on top of the carrier. Noun, verb, determiner, punctuation. These change every token.

A standard SAE mixes them together. It has no reason to separate them. T-SAEs add one constraint: some features must look like the FM station — smooth and consistent over time. The rest automatically become the AM modulation.

What Makes This Different From Prior Work

Several prior approaches tried to improve SAEs: Matryoshka SAEs (Bussmann et al., 2025) learn hierarchical features but don't enforce temporal structure. Transcoders (Paulo et al., 2025) learn causal features but still treat tokens as i.i.d. BatchTopK (Bussmann et al., 2024) improves sparsity but doesn't address the semantic/syntactic split.

T-SAEs are the first to inject a structural prior from linguistics into SAE training: the temporal consistency of meaning. The modification is a single contrastive loss term added to the standard objective. Everything else — the encoder, decoder, sparsity — stays the same.

Chapter 2: The Data Generating Process

Before building the T-SAE architecture, the authors formalize why temporal consistency should work. They propose a simple model of how language is generated, then show what it implies for dictionary learning.

The Speaker Model

Imagine a speaker producing a sequence of tokens τ₁, ..., τ_T. At each timestep t, the speaker's word choice is controlled by several factors:

• h_t — high-level latent variables: The topic ("plant biology"), the intent ("explaining a concept"), the audience ("undergraduate students"), the style ("textbook prose"). These are properties of the passage, not the individual token. They change slowly.
• l_t — low-level latent variables: The grammatical role of this specific word (noun, verb, article), the word's position in the sentence, spelling conventions. These are properties of this token specifically. They change every timestep.

The speaker produces each token as a function of context and both latent types:

where τ^t-1 = (τ₁, ..., τ_t-1) is all previous tokens. Think of h_t as "what to say" and l_t as "how to say this particular word."

Assumption 1: Temporal Consistency

The first key assumption formalizes what we just described intuitively:

This doesn't mean h_t is exactly the same at every position — it means it changes slowly. A paragraph about biology has the same topic at token 1 and token 50. The topic might shift at the paragraph boundary (say, to history), but within a passage it's approximately constant.

Assumption 2: Hierarchical Representation

The second assumption says that the LLM's internal representation x_t encodes both h_t and l_t, and they're hierarchical:

Why is this hierarchical? Because the high-level features do most of the reconstruction work. If you know the topic is "plant biology," you can already predict a lot about the activations. The low-level features clean up the residual — they capture the stuff that the topic alone can't explain (like whether this specific token is a noun or a verb).

What This Implies for SAE Design

If these assumptions hold (and the experimental evidence strongly supports them), then the optimal SAE training strategy has two parts:

1. Learn high-level features first by encouraging them to be temporally consistent. These features capture semantics and should be able to reconstruct x_t approximately.
2. Let low-level features learn the residual. Don't constrain them — they'll naturally capture whatever the high-level features miss, which is per-token syntactic information.

This is exactly the Matryoshka-style hierarchy (Bussmann et al., 2025), but with a crucial addition: the temporal contrastive loss on the high-level split. Without it, both splits learn the same kind of features. With it, they specialize.

A Worked Example

Consider a 768-dimensional activation vector x_t from layer 12 of Pythia-160m. The T-SAE has m = 16,384 features, split 20/80: h = 3,277 high-level features and 13,107 low-level features.

At token t = "Photosynthesis" and token t+1 = "is":

• The high-level features f_0:h(x_t) and f_0:h(x_t+1) should be similar — both encode "plant biology discussion." Their cosine similarity should be high (say, 0.85).
• The low-level features f_h:m(x_t) and f_h:m(x_t+1) will be different — "Photosynthesis" (capitalized noun) vs. "is" (copula verb). No consistency constraint applies.
• The high-level features alone reconstruct x_t with error ≤ ε. Adding the low-level features drives the error to zero.

Chapter 3: The Architecture

T-SAEs modify the standard SAE architecture in exactly two ways: (1) partition the feature space into high-level and low-level splits, and (2) add a contrastive loss on the high-level split. Let's walk through every component.

The Encoder

The encoder takes a model activation x_t ∈ R^d (e.g., d = 768 for Pythia-160m, d = 2304 for Gemma2-2b) and produces a sparse feature vector f(x_t) ∈ R^m (e.g., m = 16,384 features):

where W^enc ∈ R^m×d is the encoder weight matrix, b^enc ∈ R^d is the encoder bias, and σ is the activation function. The paper uses BatchTopK activation (k = 20), meaning only the top 20 features (out of 16,384) are non-zero for each token. This enforces strict sparsity.

The Feature Split

The m features are partitioned into two groups:

• High-level features f_0:h(x_t) — the first h indices. These are the features that will be encouraged to be temporally consistent. In practice, h = 0.2m (20% of features).
• Low-level features f_h:m(x_t) — the remaining m − h indices. These are unconstrained and learn to capture the residual (syntax, per-token patterns).

For a 16k SAE: h = 3,277 high-level features, 13,107 low-level features. The split ratio (20/80) is a hyperparameter — ablations in Section 4.6 of the paper show that 10/90 and 20/80 both work well, while 50/50 hurts syntax recovery.

The Decoder

The decoder reconstructs x_t from the sparse features:

where W^dec ∈ R^d×m is the decoder matrix and b^dec ∈ R^d is the decoder bias. The decoder is split to match the feature partition:

• W^dec_0:h ∈ R^d×h — decoder columns for high-level features
• W^dec_h:m ∈ R^d×(m-h) — decoder columns for low-level features

Their concatenation equals the full decoder: W^dec = [W^dec_0:h | W^dec_h:m].

Full Architecture Diagram

Implementation Details

python
import torch
import torch.nn as nn

class TemporalSAE(nn.Module):
    def __init__(self, d_model, n_features, n_high, k=20):
        super().__init__()
        self.n_high = n_high          # h: number of high-level features
        self.k = k                     # BatchTopK sparsity
        self.encoder = nn.Linear(d_model, n_features)
        self.decoder = nn.Linear(n_features, d_model)

    def encode(self, x):
        # x: (batch, d_model) -> f: (batch, n_features)
        pre_act = self.encoder(x)      # (B, m)
        # BatchTopK: keep only top-k activations per sample
        topk_vals, topk_idx = torch.topk(pre_act, self.k, dim=-1)
        f = torch.zeros_like(pre_act)
        f.scatter_(-1, topk_idx, topk_vals)
        return f

    def decode_high(self, f):
        # Reconstruct using only high-level features
        f_high = f.clone()
        f_high[:, self.n_high:] = 0   # zero out low-level
        return self.decoder(f_high)

    def forward(self, x):
        f = self.encode(x)             # (B, m), sparse
        x_hat = self.decoder(f)        # full reconstruction
        x_hat_high = self.decode_high(f)  # high-level only
        z = f[:, :self.n_high]         # high-level features for contrastive
        return x_hat, x_hat_high, z

Chapter 4: The Contrastive Loss

This is the heart of the paper. Everything else — the feature split, the Matryoshka reconstruction, the architecture — is borrowed from prior work. The contrastive loss is what makes T-SAEs work.

Manufacturing the Need

Suppose you just train with the Matryoshka reconstruction loss (L_H + L_L) without any temporal constraint. What happens? Both feature groups learn the same kind of features. The high-level split has no reason to prefer smooth, semantic features over noisy, syntactic ones. The split is meaningless.

We need a loss term that says: "the high-level features for token t should be similar to the high-level features for token t−1." But we also need to prevent collapse — we don't want ALL tokens in the batch to have the same high-level features (that would just be a constant bias term).

This is exactly the setup for contrastive learning: pull positive pairs together, push negative pairs apart.

Defining the Loss

Let z_t = f_0:h(x_t) be the high-level features of token t. Let s(x, y) be the cosine similarity between x and y. We define:

The contrastive loss is a symmetric InfoNCE-style objective:

Let's break this down piece by piece:

• First term: For each sequence i, how well can we identify the correct z_t-1⁽ⁱ⁾ among all the z_t-1^(j) in the batch? This is a "which previous token came from the same sequence?" classification problem.
• Second term: The symmetric version — for each z_t-1⁽ⁱ⁾, can we identify the correct z_t⁽ⁱ⁾ among all z_t^(j)?
• Symmetry prevents cheating: Without both terms, the model could collapse one side to a constant. The symmetry ensures both z_t and z_t-1 carry meaningful information.

Deriving Why This Works: A Worked Example

Let's trace through a concrete batch with N = 4 sequences:

Cross-sequence similarities (negatives) should be low:

• s(z_t⁽¹⁾, z_t-1⁽²⁾) = s("plants" features, "Napoleon" features) ≈ 0.1
• s(z_t⁽¹⁾, z_t-1⁽³⁾) = s("plants" features, "derivative" features) ≈ 0.15
• s(z_t⁽¹⁾, z_t-1⁽⁴⁾) = s("plants" features, "verdict" features) ≈ 0.05

For sequence 1, the first term becomes:

If the model improves the positive similarity to 0.95 while keeping negatives at 0.1:

The loss decreases. The gradient pushes the encoder to make same-sequence high-level features more similar and cross-sequence features less similar.

The Full Loss

Combining everything, the total T-SAE loss is:

where:

• L_matr = L_H + L_L — the Matryoshka reconstruction loss
• L_H = ‖x_t − W^dec_0:h f_0:h(x_t) + b^dec‖₂² — high-level reconstruction
• L_L = ‖x_t − W^dec f(x_t) + b^dec‖₂² — full reconstruction
• L_contr — the symmetric contrastive loss on z_t = f_0:h(x_t)
• α = 1.0 — contrastive weight (no tuning needed!)

Training Data: Loading Activation Pairs

In practice, activations are loaded as pairs (x_t, x_t-1) — adjacent tokens from the same sequence. The pairs are shuffled across the batch so that negative pairs come from different sequences. This means each training batch is 2× the normal size (each sample is a pair), which reduces the effective batch size by half for the same memory budget.

python
def contrastive_loss(z_t, z_prev):
    # z_t, z_prev: (N, h) high-level features for adjacent tokens
    # Compute pairwise cosine similarities
    z_t_norm = z_t / z_t.norm(dim=-1, keepdim=True)
    z_prev_norm = z_prev / z_prev.norm(dim=-1, keepdim=True)
    sim = z_t_norm @ z_prev_norm.T      # (N, N)

    # Positive pairs are on the diagonal: sim[i, i]
    # Row-wise: for each z_t^(i), classify correct z_prev
    labels = torch.arange(z_t.size(0), device=z_t.device)
    loss_row = nn.functional.cross_entropy(sim, labels)

    # Column-wise: for each z_prev^(i), classify correct z_t
    loss_col = nn.functional.cross_entropy(sim.T, labels)

    return (loss_row + loss_col) / 2

def tsae_loss(x_t, x_prev, model, alpha=1.0):
    x_hat_t, x_hat_high_t, z_t = model(x_t)
    x_hat_prev, x_hat_high_prev, z_prev = model(x_prev)

    L_H = (x_t - x_hat_high_t).pow(2).sum(-1).mean()
    L_L = (x_t - x_hat_t).pow(2).sum(-1).mean()
    L_matr = L_H + L_L

    L_contr = contrastive_loss(z_t, z_prev)
    return L_matr + alpha * L_contr

Chapter 5: Disentanglement

The whole point of T-SAEs is to separate semantic from syntactic features. But does it actually work? And how do we even measure "disentanglement"?

Probing for Semantics, Context, and Syntax

The paper uses linear probes — simple logistic regression classifiers trained on top of SAE features — to measure what information each feature split encodes.

The setup: take MMLU questions (multiple-choice academic questions from 57 subjects), encode the last 20 tokens of each question through the LLM, extract SAE features, and train probes to predict:

• Semantics: The question's subject category ("High School European History," "Professional Medicine"). This is a sequence-level property — it doesn't change across the 20 tokens.
• Context: The question ID (which specific question this is). This is also sequence-level but finer-grained than category.
• Syntax: The part-of-speech tag of each token (noun, verb, adjective, etc.). This changes every token.

The Results

And that's exactly what happens. Using Gemma2-2b SAEs on MMLU:

The T-SAE beats baselines on semantics by 9+ percentage points and on context by 8+ points. It's slightly worse on syntax — but that's expected and actually desirable, because the high-level features are now capturing semantics instead of wasting capacity on syntax.

The TSNE Visualization

The paper's Figure 2 provides a stunning visual confirmation. When you plot T-SAE high-level feature activations in 2D (via TSNE) and color by question category, you see clean clusters for each subject. Matryoshka SAE features, plotted the same way, show no clear clustering.

Even more telling: when you color by syntax (part-of-speech), T-SAE low-level features cluster cleanly by POS tag, while T-SAE high-level features don't. The split is working.

Disentanglement Within the Feature Splits

The paper goes further: it probes each split separately. For T-SAEs:

• High-level split alone: High semantics accuracy (0.89), high context (0.93), low syntax (0.72). As expected — these features capture meaning, not grammar.
• Low-level split alone: Lower semantics (0.78), lower context (0.82), but higher syntax (0.83). These features specialize in per-token patterns.

For Matryoshka SAEs, both splits perform similarly on all three tasks — no specialization. The 20/80 partition is arbitrary without the contrastive loss to drive differentiation.

Chapter 6: Experiments

A natural worry: does adding the contrastive loss hurt reconstruction quality? If T-SAEs disentangle features but can't reconstruct the input, they're useless. The paper evaluates five standard SAE metrics.

Standard SAE Metrics

Core Results (Table 1 from the paper)

Reading the Smoothness Metric

The smoothness metric deserves careful explanation. For a sequence of length T and a set of active features, we compute:

This is the maximum absolute change in feature activation, normalized by how much the underlying model activation changed. A smooth feature might have Δ_s = 0.05 (it barely changes even when the model activation changes a lot). A noisy syntactic feature might have Δ_s = 0.5 (it spikes wildly).

T-SAE high-level features have the lowest smoothness score of any method. Their low-level features are appropriately less smooth, confirming that the split works as intended.

What About Fraction Alive?

An interesting detail: T-SAEs have a higher fraction of alive features (0.87 vs. 0.84 on Pythia). This means more of the 16k features activate at least once, suggesting that the contrastive loss encourages the SAE to use more of its feature capacity. Dead features (features that never activate) are wasted parameters — having fewer of them is a benefit.

Chapter 7: The Safety Case Study

Academic metrics are nice, but do T-SAEs actually help with real problems? The paper presents two compelling case studies: understanding alignment datasets and steering model behavior.

Case Study 1: Understanding HH-RLHF

The HH-RLHF dataset (Bai et al., 2022) is Anthropic's human preference dataset used to train safety-focused models. It contains pairs of completions — one "chosen" (preferred by human raters) and one "rejected." The paper asks: what features differentiate chosen from rejected completions?

Method: For each completion pair, compute the difference in mean T-SAE feature activations (rejected − chosen). Features with the largest average difference are the ones that distinguish unsafe from safe content.

What Matryoshka SAEs find for the same analysis: "specific bicycle components," "terms related to data management," "references to ecosystem dynamics and environmental conditions." These are random noise features that happen to correlate with rejected completions for spurious reasons (like response length).

The Spurious Correlation Problem

Here's a subtle but critical finding. Some T-SAE features that show high activation differences are actually spuriously correlated with response length, not with actual safety content. The paper identifies these by computing the Pearson correlation between feature activation difference and response length difference.

The semantically relevant features (green) have low correlation with length — they genuinely capture unsafe content, not just the fact that rejected responses tend to be longer. The length-related features (orange) are spurious confounders. T-SAEs recover both, but critically, the semantically relevant ones are clearly identified and separable.

Case Study 2: Steering

Can T-SAE features be used to steer model generation? The paper intervenes on features during inference by adding α · d_i to the model's residual stream (where d_i is the decoder column for feature i and α is the intervention strength).

The key finding: steering with high-level (semantic) features is dramatically more effective than steering with low-level (syntactic) features.

• High-level steering: Changes the semantics of the generation. The model talks about different topics, with coherent text. Coherence is preserved even at high intervention strengths.
• Low-level steering: Causes token repetition and degenerate text. Because syntactic features are per-token, amplifying them just repeats specific tokens or corrupts grammar.

Chapter 8: Ablations

The paper ablates three key design choices in the T-SAE training pipeline. Each ablation answers a specific question about why the method works.

Ablation 1: Feature Split Ratio

How much of the feature space should be "high-level"? The paper tests 10/90, 20/80, and 50/50 splits.

Ablation 2: Contrastive Window

Instead of contrasting with the immediately previous token (t-1), what if we contrast with a random token from the past context? The paper samples the contrastive partner uniformly from x₁, ..., x_t-1, where the random token x_t-r has r < 25.

Random contrasting boosts context accuracy significantly (+11%) because it encourages longer-range consistency. But it hurts syntax accuracy (−10%) because the low-level features now have less capacity — the high-level features are "greedier," capturing more information.

Ablation 3: No Contrastive Loss (Just Matryoshka)

What if you remove the contrastive loss entirely and just use the Matryoshka reconstruction objective with the 20/80 split?

Without the contrastive loss, semantics drops by 7 points and context by 10 points. The Matryoshka split alone is not enough — the contrastive loss is essential for driving specialization.

The "naive L2 smoothness" ablation replaces the contrastive loss with a simple per-sample L2 penalty: ℓ = α‖z_t − z_t-1‖₂². This enforces smoothness but without the contrastive negative pairs. It performs worse on semantics and much better on syntax — because without negatives to prevent collapse, all high-level features converge to similar values, losing discriminative power.

Chapter 9: Connections

T-SAEs sit at the intersection of three research areas: mechanistic interpretability, contrastive representation learning, and computational linguistics. Let's map the connections and limitations.

Cheat Sheet: Every Key Equation

How T-SAEs Relate to Other Methods

Limitations

• Memory cost: Each batch requires 2x activations (both x_t and x_t-1), halving the effective batch size for the same memory budget.
• Binary split: Only two levels (high/low). Language has more than two timescales — book-level, paragraph-level, sentence-level, token-level. Multiple temporal hierarchies (multi-scale splits) are an obvious extension.
• Contrastive formulation: SAE features are sparse and non-negative. Cosine similarity is designed for dense vectors in R^d. Alternative similarity measures for sparse non-negative spaces might work better.
• Assumes adjacent tokens share semantics: This assumption breaks at topic boundaries (e.g., a paragraph about biology followed by one about history). The paper mitigates this by shuffling pairs, but the assumption is imperfect.
• Only trained at one layer: The paper trains on layer 8 (Pythia) and layer 12 (Gemma). Different layers may benefit from different split ratios or contrastive strategies.

Future Directions

• Multi-scale temporal hierarchy: Instead of 2 splits, use 4+ splits corresponding to book, paragraph, sentence, and token timescales. The contrastive window would differ for each level.
• State tracking: Smooth, temporally consistent features could serve as "state trackers" for monitoring model behavior over long generations, detecting when the model shifts topic or intent.
• Linearized contrastive loss: The paper derives (in Appendix A) a linear approximation to the contrastive loss that exploits the sparse, non-negative geometry of SAE features: −(1/N)∑s_ii + (1/N²)∑(1 + s_ij). This could be faster to optimize and more geometrically appropriate.

Related Lessons

"What I cannot create, I do not understand." — Richard Feynman. T-SAEs let us decompose what LLMs create into parts we can understand: meaning and grammar, separately.