Source-linked AI summary

Beyond Token Scale: Chunk-Level Sparse Autoencoders for Reliable Semantic Feature Discovery

Xu Wang, Yifan Yang, TingHao YU, Difan Zou

arXiv:2609.35521v1cs.CLcs.AIcs.LG

TL;DR

Token-level SAE reconstruction can reward surface details without guaranteeing informative semantic features. The paper introduces chunk-level SAEs that reconstruct observed chunk means, predict independently processed neighbors, or combine both targets. They retain strong target-relative fidelity while showing complementary strengths across semantic discovery and downstream tasks.

  • Problem

    Token-level SAE objectives reward lexical and formatting details alongside semantic content, and faithful reconstruction does not guarantee informative concepts.

    Source: §1 · p. 2
  • Method

    The paper compares chunk-level SAEs that reconstruct a chunk’s mean activation, predict an independently processed same-document neighbor, or combine both targets.

    Source: §1 · p. 2
  • Results

    Chunk-level SAEs retain strong target-relative fidelity and reliable semantic features; Cross-Chunk leads retrieval and transfer, while Mean-Chunk leads cue-free reasoning detection and steering.

    Source: §1 · p. 3, §6 · p. 9
  • Takeaways & Limitations

    The features can support semantic retrieval across different wording and help recognize reasoning relations beyond familiar cue words.

    Source: §5.1 · p. 8, §5.2 · p. 8

Abstract

from arXiv · show

Sparse autoencoders (SAEs) expose features that help us understand and steer language models, but faithful reconstruction does not guarantee informative concepts. Token-level objectives reward lexical and formatting details alongside semantic content, all competing for a limited sparse budget. We introduce a family of chunk-level SAEs that encode mean-pooled activations over chunks, each a contiguous span of tokens: Mean-Chunk reconstructs the observed chunk, Cross-Chunk predicts an independently processed neighbor, and Joint-Chunk combines both targets. These designs separate the effect of a larger observation unit from that of predicting information shared across passages. With matched training data, chunk-level SAEs remain powerful interpretability tools while learning reliable semantic features that capture high-level concepts and respond selectively to relevant content. Their strengths are complementary: Mean-Chunk improves high-level feature discovery, reasoning detection beyond surface cues, and steering; Cross-Chunk leads document retrieval and classification transfer while producing selective, persistent features. Changing what an SAE sees and predicts yields reliable semantic features for more meaningful tasks. We demonstrate their practical value through gains across downstream tasks such as retrieval, reasoning detection, and steering.

1 INTRODUCTION

Token-level reconstruction can reward surface cues without ensuring semantic features, motivating chunk-level SAEs that vary what they encode and predict. These SAEs retain strong target-relative fidelity while showing complementary strengths across semantic feature discovery and downstream tasks.

1 INTRODUCTION

  • Token-level SAEs can mistake lexical cues for reasoning: “wait” appears in both backtracking and non-reasoning contexts, while backtracking can occur without that cue.A semantically coherent feature should recognize plan revision across different wordings and reject isolated surface cues.Source: §1 · p. 2
  • The paper introduces Mean-Chunk, Cross-Chunk, and Joint-Chunk SAEs, which reconstruct a chunk mean, predict an independently processed neighbor, or combine both targets.Using a shared sparse-coding backbone, token stream, and range of chunk lengths makes the effects of pooling and partner prediction explicit.Source: §1 · p. 2
  • Chunk-level SAEs retain strong target-relative fidelity and reliable semantic features; Cross-Chunk leads retrieval and transfer, while Mean-Chunk leads cue-free reasoning detection and steering.Across related passages, Cross-Chunk also achieves the strongest feature persistence and over five times BatchTopK’s dictionary utilization.Source: §1 · pp. 2–3
  • Mean-Chunk more than doubles the high-level feature fraction to 34.8% from BatchTopK’s 17.2%; Cross-Chunk reaches 77.0% Recall@5 versus Temporal’s 64.1%.Mean-Chunk also detects backtracking with 66% recall after reasoning cues are removed, versus under 5% for both token baselines.Source: §1 · p. 3

2 TOKEN-LEVEL SAES FAVOR SURFACE-LEVEL FEATURES

