GRAIN: Gradient-Referenced Alignment of Informative Tokens for LLM Supervised Fine-Tuning

Token-level selection has emerged as an effective approach to improving data efficiency in supervised fine-tuning (SFT) of large language models (LLMs). Its benefit hinges on the scoring signal: a useful score should reflect how much a token contributes to the target task, yet remain cheap enough to compute during training. Existing selectors score tokens by loss comparisons against reference or historical model states, which capture learning dynamics but do not explicitly encode downstream task objectives. In contrast, gradient-based influence offers a task-aware alternative, yet incurs substantial computational overhead at the token level. To fill this gap, we introduce Grain, an efficient task-aware token selector that approximates gradient-based influence using last-hidden-state gradients computed directly from the forward pass. Specifically, each token is scored by the gated alignment of this gradient with task directions estimated from a small held-out set, complemented by a lightweight excess-loss term computed from cached base-model losses. Together, these components enable online selection without additional backward passes or a separately trained reference model. Theoretical analysis shows that, under stated conditions, the ungated alignment term approximates first-order token-level influence. Extensive experiments demonstrate the superiority of our proposed method over existing token selectors.