Research map / Large language models / Reasoning and chain of thought

Unlocking Full Efficiency of Token Filtering in Large Language Model Training

ICLR 2026 · Large language models › Reasoning and chain of thought · approach: Efficient LLM serving · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Centrifuge co-designs an attention-backward kernel that filters K/V gradient activations and a graph-rewriting workflow that turns sparse GEMM into dimension-reduced dense GEMM, speeding up token-filtered LLM training.

Problem

Token filtering in LLM training reduces tokens but yields little real speedup because backward sparsity does not propagate past the output layer and existing sparse GEMM kernels only pay off above ~95% sparsity, while token filtering uses 30-50%.

Approach

Filter K/V gradient activations inside a memory-efficient attention backward kernel by splitting computation into a kept-token pass for dK/dV and a full-token pass for dQ, then shrink the sequence dimension across the autograd graph so sparse GEMMs become smaller dense GEMMs.

Key results

• With 50% of tokens filtered, Centrifuge reduces backward time by 40.0%-49.9% and end-to-end time by 17.9%-34.7% across models from 1.1B to 40B.
• Centrifuge improves average math-task accuracy of TinyLlama-1.1B from 17.8 (regular finetuning) to 27.3, with up to 26.6 points gain on MAWPS.
• The strawman softmax/QKV filtering approach fails to converge on OWM training loss, while Centrifuge matches loss-only token filtering.

Surprising findings

• Existing token filtering with sparse GEMM from PyTorch is about 10x slower than dense GEMM at 40% filtering.
• Loss-only token filtering yields only about 1.2% end-to-end speedup even when 40% of tokens are dropped.

Limitations

• Load balancing with MoE remains unresolved and the filtering operator overhead grows with model size because graph updates are layer by layer.
• Evaluation focuses on math reasoning and fine-tuning of pretrained models rather than pretraining at scale.

Experimental setup

Datasets: open-web-math (OWM) synthetic and manually curated math-related tokens

Benchmarks: GSM8K MATH SVAMP ASDiv MAWPS TabMWP MathQA SATMath MMLU

Metrics: accuracy training time per 1M tokens backward time reduction end-to-end time reduction throughput

Baselines: Regular training Token filtering (loss-only backward filtering) Strawman softmax/QKV activation filtering

Builds on: FlashAttention backward token filtering (reference-model loss selection) PyTorch autograd LoRA

Closest papers by meaning

• Beyond Next Token Prediction: Patch-Level Training for Large Language Models (ICLR 2025 · similarity 0.88)
• Better & Faster Large Language Models via Multi-token Prediction (ICML 2024 · similarity 0.88)
• Over-Tokenized Transformer: Vocabulary is Generally Worth Scaling (ICML 2025 · similarity 0.87)
• Token Cleaning: Fine-Grained Data Selection for LLM Supervised Fine-Tuning (ICML 2025 · similarity 0.87)
• Enhancing Large Language Model Performance with Gradient-Based Parameter Selection (AAAI 2025 · similarity 0.87)
• Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum (NeurIPS 2024 · similarity 0.87)