Token-level SAEs encode individual token states and share token reconstruction as their target. Even when objectives add constraints across tokens, they still reproduce individual states, rewarding surface details alongside semantic content.

2 TOKEN-LEVEL SAES FAVOR SURFACE-LEVEL FEATURES

  • BatchTopK and Temporal are representative token-level baselines that reconstruct individual token states under a limited sparse budget.Temporal additionally encourages consistency between neighboring codes.Source: §2 · p. 3
  • BatchTopK reconstructs each contextualized token, while Temporal adds prefix reconstruction and symmetric contrastive alignment between adjacent-token prefix codes.The prefix loss measures how well a designated feature subset reconstructs the current token; the contrastive term encourages similar information across neighboring tokens.Source: §2 · p. 3
  • Both baselines still reproduce individual token states, so their objectives reward surface details alongside semantic content.This motivates observing and predicting longer spans.Source: §2 · p. 4

3 DESIGNING CHUNK-LEVEL SAES TO LEARN MEANINGFUL FEATURES

The design compares chunk-level sparse codes across chunk lengths and three targets: reconstructing the observed chunk, predicting its neighbor, or combining both. Each design uses chunk-mean activations as its starting representation.

3 DESIGNING CHUNK-LEVEL SAES TO LEARN MEANINGFUL FEATURES

  • Adjacent, non-overlapping chunks from the same document are processed independently; varying chunk lengths covers different context scales, and each chunk mean is encoded as the shared starting point.The design compares self reconstruction, neighbor prediction, and a combination of the two targets.Source: §3 · p. 4

3.1 MEAN-CHUNK: RECONSTRUCTING THE CHUNK MEAN

  • Mean-Chunk encodes a chunk’s pooled activation and reconstructs that same mean, giving each sparse code a broader view than token-by-token reconstruction.The design isolates pooling while retaining both semantic content and surface details.Source: §3.1 · p. 4

3.2 CROSS-CHUNK: PREDICTING A SAME-DOCUMENT NEIGHBOR

  • Cross-Chunk replaces self reconstruction with prediction of the neighboring chunk’s mean, averaging prediction in both directions.Both chunks share the encoder and decoder; independent processing prevents the encoder from reading its target passage, making shared topics more useful than passage-specific wording.Source: §3.2 · p. 4

3.3 JOINT-CHUNK: BALANCING SELF AND PARTNER TARGETS

  • Joint-Chunk encodes each observed chunk mean once as one sparse code; the self decoder uses the full code, while the partner decoder uses only its shared prefix.Each prediction error is divided by its own constant-predictor normalizer; α weights the partner term relative to the self term, and the main comparisons use α = 0.25.Source: §3.3 · pp. 4–5

4 CHUNK-LEVEL SAES REVEAL CLEARER SEMANTIC STRUCTURE

The evaluation tests chunk-level SAEs for fidelity, semantic feature discovery, retrieval, and feature structure. They retain target-relative fidelity while exposing high-level concepts and selective, persistent activations.

4.1 EXPERIMENTAL SETUP

  • Across five SAE families, experiments use layer-21 activations from Qwen3.5-9B-Base, one billion Pile token occurrences, and chunk lengths from 32 to 512.Paired chunks are non-overlapping and independently processed; token-level SAEs encode tokens before averaging, whereas chunk-level SAEs average hidden states before encoding.Source: §4.1 · p. 5

4.2 FIDELITY: RETAINING PREDICTABLE INFORMATION

  • Mean-Chunk reaches 93.8% RFVE and leads reconstruction cosine; Cross-Chunk approaches its dense partner reference despite predicting content it never directly observes.Reference-normalized FVE compares an SAE with a reference using the same inputs and targets, accounting for the greater difficulty of partner prediction.Source: §4.2 · pp. 5–6

4.3 INTERPRETABILITY: SELECTIVE EXPLANATIONS AND HIGH-LEVEL CONCEPTS

  • Chunk-level SAEs outperform token-level baselines on same-document Recall@5 beyond lexical overlap, with Cross-Chunk gaining further; the example shares no content words.Figure 5 compares word-set Jaccard matching and SAE representations against a dense hidden-state baseline, and shows a volcanic-activity pair with no shared content words.Source: Figure 5 · p. 7
  • 34.8% of Mean-Chunk coordinates pass the high-level semantic test, versus 17.2% for BatchTopK; the census samples 1,000 coordinates and excludes surface-only rules.A feature passes when at least six of ten strong contexts from distinct held-out documents share a stable semantic or functional concept.Source: §4.3 · p. 7

