Search papers, labs, and topics across Lattice.
This paper introduces LESSER, a novel method for efficient post-training data selection in large language models by leveraging output-layer gradients instead of full-parameter gradients. This approach significantly reduces the computational cost of data selection, enabling practical application to large candidate pools without sacrificing downstream performance. LESSER achieves substantial FLOP reductions ($9.7\times$ for SFT, $3.0\times$ for RL) while maintaining performance comparable to full-gradient selection.
Output-layer gradients alone are sufficient for effective post-training data selection, slashing feature-extraction costs by up to 10x.
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.
