PaperlayerView exact revision on arXiv

Source-linked AI summary

PULSE: Identifying Demonstration-Utility Features with Sparse Autoencoders

Chenduo Hao, Chuanbao Gao, Pinjun Zeng, Jingze Zhu, Chonghan Liu, Zidong Liu, Xu Yang

arXiv:2609.32469v1cs.AIcs.CLcs.LG

TL;DR

In-context learning depends strongly on demonstration choice, while many selection methods rely on external query–demonstration similarity. PULSE identifies SAE features associated with demonstration utility and uses them for ranking and retrieval; experiments report improvements across classification, generation, and reasoning tasks.

  • Problem

    Demonstration choice strongly affects in-context learning, while common selection methods rely on external query–demonstration similarity.

    Source: §1 · pp. 1–2

  • Method

    PULSE uses paired candidate-set utility and SAE activation differences from a small labeled discovery set to identify features for complete-set ranking and pool-scale retrieval.

    Source: §1 · p. 2

  • Results

    PULSE-Retriever improves demonstration selection across classification, generation, and reasoning tasks, while controlled ranking validates a predictive set-level utility signal.

    Source: §1 · pp. 2–3, §6 · p. 10

  • Takeaways & Limitations

    The results support using explicit model-internal utility-related features alongside external query–demonstration similarity for ICL demonstration selection.

    Source: §6 · p. 10

    PULSE requires reliable SAEs and a small labeled discovery split, and adds offline compute and storage costs beyond lightweight embedding retrievers.

    Source: J Limitations · p. 22

Abstract

from arXiv · show

In-context learning is highly sensitive to demonstration choice, yet most methods select demonstrations using external query-demonstration similarity. Such criteria can miss model-specific signals: Similar demonstrations may activate different internal features and downstream behaviors. We introduce PULSE (Paired Utility Localization over Sparse Encodings), an SAE-based framework for identifying model-internal features associated with demonstration utility and using them for demonstration selection. Using a small labeled discovery set, PULSE samples candidate demonstration sets, measures their zero-shot-relative utility under the target model, and scores SAE features by how their activation differences align with utility differences. The top positive and negative coordinates form a sparse utility-localization vector. We use this vector in two complementary ways: as a signed score for controlled complete-set ranking, and as PULSE-Retriever, which converts its magnitude into a feature-relevance mask for scalable pool-scale retrieval. Across classification, generation, and reasoning benchmarks, PULSE-Retriever improves over the strongest baseline by 2-3 accuracy points, 0.6-0.9 BLEU-4, and 3.2 exact-match points, respectively, while controlled ranking validates the identified features encode a predictive set-level utility signal. Feature inspection and cross-dataset experiments suggest that the identified features capture task-relevant, dataset-conditioned patterns, yet retain utility signals that partially transfer across datasets. Our code is available at https://github.com/aohenuo/PULSE.

1 Introduction

PULSE addresses the gap between external query–demonstration matching and identifying model-internal features associated with utility. It uses those features for demonstration selection, with controlled ranking and retrieval experiments supporting their predictive utility signal.

1 Introduction

  • ICL performance is sensitive to demonstration configuration, making demonstration selection a central problem.For the same query, some demonstrations improve performance while others add little benefit or degrade it.Source: §1 · p. 1
  • Existing methods mainly study query–demonstration matching rather than which model-internal features are associated with useful demonstrations.The paper notes that demonstration utility also depends on which internal features a model activates.Source: §1 · p. 2, Figure 1 · p. 2
  • Dense activations make utility-related features difficult to isolate because utility signals are entangled with other factors.The paper presents sparse autoencoders as a way to identify explicit features associated with useful demonstrations.Source: §1 · p. 2
  • PULSE compares candidate demonstration sets with different downstream utilities and uses SAE activation differences to identify a sparse utility-localization vector.The vector emphasizes features predictive of demonstration utility while suppressing unrelated variation.Source: §1 · p. 2
  • PULSE-Ranking finds that top-ranked demonstration sets outperform lexical, embedding-based, and SAE-space baselines, while PULSE-Retriever improves selection across tasks and settings.The experiments test complete-set ranking and pool-scale retrieval, respectively.Source: §1 · pp. 2–3

