TRIM: Token-wise Attention-Derived Saliency for Data-Efficient Instruction Tuning

TL;DR

TRIM introduces a token-centric, forward-only data-selection method for instruction tuning that leverages attention-derived token fingerprints from a small target set to identify high-signal examples. It builds fingerprints by aggregating token saliency from multi-layer attention and matches candidate samples to these fingerprints via cosine similarity, yielding a robust, one-pass scoring process. Empirically, TRIM achieves up to 9% improvement with only 5% coresets, surpasses some full-data fine-tuning results, and transfers effectively across models, while reducing length bias and enabling task-adaptive data selection. This approach offers a scalable, efficient alternative for constructing high-quality instruction-tuning datasets in large-scale settings, reducing computational cost without sacrificing performance. The work demonstrates substantial practical impact for data selection in LLM fine-tuning, particularly in resource-constrained or rapidly evolving scenarios where data quality matters more than sheer quantity.

Abstract

Instruction tuning is essential for aligning large language models (LLMs) to downstream tasks and commonly relies on large, diverse corpora. However, small, high-quality subsets, known as coresets, can deliver comparable or superior results, though curating them remains challenging. Existing methods often rely on coarse, sample-level signals like gradients, an approach that is computationally expensive and overlooks fine-grained features. To address this, we introduce TRIM (Token Relevance via Interpretable Multi-layer Attention), a forward-only, token-centric framework. Instead of using gradients, TRIM operates by matching underlying representational patterns identified via attention-based "fingerprints" from a handful of target samples. Such an approach makes TRIM highly efficient and uniquely sensitive to the structural features that define a task. Coresets selected by our method consistently outperform state-of-the-art baselines by up to 9% on downstream tasks and even surpass the performance of full-data fine-tuning in some settings. By avoiding expensive backward passes, TRIM achieves this at a fraction of the computational cost. These findings establish TRIM as a scalable and efficient alternative for building high-quality instruction-tuning datasets.

Figures (5)

• Figure 1: Overview ofTRIM. (a) From multi-layer attention, we derive an aggregated token saliency signal that combines row (allocation sharpness) and column (received attention) signals. (b) Using a small target set, we compute one fingerprint per token class by saliency-weighted averaging of last-layer hidden states, capturing task-defining patterns (\ref{['sec:method_fingerprinting']}). (c) For each candidate sample, token states are matched to their class fingerprints (cosine similarity), and token scores are pooled into a single relevance score for ranking (\ref{['sec:method_scoring']}); top-ranked samples form the instruction-tuning coreset.
• Figure 2: TRIM vs. top baselines on Llama-3.2-1B across coreset budgets for CommonsenseQA (left) and SocialIQA (right). TRIM maintains a consistent advantage and, on SocialIQA, exceeds the full-data baseline with coresets as small as 1%.
• Figure 3: Length distribution of selected coresets for GSM8K. The histogram shows the percentage of selected samples by length bucket. Sample-level methods (LESS, TAGCOS) skew short, whereas TRIM selects a broader distribution with a higher mean length.
• Figure 4: Distribution of selected training examples across data subsets for GSM8K and CommonsenseQA.
• Figure 5: Distribution of selected examples across training subsets for different methods and target tasks. Each bar shows the percentage breakdown of data selected from CoT, Dolly, Flan_v2, and Oasst1 datasets.