Language models show consistent feature use to solve tasks efficiently
How Language Models Organize Task Information
Artificial Intelligence
Summary
Understanding how language models organize information for different tasks can help improve their efficiency. The authors found that many language models use a similar small number of important features to make predictions, and that instruction tuning can reduce this number even further. They also showed that models trained to learn new tasks need fewer examples when the features they rely on are sparse and well-organized. Their results help distinguish between how easy a task is to read from a model and how many features it actually depends on.
What this means in practice
- •For machine learning engineers: Optimize language model fine-tuning by focusing on fewer key features to reduce training data needs and speed up adaptation to new tasks.
- •For ai deployment teams: Diagnose and predict performance changes after pruning or continual learning by monitoring feature readout sparsity in language models.
Authors
Kaiyang Wan, Yuxia Wang
Abstract
Dense linear probes assess task readability from language model (LM) representations, but do not directly reveal how many features suffice for prediction. We measure readout sparsity with \(k^*\), the smallest feature count found by a fixed search protocol that recovers 95% of dense-probe AUROC above chance. We evaluate 19 LMs from five labs on 14 core tasks. Ten base models give the nine reliably readable tasks similar feature counts and nearly the same order. This agreement extends to the selected features themselves, which on three tasks show cross-model response correspondence beyond class-mean differences and a label-correlation-matched null. Instruction-tuned deployment selectively reduces \(k^*\), with both weights and input format contributing. In controlled synthetic experiments across 20 models, successfully learned targets admit single-feature readouts, whereas readable controls require tens of features. Reference-model \(k^*\) is associated with label requirements across eight students, following an empirical scaling of \(N^*\propto(k^*)^{0.65}\). After continual learning, saved heads can fail while linear readability remains high, and recovery requires fewer labels than initial task learning under the tested procedures. Pruning reduces linear readability, with wider pre-pruning readouts generally associated with larger losses. These findings distinguish task readability from readout sparsity and establish their complementary value for predicting learning requirements and diagnosing LM changes.