2 Related Work

Prior work selects in-context demonstrations through external similarity, learned retrieval, and set-level criteria. PULSE builds on model-aware utility and sparse-autoencoder research to identify internal features associated with useful demonstrations.

In-context learning and demonstration selection

  • ICL demonstrations are selected using lexical or semantic similarity, learned retrievers, and set-level criteria such as diversity and coverage [19] [29] [41] [14] [35] [13] [38] [26] [34] [6].These methods primarily use external query–demonstration or demonstration-set signals.Source: §2 · p. 3
  • Model-dependent demonstration preferences and mechanistic findings motivate identifying utility-related features inside the target model [24] [21].Prior analyses also relate ICL to specific internal components and the balance between learning from demonstrations and retrieving internal knowledge.Source: §2 · p. 3

Sparse autoencoders for interpretable model features

  • Sparse autoencoders decompose dense, entangled activations into sparse feature representations that can be used to study model-specific behavior [8] [9] [27] [15].PULSE uses SAE features to distinguish higher- from lower-utility demonstration sets.Source: §2 · p. 3

3 Method

PULSE estimates a sparse utility-localization vector from paired demonstration sets, validates it by ranking complete sets, and uses it as a feature-relevance mask for scalable retrieval. Its utility estimates use a small labeled discovery set and zero-shot-relative task utility.

3 Method

  • The method represents a k-shot set as support examples paired with task targets, appended before a target-hidden query in the task prompt.SAE activations are mean-pooled over prompt tokens to form the sparse feature vector.Source: §3 · p. 3
  • For classification, PULSE uses a gold-label margin; for generation, it uses length-normalized teacher-forced likelihood, with utility measured relative to zero-shot.Labels are used only in the small discovery set to estimate the utility-localization vector.Source: §3 · p. 4
  • PULSE pairs demonstration sets within each discovery query and scores SAE features by the normalized alignment of activation differences with utility differences.It retains the K+ largest positive and K− most negative feature scores in the sparse utility-localization vector.Source: §3 · pp. 4–5
  • The utility-localization vector is tested by ranking the same pre-sampled complete demonstration sets across methods; higher scores indicate activation of features associated with higher utility.This controlled ranking experiment validates the identified features, rather than serving as the practical pool-scale retriever.Source: §3 · p. 5
  • PULSE-Retriever uses a feature-relevance mask to rank individual support examples, blending masked and standard SAE similarity before greedy context selection.It uses the mask magnitude rather than its sign for pointwise retrieval, because the sign is not calibrated as an individual-candidate preference.Source: §3 · p. 5

4 Experiments

Experiments evaluate PULSE on classification, generation, and reasoning using controlled candidate-set ranking and pool-scale retrieval. Results support predictive utility signals, practical retrieval gains, recognizable feature patterns, and partial cross-dataset transfer.

4.1 Experiment Setup

  • Experiments use six benchmarks across classification, generation, and reasoning; utility features are identified on training data and evaluated on held-out test data.PULSE uses 64 discovery queries and 32 candidate sets per query; hyperparameters are calibrated on a held-out validation split.Source: §4.1 · pp. 5–6

4.2 Main Results

  • Across 32 dataset–shot–backbone cells, PULSE-Ranking is best or tied-best in 30, validating a predictive set-level utility signal.It improves over KNN-SAE by 3.89 points on Gemma2-2B and 3.58 points on Llama3.1-8B.Source: §4.2 · p. 6, Table 1 · p. 6
  • PULSE-Retriever improves average accuracy over KNN-SAE by 2.93 points on Gemma2-2B and 1.99 points on Llama3.1-8B.It achieves the best accuracy in 25 of 32 dataset–shot–backbone cells.Source: §4.2 · p. 6, Table 2 · p. 7
  • +0.93 BLEU-4 on Gemma2-2B and +0.61 BLEU-4 on Llama3.1-8B for CommonGen, and +3.16 exact match on GSM8K over the strongest baseline.PULSE-Retriever achieves the best result in all 12 setting–shot cells.Source: §4.2 · p. 6, Table 3 · p. 7