4.4 FEATURE STRUCTURE: PERSISTENCE AND DICTIONARY UTILIZATION

  • Cross-Chunk has the strongest feature persistence lift and 32.5% dictionary utilization, over five times BatchTopK’s level, with content-selective responses.Persistence lift measures adjacent same-document feature firing relative to shuffled cross-document chunks; utilization reflects how evenly activation mass is spread across active features.Source: Figure 4 · p. 6, §4.4 · p. 7

5 CHUNK-LEVEL SAES IMPROVE DOWNSTREAM UTILITY

Chunk-level SAEs support retrieval, reasoning detection, classification transfer, and steering, with different variants leading different tasks. Cross-Chunk excels at retrieval and classification, while Mean-Chunk leads reasoning detection beyond surface cues and steering.

5.2 REASONING: RECOGNIZING RELATIONS BEYOND SURFACE CUES

  • 66% cue-free backtracking recall for Mean-Chunk compares with under 5% for both token baselines; Temporal leads native reasoning recall.The controls test detection with reasoning structure preserved but surface cues removed, and false activation when cues remain but reasoning structure is removed.Source: Figure 6 · p. 8, §5.2 · p. 8

5.1 RETRIEVAL: CONNECTING PASSAGES BEYOND SHARED WORDS

  • 77.0% document recall for Cross-Chunk exceeds 64.1% for Temporal, the strongest token baseline, when query–partner word overlap is at most 0.1.Candidates are matched to the target’s token length, limiting the usefulness of shared wording and passage length.Source: §5.1 · p. 8

5.3 CLASSIFICATION: FEW LABELS AND FUTURE DOCUMENTS

  • Cross-Chunk leads low-label AUC and future-year OOD accuracy, reaching 82.9% OOD accuracy.The evaluations test probe learning from few examples and transfer to later-year documents, respectively.Source: Figure 7 · p. 9, §5.3 · p. 9

5.4 STEERING: GUIDING GENERATION WITH SEMANTIC FEATURES

  • Mean-Chunk achieves the strongest causal steering score, which accounts for movement toward the intended concept and coherence preservation.The score combines concept direction and coherence preservation; unchanged continuations score 50.Source: Figure 7 · p. 9, §5.4 · p. 9

6 CONCLUSION AND DISCUSSION

The paper concludes that chunk-level SAEs retain target-relative fidelity while producing semantic features useful across tasks. The accompanying appendices specify shared training conditions, representation and objective details, and fidelity and uncertainty procedures.

6 CONCLUSION AND DISCUSSION

  • Chunk-level SAEs retain strong target-relative fidelity while learning semantic features that generalize across wordings and respond selectively to relevant content.Mean-Chunk is strongest for reasoning detection and steering; Cross-Chunk leads retrieval and classification transfer.Source: §6 · p. 9

A TRAINING AND EVALUATION PROTOCOLS

A.1 MODEL, CORPUS, AND TRAINING SCHEDULE

  • All SAE families share a one-billion-token Pile stream and training schedule, while chunk methods use paired non-overlapping chunks of lengths 32–512.The cache includes 2,520,159 training pairs; each chunk is processed separately with position indices reset.Source: §A.1 · p. 14, Table 1 · p. 14

A.2 REPRESENTATIONS, TARGETS, AND JOINT-CHUNK NORMALIZATION

  • Chunk-level codes pool before encoding and thresholding, and retrieval and classification use each observed chunk’s own code.Joint-Chunk uses one full code for self prediction and its first 32,768 coordinates for partner prediction, without a second sparse budget.Source: §A.2 · pp. 14–15

A.3 TOKEN-LEVEL BASELINE OBJECTIVES

  • Token-level baselines reconstruct token states; Temporal additionally aligns neighboring token codes while preserving token reconstruction.Its full-code term includes chunk-first tokens, while the prefix term averages over tokens with a predecessor in the same chunk.Source: §A.3 · pp. 15–16

