LESSER: Post-Training Data Selection with Output-Layer Gradients
LESSER cuts post-training data-selection cost by using output-layer gradients from forward passes, reducing FLOPs 9.7x for SFT while matching full-gradient performance.
LESSER is a drop-in wrapper for gradient-based post-training data selection that ranks samples using output-layer gradients computed via cheaper forward passes instead of full backward passes. It reduces feature-extraction FLOP cost by 9.7x for SFT and 3.0x for RL benchmarks while tracking full-gradient downstream performance. The authors find that even when output-layer and full gradients rank individual samples differently, the selected batches remain gradient-aligned.
- Output-layer gradients from forward passes replace expensive full backward passes
- 9.7x FLOP reduction on SFT, 3.0x on RL selection benchmarks
- Downstream performance matches full-gradient selection methods
- Works as drop-in wrapper for existing selection methods
Full article144 words · extracted from arxiv.org · click to collapse
The choice of post-training data for large language models substantially affects downstream performance. Gradient-based data selection is a popular approach that ranks training data by how well their gradients align with those of a small validation set. However, ranking with full-parameter gradients requires an expensive backward pass on every sample, making computation intractable for large candidate pools. This raises a natural question: can we approximate full-gradient features at a fraction of the cost? Conveniently, we find that output-layer gradients suffice for effective data selection, yet require only the cheaper forward pass. We implement this as LESSER, a drop-in wrapper for selection methods that reduces the feature-extraction FLOP cost by $9.7\times$ for SFT and $3.0\times$ for RL benchmarks, while tracking full-gradient performance on downstream tasks. Empirically, we find that even when output-layer and full gradients rank individual samples differently, they select batches with aligned gradients.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2610.03702