4.3 Ablation Studies

Do SAE features provide a better internal basis?

  • SAE features outperform PCA by 3.47 points and Raw hidden states by 3.96 points on the 4-shot classification average, and improve CommonGen BLEU-4 from 10.23 and 10.34 to 11.09.The paired identification objective and retrieval pipeline are held unchanged across the SAE, PCA, and Raw comparisons.Source: Do SAE features provide a better… · p. 7, Table 4 · p. 7
  • PULSE-MASK reaches 73.46 average accuracy versus 48.52 for ANTI-RETRIEVAL, while exceeding SPARSE-RANDOM by 1.32 points and BOTTOM-FEATS by 2.15.PULSE-MASK uses the identified feature-vector magnitude without the standard SAE cosine component.Source: Do SAE features provide a better… · p. 7, Table 5 · p. 8

5.1 What Do High-Weight PULSE Features Capture?

  • High-weight features correspond to topic anchors, aspect–opinion bindings, affective cues, and off-aspect patterns; the inspection does not establish that individual features causally mediate ICL performance.The feature inspection and transfer analyses focus on classification under a shared evaluation protocol.Source: §5.1 · p. 8, Do SAE features provide a better… · p. 8

6 Conclusion

PULSE uses SAE features associated with demonstration utility, applying signed utility-localization scores to complete-set ranking and a magnitude mask to scalable retrieval. The retrieval proxy avoids exhaustive set scoring but is not an exact decomposition of complete-context utility.

6 Conclusion

  • PULSE estimates a sparse utility-localization vector from paired differences in candidate-set utility and SAE feature activations.The vector is used as a signed score for complete-set ranking and as a magnitude-based feature-relevance mask for pool-scale retrieval.Source: §6 · p. 9, §B.1 · p. 14

B.1 From Complete-Set Scoring to Per-Example Retrieval

  • The complete-set score tracks activation shifts associated with discovery-time utility, but exhaustive scoring of all k-shot subsets in a large pool is infeasible.The search space grows combinatorially.Source: §B.1 · p. 14

B.2 Why the Retriever Uses a Magnitude Mask

  • PULSE-Retriever uses zero-shot, input-only query and candidate encodings with |wl| as a feature-relevance mask, rather than treating signed weights as per-candidate preferences.This proxy makes candidate encodings cacheable and avoids a full-prompt forward pass for each query-candidate pair, but is not an exact decomposition of complete-context utility.Source: §B.1 · p. 15, §B.2 · p. 15
  • In 4-shot pool-scale retrieval, the magnitude mask performs better on average than the sign-sensitive score.The comparison supports using feature magnitude, rather than the set-level sign, as the per-example relevance signal.Source: §B.2 · p. 16, Table 12 · p. 16

C Greedy Demonstration Selection

  • Greedy context composition selects the highest-scoring shortlist candidate after subtracting an SAE-space redundancy penalty.The penalty strength is controlled by λr, and selected demonstrations are excluded from later choices.Source: C Greedy Demonstration Selection · p. 16

D.1 Visualization Protocol

  • Feature visualizations use max-activating training-pool examples and omit features whose maximum activation is below 0.5.The displayed snippets are illustrative rather than manually selected interpretations.Source: §D.1 · p. 17

D.2 Representative Features Across Datasets

Inspected high-weight PULSE features align with dataset-specific patterns, while feature sign indicates association with discovery-time utility rather than label or causal polarity.