A.4 FIDELITY REFERENCES AND STATISTICAL SUMMARIES

  • Fidelity scores are normalized against references aligned to each SAE’s inputs, targets, and evaluation examples, and uncertainty summaries use paired resampling.The partner-prediction reference is a dense residual MLP; intervals generally use 10,000 paired bootstrap resamples, with stated exceptions for retrieval and other analyses.Source: §A.4 · p. 16

A.5 RETRIEVAL AND CLASSIFICATION DATA

The appendix describes retrieval galleries and ArXiv classification data used to assess SAE representations. It also reports Joint-Chunk’s lower effective L0 and higher empty-code rate across its full-code and partner-prefix views.

A.5 RETRIEVAL AND CLASSIFICATION DATA

  • Retrieval galleries contain 749 candidates, including one same-document partner and 748 length-matched distractors, with query–partner Jaccard overlap capped at 0.10.All methods rank the same gallery; SAEs use cosine similarity between codes.Source: §A.5 · p. 16
  • Classification uses eight balanced ArXiv domains, while the OOD set contains papers first submitted in 2023 or later.Each class has 1,024 training, 256 validation, and 256 test documents; the low-label evaluation uses budgets from 1 to 256 documents per class across five seeds.Source: §A.5 · pp. 16–17
  • Joint-Chunk shows smaller effective L0 and a larger empty-code rate across full-code and partner-prefix views.The figure also reports training-batch and fixed-monitor RFVE, inactive-coordinate fraction, and validation empty-code rate.Source: Figure 8 · p. 17

B TRAINING STABILITY AND FEATURE ACTIVITY

The selected checkpoints maintain active sparse codes and continued dictionary use, while the interpretability analyses assess explanation consistency and the prevalence of high-level concepts.

B TRAINING STABILITY AND FEATURE ACTIVITY

  • Zero validation empty-code rate and fewer than 0.025% recently inactive coordinates show that selected checkpoints produce nonempty codes while using nearly the whole dictionary.Recent inactivity is measured over a preceding ten-million-occurrence window once that window is available.Source: B TRAINING STABILITY AND FEATURE… · p. 17

C.1 FOCUSED INTERPRETABILITY

  • Focused interpretability scores whether frozen explanations match active contexts and reject inactive ones, then averages feature scores across 1,000 sampled features.A feature scores zero if either active-context coverage or inactive-context rejection is zero; the dictionary average is reported on a 0–100 scale.Source: §C.1 · p. 18

C.2 HIGH-LEVEL FEATURE CENSUS

  • The high-level census uses a separate ten-context criterion and does not contribute to InterpScore.A blinded judge assesses ten strong contexts from distinct held-out documents; a concept must explain at least six without a sufficient surface-only rule.Source: §C.1 · p. 18, §C.2 · p. 18
  • Mean-Chunk has a 34.8% high-level feature fraction versus BatchTopK’s 17.2%, with all sampled coordinates retained in the denominator.These fractions correspond to 348 and 172 passing coordinates, respectively; Wilson 95% intervals are shown in Figure 3.Source: §C.2 · p. 18

D DICTIONARY UTILIZATION AND FEATURE DYNAMICS

The dictionary audit measures persistence across related passages and how activation mass is distributed among alive features. The displayed domain neighborhoods show Cross-Chunk leading overall purity.

D DICTIONARY UTILIZATION AND FEATURE DYNAMICS

  • Persistence lift compares feature coactivation in adjacent same-document chunks with length-matched chunks shuffled across documents.The audit uses 2,048 validation pairs and 8,192 sampled coordinates, with a minimum support threshold of eight.Source: D DICTIONARY UTILIZATION AND FEATURE… · p. 18
  • Dictionary utilization measures activation-mass distribution as an entropy-equivalent feature count divided by the sampled alive-feature count.Equal activation shares yield 100% utilization; concentration on fewer coordinates lowers the measure even if many coordinates have fired.Source: D DICTIONARY UTILIZATION AND FEATURE… · p. 18
  • Cross-Chunk has the highest overall purity among the displayed methods in 50-dimensional cosine neighborhoods of held-out ArXiv document codes.Figure 9 shows independently fitted t-SNE projections; its purity values are computed in the cached 50-dimensional space.Source: Figure 9 · p. 19
  • Figure 4’s colored traces link five selected features and their frozen explanations to activation patterns as passages change domain.These examples complement aggregate persistence and utilization measurements across the sampled dictionary.Source: D DICTIONARY UTILIZATION AND FEATURE… · p. 19

