PULSE finds model features that improve demonstration selection in AI tasks
PULSE: Identifying Demonstration-Utility Features with Sparse Autoencoders
Artificial IntelligenceComputation and LanguageMachine Learning
Summary
Choosing the best examples to teach AI models how to perform tasks is tricky because similar examples might affect the model differently. The authors introduce PULSE, a technique that looks inside the AI model to find specific features related to how useful example sets are. By using a small set of known examples, PULSE identifies features that predict which examples will help the model the most, improving performance in tasks like classification, text generation, and reasoning.
What this means in practice
- •For machine learning engineers: Improve AI model performance by selecting more effective teaching examples based on internal model features identified by PULSE.
- •For natural language processing developers: Enhance text generation, classification, and reasoning tasks with better example selection informed by model-internal utility features.
Authors
Chenduo Hao, Chuanbao Gao, Pinjun Zeng, Jingze Zhu, Chonghan Liu, Zidong Liu, Xu Yang
Abstract
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.