D.2 Representative Features Across Datasets

  • AGNews features often align with topic cues, EMOC features with affective or conversational patterns, and REST14/LAP14 features with aspect–opinion bindings or generic expressions.The feature tables show max-activating snippets, with interpretations describing visible patterns in those snippets.Source: §D.2 · p. 17, Table 14 · p. 18, Table 15 · p. 18
  • A feature’s sign marks its association with paired utility differences, not sentiment polarity, label polarity, or a causal effect.Negative weights associate higher activation with lower discovery-time utility, not necessarily negative task labels.Source: §D.2 · p. 17

E Cross-Dataset Transfer

Cross-dataset experiments find that utility signals partly transfer despite small feature overlap and task-dependent signs. Shared SAE coordinates therefore do not define a single task-independent utility direction.

E.1 Source–Target Transfer

  • Off-diagonal transfer accuracy averages 76.05, 1.29 points below the within-dataset SELF average of 77.34.LOO∪ matches SELF on REST14, slightly exceeds it on LAP14, and is 0.39 points lower on EMOC; AGNews has a larger gap.Source: §E.1 · p. 18

E.2 Feature Overlap and Sign Consistency

  • REST14 and LAP14 share the largest active-feature overlap, 207 features with Jaccard overlap 0.112; all other dataset pairs are below 0.075.The two datasets share aspect-based sentiment structure.Source: §E.2 · p. 19, Table 17 · p. 19
  • Mixed signs occur for 44.1% of features active in two datasets, 60.9% in three, and 75.0% in all four.A shared SAE coordinate can be utility-positive for one task and utility-negative for another, so transfer is partial rather than task-independent.Source: §E.2 · p. 19

F Learned Retriever Baseline and Feedback Budget

PULSE-Retriever outperforms a supervised learned retriever on classification averages across both backbones and more clearly on generation and reasoning. It does so with half the query budget and no gradient-based retriever training.

F Learned Retriever Baseline and Feedback Budget

  • PULSE-Retriever uses 64 discovery queries and closed-form feature scoring without gradient-based retriever training, versus 128 training queries and 2,048 singleton feedback labels for the learned retriever.The authors attribute the comparison to an SAE feature-space utility signal not captured by the standard learned retriever, particularly for generation and reasoning.Source: F Learned Retriever Baseline and… · pp. 19–20
  • PULSE-Retriever exceeds the learned retriever on CommonGen by +1.31 BLEU-4 on Gemma2-2B and +0.71 on Llama3.1-8B, and on GSM8K by +10.13 exact match.On GSM8K, the learned retriever scores 44.53 versus the random baseline’s 47.61.Source: F Learned Retriever Baseline and… · p. 20, Table 19 · p. 20
  • Classification averages favor PULSE-Retriever on both backbones: 77.34 vs. 77.27 on Gemma2-2B and 80.77 vs. 80.22 on Llama3.1-8B.Per-dataset strengths are complementary: the learned retriever leads on AGNews and LAP14, while PULSE-Retriever leads on REST14 and EMOC across both backbones.Source: F Learned Retriever Baseline and… · p. 20, Table 19 · p. 20

G Hyperparameter Sensitivity

Sensitivity analyses show PULSE-Retriever remains effective across feature counts and SAE layers, while discovery-set size and the sparse–dense blend affect average accuracy. The default layer remains close to the best-performing layer.

Number of identified features

  • Average accuracy remains between 76.62 and 77.65 across a wide range of selected feature counts K.The results suggest retrieval does not depend on a finely tuned feature count once enough high-utility features are retained.Source: G Hyperparameter Sensitivity · p. 20, Table 20 · p. 21

Discovery set size

  • Average accuracy rises from 75.53 with 16 discovery queries to 77.34 with 64, then saturates at 128 queries.This indicates a useful sparse utility vector can be identified from a relatively small labeled discovery set.Source: G Hyperparameter Sensitivity · p. 20, Table 21 · p. 21

Sparse–dense blend weight

  • β = 0.3 gives the best average performance; pure PULSE-weighted and pure SAE-cosine retrieval are both weaker than intermediate blends.The blend combines PULSE-masked similarity with standard SAE-space cosine similarity.Source: G Hyperparameter Sensitivity · p. 21