E SEMANTIC NEIGHBORHOODS ACROSS DOMAINS AND TIME

Cross-Chunk leads overall domain-neighborhood purity and varies in advantage across scientific domains and metrics. The limited-label classification setup tests frozen representations as probe supervision increases.

E SEMANTIC NEIGHBORHOODS ACROSS DOMAINS AND TIME

  • Within-test purity uses non-self ranks 2–11, while out-of-domain purity uses the ten nearest earlier-test codes; Table 2 reports cross-time majority-vote accuracy.These measures distinguish local label agreement from the class assigned by a neighborhood vote.Source: E SEMANTIC NEIGHBORHOODS ACROSS… · p. 19, Table 2 · p. 19
  • 70.6% overall purity for Cross-Chunk compares with 66.7% for BatchTopK, while Quantitative Biology and Quantitative Finance improve by 9.7 and 3.9 percentage points over the strongest other displayed method.Figure 10 reports neighborhood purity; its within-test and out-of-domain query collections each contain 2,048 documents balanced across eight classes.Source: E SEMANTIC NEIGHBORHOODS ACROSS… · p. 19, Figure 10 · p. 20
  • Temporal has the highest silhouette, whereas Cross-Chunk leads neighbor purity and cross-time nearest-neighbor accuracy, with advantages varying by domain.Economics and Mathematics favor token baselines, while Quantitative Biology and Quantitative Finance show larger Cross-Chunk gains.Source: E SEMANTIC NEIGHBORHOODS ACROSS… · p. 20

F CLASSIFICATION WITH LIMITED LABELS

  • Frozen-code classification varies labeled documents per class across seven budgets, using five seeds and identical labeled subsets for all methods.The probe settings, representation, and evaluation split remain fixed as supervision changes.Source: F CLASSIFICATION WITH LIMITED LABELS · p. 20

G JOINT-CHUNK ALPHA SWEEP ACROSS THE EVALUATION SUITE

The tested Joint-Chunk weights preserve stronger semantic feature discovery and downstream transfer than token baselines, while changing the balance among capabilities. The main comparisons use α = 0.25 as a representative setting.

G JOINT-CHUNK ALPHA SWEEP ACROSS THE EVALUATION SUITE

  • The sweep tests α values 0.25, 0.5, 1.0, and 1.5 under the same architecture, token budget, caches, and downstream procedures.The main comparisons use α = 0.25 as a randomly chosen representative setting, and all Joint-Chunk runs use three random seeds.Source: G JOINT-CHUNK ALPHA SWEEP ACROSS THE… · pp. 20–21
  • Across three additional weights, high-level feature fractions range from 29.5% to 34.4%, versus 16.8–17.2% for token methods, and reasoning generalization reaches 59.2–64.0%, versus 46.6% for BatchTopK and 25.9% for Temporal.Document retrieval and classification transfer also improve across all three additional evaluated settings.Source: G JOINT-CHUNK ALPHA SWEEP ACROSS THE… · p. 20
  • Among additional settings, α = 0.5 has the highest high-level fraction and focused interpretability; larger partner weights reduce reconstruction cosine and native reasoning recall.Retrieval and classification vary less across the tested weights.Source: G JOINT-CHUNK ALPHA SWEEP ACROSS THE… · p. 20
  • Across tested weights, the observation unit and prediction target determine complementary strengths within a shared design.The partner weight controls the relative emphasis within that design.Source: G JOINT-CHUNK ALPHA SWEEP ACROSS THE… · p. 21

H LLM JUDGE PROMPTS AND EVALUATION CONTROLS

The appendix specifies frozen-context judging protocols for semantic feature scoring, reasoning detection, and steering. Its reasoning tests separate relations from familiar wording and evaluate feature decisions under native and controlled views.

H LLM JUDGE PROMPTS AND EVALUATION CONTROLS

  • Judging uses fixed-schema prompts and frozen evidence with SAE names hidden; DeepSeek V4 Flash scores at temperature zero.DeepSeek V4 Pro generates reasoning passages and paired controls.Source: H LLM JUDGE PROMPTS AND EVALUATION… · p. 21

H.1 FROZEN EXPLANATIONS AND HIGH-LEVEL SCORING

  • High-level feature scoring freezes a concise explanation before judging whether complete 128-token contexts match it, avoiding trigger-word scoring.The judge sees four active and four inactive contexts; explanations are generated from activating contexts as a general semantic or functional pattern.Source: §H.1 · pp. 21–22
  • The high-level census counts a feature only when at least six of ten contexts support a reusable semantic or functional concept without a sufficient surface-only rule.The fraction is computed across all 1,000 sampled coordinates, including those without explanation support.Source: §H.1 · p. 22

H.2 REASONING RELATIONS, SELECTION, AND CALIBRATION

  • The five reasoning relations are tested with 500 native examples each, split into 400 discovery and 100 held-out examples, using features and cutoffs frozen before evaluation.Cutoffs are calibrated on 208 confirmation-fold Pile documents to an empirical background false-positive rate of at most 5%.Source: §H.2 · p. 22

H.3 SEPARATING REASONING RELATIONS FROM SURFACE CUES

  • Reasoning controls test native recognition, recognition after familiar wording is removed, and rejection when familiar words remain but the relation is absent.Cue-free positives preserve the relation and broad topic; cue-only negatives preserve observed lexical anchors while replacing events to remove the relation.Source: §H.3 · p. 23
  • The 66% versus below-5% backtracking comparison measures recognition under a shared, baseline-anchored lexical stress test.Cue-free positives remove observed anchors while retaining the relation; cue-only negatives retain anchors while removing the relation.Source: §H.3 · pp. 23–24

H.4 MATCHED POOLING AND CUTOFF SENSITIVITY

  • Matched pooling changes BatchTopK backtracking native recall by −31 percentage points and cue-only rejection by +24, while Mean-Chunk and Cross-Chunk retain their decisions.Temporal changes remain within five points; Figure 12 compares paired decisions across 100 held-out documents per relation and view.Source: §H.4 · p. 24, Figure 12 · p. 24

H.5 STEERING GENERATION AND SCORING

  • The steering evaluation adds each feature’s normalized decoder direction at layer 21, then averages feature-wise maximum scores across 1,000 features.It searches a shared strength grid, with an ordered rescue grid for unchanged continuations; the score combines concept direction and coherence preservation.Source: §H.5 · pp. 24–25

I JOINT-CHUNK SAE ALPHA ABLATION

The alpha ablation varies Joint-Chunk’s partner-target weight while holding its training setup fixed, then reports target-specific fidelity and a composite score. The accompanying appendix sections define the fidelity, feature-structure, retrieval, reasoning, and steering measures used in evaluation.

I JOINT-CHUNK SAE ALPHA ABLATION

J.1 FIDELITY AND DIRECTIONAL AGREEMENT

  • FVE measures improvement over a constant predictor, while RFVE divides FVE by a reference FVE on the same examples and target.Self reconstruction uses an identity reference; partner prediction uses a train-fitted dense predictor selected on validation.Source: §J.1 · p. 26

J.2 FEATURE STRUCTURE

  • Feature persistence compares coactivation in adjacent same-document chunks against length-matched pairs shuffled across documents.Positive lift indicates coactivation associated with same-document continuity beyond the shuffled comparison.Source: §J.2 · p. 26

J.3 DOCUMENT RETRIEVAL AND CLASSIFICATION

  • Recall@5 averages whether each query’s correct partner ranks among the top five candidates, with methods using the same partners and galleries.Low-label AUC summarizes probe accuracy across label budgets; OOD accuracy measures future-year domain prediction.Source: §J.3 · pp. 26–27

J.4 REASONING SELECTION, RECALL, AND GENERALIZATION

  • Reasoning evaluation freezes each selected feature and calibrated cutoff across native and control views, measuring native recall and cue-free detection alongside cue-only rejection.The generalization measure combines cue-free detection with rejection of cues without the relation.Source: §J.4 · p. 27

J.5 CAUSAL STEERING AND AGGREGATION

  • The dictionary-level steering score averages one selected maximum score per feature, with ties resolved in favor of the lower intervention strength.The panel contains 1,000 features, and eligible rescue strengths are included.Source: §J.5 · p. 27
Loading 2609.35521v1…