SAE layer choice

  • Across SAE layers, default Layer 12 stays within 1.14 percentage points of the best average accuracy, although the best layer varies by shot setting.The layer sensitivity analysis is on Gemma2-2B; Table 23 summarizes the sweep and Tables 24–27 give per-layer results.Source: G Hyperparameter Sensitivity · p. 21

H Compute Resources

Experiments used frozen language models and pretrained SAEs in local single-GPU jobs. Compute was dominated by discovery and downstream evaluation forward passes, with cached activations reused where possible.

Compute resources

  • Reported runs used one GPU at a time on workstations containing eight RTX 5090 (32GB) and eight RTX 3090 (24GB) GPUs.The language models were frozen, and the SAEs were pretrained rather than trained for this work.Source: H Compute Resources · p. 21
  • For Gemma2-2B classification, building a static SAE cache took about 141–277 seconds per dataset, while discovery took about 448–884 seconds.The discovery phase used 64 discovery queries and 32 candidate discovery sets per query.Source: H Compute Resources · p. 21
  • Approximate Gemma2-2B costs were 0.25–0.6 PFLOPs for one static pool cache and 7–16 PFLOPs for one discovery pass.These are ranges because total compute depends on prompt length, label-word batching, generation length, and cache reuse.Source: H Compute Resources · p. 21

I Seed Stability of PULSE-Retriever

Across three seeds, PULSE-Retriever was more stable on Llama3.1-8B than Gemma2-2B, with Llama3.1-8B per-seed results generally near their means. The authors caution that the seed sweep provides only a coarse variance estimate and does not match baselines across seeds.

Seed stability

  • Standard deviations range from 0.53 to 1.81 across six benchmarks on Llama3.1-8B, versus 0.62 to 2.60 on Gemma2-2B.The authors attribute the difference primarily to Gemma2-2B's smaller capacity, which makes its utility-localization vector more sensitive to discovery-set composition.Source: I Seed Stability of PULSE-Retriever · p. 22
  • On Llama3.1-8B, every benchmark's per-seed values fall within roughly one standard deviation of the mean; GSM8K has std 0.91 over a +3.16-point gain.The authors cite this as support for the robustness of the reasoning-task improvement.Source: I Seed Stability of PULSE-Retriever · p. 22

Limitations

  • Three seeds provide only a coarse variance estimate, and this sweep varies PULSE-Retriever sampling rather than matching baseline protocols across methods.The authors identify matched multi-seed comparisons and paired per-query significance tests as future work.Source: I Seed Stability of PULSE-Retriever · p. 22

J Limitations

PULSE depends on reliable target-backbone SAEs and a labeled discovery split, and adds offline compute and storage beyond lightweight embedding retrieval. PULSE-Retriever is a proxy for set utility, while mean-pooled activations and reasoning-task utility impose additional scope constraints.

Limitations

  • PULSE requires reliable SAEs for the target backbone and a small labeled discovery split, and adds offline compute and storage costs above lightweight embedding retrievers.The costs avoid exhaustive scoring of all k-shot subsets; feature discovery uses full-prompt backbone and SAE passes, and retrieval caches support-pool zero-shot encodings.Source: J Limitations · p. 22
  • PULSE-Retriever is not an exact set-utility decomposition: it turns a vector estimated from set-level activation shifts into a per-demonstration magnitude mask.Mean-pooled SAE activations may also be affected by prompt-length differences; suggested extensions include query-segment pooling, length-matched sets, or length-residualized activation differences.Source: J Limitations · pp. 22–23
  • For GSM8K-style reasoning, discovery utility uses teacher-forced likelihood of supervised solution sequences and assumes access to reasoning traces.Settings with only final answers may require alternative proxies; feature analyses are post-hoc and associational, with causal interventions left for future work.Source: J Limitations · p. 23

Related papers

Loading 2609.32469v1…