Unlocking Full Efficiency of Token Filtering in Large Language Model Training
Authors: Di Chai
LI Pengbo
Feiyuan Zhang
Yilun Jin
Han Tian
Kaiqiang Xu
Binhang Yuan
Dian Shen
Junxue Zhang
Kai Chen
Venue: ICLR 2026 Poster (2026)
Keywords: Efficient LLM Training; Token Filtering;
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: This paper presents Centrifuge, a system that leverages algorithm and system co-design to unleash the full efficiency of token filtering in LLM training.

Abstract: Token filtering has been proposed to enhance the utility of large language models (LLMs) by eliminating inconsequential tokens during training. While using fewer tokens is expected to reduce computational workloads, existing methods have not yet achieved a real-world efficiency boost. This is primarily due to two factors: (1) existing work has inadequate sparsity for speedup, and (2) token filtering operates within a sparsity range that is non-standard in existing machine learning (ML) libraries and thus cannot be efficiently supported. This paper presents Centrifuge, a system that leverages algorithm and system co-design to unleash the full efficiency of token filtering in LLM training. At the algorithm level, Centrifuge filters activations of inconsequential tokens in the attention backward kernel to amplify the sparsity in backward computation. At the system level, Centrifuge proposes an automatic workflow that transforms sparse GEMM into dimension-reduced dense GEMM for optimized efficiency using standard ML libraries. Evaluations on models with various scales—from 1.1B to 40B—demonstrate that Centrifuge reduces backpropagation time by up to 49.9\% and end-to-end training time by up to 34.7\% when filtering 50\% of tokens. Utility assessments indicate that Centrifuge preserves the utility benefits of token filtering and significantly enhances model performance by up to 26.6\% compared to standard training. Centrifuge is designed for seamless integration into existing LLM training frameworks, enabling systems already utilizing token filtering to accelerate training with just one line of code.

Published as a conference paper at ICLR 2026
UNLOCKING FULL EFFICIENCY OF TOKEN FILTERING IN LARGE LANGUAGE MODEL TRAINING
Di Chai', Pengbo Li,, Feiyuan Zhang", Yilun Jin Han Tian', Kaiqiang Xu², Binhang Yuan Dian Shen', Junxue Zhang Kai Chen
1 Shhnghha University of Finance and Economics
Hong Kong University of Science and Technology 3 University of Science and Technology of China
SSoutheast University
ABSTRACT
Token filtering has been proposed to enhance the utility of large language mod- els (LLMs) by eliminating inconsequential tokens during training. While using fewer tokens is expected to reduce computational workloads, existing methods have not yet achieved a real-world efficiency boost. This is primarily due to two factors: (I) existing work has inadequate sparsity for speedup, and (2) to- ken filtering operates within a sparsity range that is non-standard in existing ma- chine learning (MI.) libraries and thus cannot be efficiently supported. This paper presents CENTRIFUGE[] system that leverages algorithm and system co-design to unleash the full effiiency of token filtering in LLM training. At the algo- rithm level, CENTRIFUGE filters activations of inconsequential tokens in the at- tention backward kernel to amplify the sparsity in backward computation. At the system level, CENTRIFUGE proposes automatic workflow that transforms sparse GEMM into dimension- ing standard MI. libraries. Evaluations models with various scales-from 1.1B to demonstratee that CENTRIFUGE reduces backpropagation time by up to 49.9% and end-to-end training time by up to 34.7% when filtering 50% of to- kens. Utility assessments indicate that CENTRIFUGE preserves the utility benefits of token filtering and significantly enhances model performance by up 26.6% compared to standard training. CENTRIFUGE is designed for seamless integration into existing LLM training frameworks, enabling systems already utilizing token filtering to accelerate training with just one line of code.
dense GEMM for optimized efficiency us-
INTRODUCTION
Training high-quality large language models (LLMs) is notably resource-intensive, requiring sub- stantial investments in both data and computational power. For example, the training process of Llama3-70B occupied approximately 7 million GPÜ hours over 15 trillion high-quality tokens (Dubey et al. 2024a). Token filtering represents an emerging paradigm aimed enhancing the cost- efficiency of LLM training by systematically discarding less significant tokens early in the training process This methodology enables the model to concentrate on the most pertinent tokens, resulting in to 30% absolute improvement in model utility across various tasks (Lin et al. 2024).
While the effectiveness of token filtering in enhancing model utility is well recognized, its potential to improve computational efficiency in training remains largely unexplored. In principle, by signifi- cantly reducing the number of tokens processed in the computational pipeline, token filtering should decrease computational demands and expedite training. However, our analysis reveals that combin- ing token filtering with existing LLM training systems can only introduce a marginal 1.2% speedup in end-to-end training time, even when 40% of the tokens are eliminated ($3.I). This limited en- hancement training efficiency constrains the broader advantages of token filtering for large-scale
Junxue Zhang and Kai Chen are the corresponding authors.
A centrifuge is a modern system used in chemistry lahoratories to efficiently filter different elements.
This paper primarily focuses backward filtering, it demonstrates superior performance in enhancing the capabilities of T.I.M. Ms.. Further details can be found in 2.1

---

Published as a conference paper at ICLR 2026
LLM training. Therefore, we pose the question: Can we fully unlock the efficiency of token filtering while simultaneously achieving better utility than conventional training?
To answer this question, first investigate the key limitations of the existing token filtering sys- tem: (1) Inadequate sparsity and the naive approach of amplifying sparsity cannot work. Existing approaches (Lin et al. 2024) filter tokens the output layer during loss calculation, which does not create true sparsity. The gradients of these dropped tokens are still computed in the attention backward kernel and propagated to the front layers, leading dense matrix operations and negating potential efficiency gains. A strawman approach of filtering softmax activations does retain spar- sity but cannot work with mainstream memory- -eforccienfin attention implementations since softmax outputs are not explicitly stored, and naively filtering activations used for recomputing softmax un- intentionally causes interference and harms the gradients. (2) Non-standard sparsity range. Token filtering potentially brings 30~50% sparsity (Lin et al. 2024), while current sparse matrix multipli- cation (GEMM) implementations in machine learning (ML) focus on high sparsity range and require >95% spansity ($3.2 to be effective. Thus, token filtering cannot be supported by existing sparse computations implementations because of mismatching sparsity range. Using these implementations in token filtering actually slows down training efficiency rather than accelerating it ($3.2).
To fully unlock the training efficiency of token filtering and simultaneously achieving better util- ity than conventional training, propose CENTRIFUGE. At its core, CENTRIFUGE integrates an algorithm and system co-design:
1. At the algorithm level, CENTRIFUGE carefully analyzes the backward computation and proposes further filtering activations of inconsequential tokens in the attention backward kemeel to amplify the sparsity. Our solution designed be compatible with mainstream memory-efficiet at- tention implementations (e.g., FlashAttention (Dao et al. 2022) by separately processing each gradient output to avoid interference. CENTRIFUGE sustains the utility advancements of existing token filtering methods (Lin et al. 2024) and unlocks the potential for efficiency improvement.
2. At the system level, CENTRIFUGE leverages the characteristics of token sntering-specifein the sparsity of matrices in either columns or rows transform sparse GEMM into dimension- reduced dense GEMM, maximizing efficiency on existing machine learning libraries. However, PyTorch's 's dynamic graph nature complicates global updates to dimensions and variables, as graph variability and node differences prevent static rules for correctness. To overcome this, we design automatic workflow leveraging the runtime stability (i.e., the graph remains stable during the training) to dynamically identify and update the necessary dimensions and variables before backpropagation.
We implement CENTRIFUGE to be easily integrated into existing training pipelines with minimal code changes. Systems already using backward token filtering only need to add one line of code to achieve efficiency improvement. To better demonstrate the versatility of CENTRIFUGE and its seam- less integration with existing training systems, we have broadly tested CANTRIFUGE training scenarios, such as distributed training using tensor parallel (TP) and parameter-efftcieen fine- tuning using low-rank adapters (LoRA), and CENTRIFUGE has shown consistent efficiency gain.
various
We comprehensively evaluate CENTRIFUGE in terms of both utility and efficiency using four com- pact yet powerful models: TinyLlama- 1. (Zhang et al. 2024), Qwen2.5-1.5B (Yang et al., 2024), Llama3.2- 3.3BBB and Llama3. 1-13..1888 1-8B (Dubey et al. 2024b). The evaluation results indicate that CEN- TRIFUGE fully preserves the utility benefits of token filtering while increasing model performance by up to 26.6% over standard training methods on nine tasks. With 50% of tokens filtered, CEN- TRIFUGE reduces backward and end-to-end training cost by up to 49.0% and 31.7%, respectively. Notably, CENTRIFUGE achieves greater efficiency gains in scenarios with high computational de- mands (e.g., long-context training) and intensive communication (e.g., tensor parallel operations), demonstrating its high practicality in real-world LLM training.
PRELIMINARY AND RELATED WORK
2.1 PRELIMINARY OF TOKEN FILTERING IN LLM TRAINING
Token filtering is recently proposed technology that has been well recognized by the AI commu- nity. The core idea is to identify and filter out tokens that are either noisy or unlikely to contribute meaningfully to the LLM training process, which implicitly improves the quality of training data
2

---

Published as a conference paper at ICLR 2026
XN
Sequence
Buippoque
Selected states
Antention
WM Peol
Loss (Filtered)
EEN
Dropped states
Label
(a) Filtering training states during forward.
Filtering training loss during backward.
Figure 1: An overview of existing token filter studies.
to benefit the model utility. Moreover, by reducing the total number of tokens to be trained, token filtering also brings the opportunity for efficiency improvement. Existing token filtering work can be categorized into two types: forward token filtering and backward token filtering. As illustrated in Figure forward token filtering techniques remove training tokens during the forward pass, whereas backward token filtering methods eliminate tokens exclusively during the backward pass.
Forward token filtering methods have been extensively studied in previous works (Hou al. 2022
Zhong et al.] 2023 Yao et al. 2022 Ataiifard et al. 2024). However, they typically underperform compared to backward filtering methods due to semantic losses (Zhong et al., 2023 Yao al. [022 Lin et al. 2024). As shown in Figure 1 forward token filtering methods filter tokens at each layer of the forward computation, such that each layer of the model only processes partial context. However,
this approach has been shown to cause semantic loss and potentially harm model utility (Zhong et al.)
2023
Yao et al. 2022). Evaluations in existing forward filtering studies (Hou al.) [0022 Zhong 2023 Yao et al. 2022 Ataicfard et al. 2024) report only similar or lower model utilities and fail to achieve the improvements in utility seen with backward filtering methods (Lin al., 2024).
et al.
In this paper, we focus on backward token filtering due to its impressive advantages in improving model utility. As illustrated Figure [ the backward filtering method maintains standard forward computation while performing selective token training in the output layer. Existing studies leverage a reference model assess the importance of each token. Mathematically, backward token filtering can be formulated as follows (Lin et al. 2024);
Σ Ia% (x:) log (o(xx(Xxx::)
fitter
N x 6%
1, i E top 1% of (Co(x;) Cres(xi))
Tx%(x,)
(2)
10, otherwise
where La is the loss of the target model, Lrey is the loss of the reference model, and Litter is the actual loss to train the target model while keeping 4% of tokens.
2.2 RELATED WORK
Data selection. Data selection is a pre- pre-processing technique (iee, applied before training) aimed
improving data quality. Typically, involves selecting diverse and high-quality training sam-
ples (Bai et al. 2024: Wettig et al. [2024] Xie et al. [2023] Fan et al. 2024] Ye et al. 2025 Liu et al.)
2025 Wenzek et al. 2020) [hrush al. 2025). Data selection occurs the sample level prior to
training which can introduce biases that can negatively impact model utility (Lin et al. 2024). In
contrast, CENTRIFUGE functions as fine- -reeggraiedd data selection method the token level and, more
importantly, selects tokens in a model-adaptive manner by evaluating whether the tokens contributes significantly improving the model's performance.
Parameter-efficieent training. Apart from efficient training methods that focus on data sparsity, another branch of research is parameter-efficient training (Ding al., 2023). which emphasizes spar- sity in model parameters. Typical techniques include low-rank adapters (Hu et al., 2022 Dettmers) jet al. 2023 Karimi Mahabadi et al.) 2021 Loeschcke et al. 2024 Valipour et al. 2023), prefix tuning (Li & Liang 2021), etc. CENTRIFUGE and parameter-efficient training differ in their focus on improving training efficiency at the data and parameter levels, respectively. In particular, CEN- TRIFUGE can work with parameter-efficient training methods to further enhance efficiency. Our experiments demonstrate that CENTRIFUGE is highly generalizable and seamlessly integrates into LoRA training, accelerating backward computation by up to 43. 5% when filtering 50% tokens.
3

---

Published as a conference paper at ICLR 2026
SpArGGEMM (batched)
GEMM (batche])
Filtered tokens have zero gradicnts
10
SpicseCEMM (batch
SpereGEMM (Banch-2) --- GEMM (barch
(batch
Fltering
95%
Thme (s)
18'
Gy = Softmar" x
Gy (sparse gradients from output layer)
10'
Gradients passed freat layers are not sparse.
16"
Token Filtering
(a) Leaving the activation (e.g., softmar) of filtered to- kens unchanged makes the gradients of attention block sity >95% and cannot improve efficiency of token being dense matrices, which leads vanishing gradients. filtering which typically drops 30%~50% tokens.
(b) PyTorch sparse GEMM is effective when spar-
Figure 2: Existing token filtering methods exhibit inadequate sparsity, and even when sufficient sparsity is present, sparse GEMM implementations are unable efficiently support token filtering.
3 OBSERVATIONS ABOUT TOKEN FILTERING
We present two key observations that inspire the design of CENTRIFUGE: (1) deeper understanding of why existing token filtering method has inadequate sparsity and the impracticality of intuitive strawman approach of amplifying sparsity; (2) profiling results showing the impact of sparsity mis- matching between token filtering and sparse computations in ML libraries.
3.1 INADEQUATE SPARSITY OF EXISTING TOKEN FILTER
Inadequate sparsity after token filtering. Current methods filter the loss of unimportant tokens at the output layer, resulting in sparse gradients. However, they leave all dense activations unchanged. Consequently, after being multiplied by these dense activations, the gradients are no longer sparse once they pass through the first attention block. Therefore, existing backward filtering methods (Lin jet al., 2024) exhibit inadequate sparsity, even after filtering the loss at the output layer. Figure ((a) illustrates the process of computing gradients for V (i.e., Gv) using sparse gradients while main- taining unchanged activations (i.e., activations of all tokens are retained). After filtering the tokens based loss, the gradients of the corresponding tokens become zero, as depicted in Figure 2(a) However, because the activations of the filtered tokens remain unchanged, the gradients of V are no longer sparse. Consequently, the backward computation following the first attention block lacks sparsity, limiting efficiency improvements solely within the output layer.
Following the setting in existing work (Lin al. 2024), we can estimate the upper bound of ef- ficiency improvement with existing token filtering schemes. Taking TinyL a model with 22 layers and I.IB parameters, as an example. Filtering 40% tokens will only linearly improve the efficiency on backward propagation of the last layer, while no front layers can be improved. Thus, the backward efficiency can only be improved by 1.8%. Given that backpropagation consumes 66% of the whole training (Narayanan et 2021), the end-to-end efficiency improvement is only 1.2%.
Strawman approach of creating sparsity cannot work. A simplistic strawman approach to amplify sparsity is that we further filter the activations accordingly. Specifically, the activation of softmax determines the sparsity of V's gradient and implicitly impacts the gradients of Q and K through G (i.e., GA is also computed based on softmax). We can filter the activation of softmax, i.e,, set the data corresponding to the filtered tokens to zero, to retain the sparsity.
However, this strawman approach cannot work in existing LLM training system that typically uses memory-efficiency attention implementations (e.g., FlashAttention (Dao et al. 2022) ). Because: (1) memorr--effcieen attention does not explicitly store the softmax output, making impossible to directly filter the softmax activations; (2) if we alternatively filter the Q, K, V activations, which are used for recomputing the softmax and its gradients in backward of memory-efficient attention, an interference will occur between the kernel outputs, since not all gradient outputs impact sparsity and may not require activation filtering. Specifically, the gradients aQ will be unintentionally harmed (...). Based on our experiments, naively using the strawman approach with memory-efficient attention significantly harms model training, causing non-converging training loss (5.2).
Observation 1. Existing token filtering methods exhibit inadequate sparsity because the sparsity fails to propagate from the output layer to the front layers. We require algorithm that ampli-
4

---

Published as a conference paper at ICLR 2026
fies sparsity and, importantly, is compatible with mainstream memory-efficient attention kernels to achieve true efficiency gains.
3.2 MISMATCHING SPARSITY RANGE HINDERS EFFICIENCY GAIN
Inefficient sparse GEMM. Existing sparse GEMM implementations are not well-suited for the token filtering training. Although sparse GEMM is a hot research topic and PyTorch has provided sparse tensor implementation (i.e., torch. sparse). the efficiency of existing sparse GEMM is only improved when the data has very high sparsity (e.g., 95%). To demonstrate the problem, we perform experiments on our testbed (details in $5.1 We compare the efficiency of sparse GEMM in Py TyTochh and regular GEMM in the scenario of token filtering, i.e., the matrix is sparse by row or columns. Figure[2(]
the comparison results under different ratios of token filtering and batch sizes. The sparse GEMM is more efficient only when over 95% of all tokens are filtered, which is unrealistic for token filtering. Under the typical filtering rate of 40%, sparse GEMM even 10x slower than regular GEMM.
Observation The valid sparsity range of existing sparse GEMM in ML libraries exceeds 95%, which is misaligned with the 30%~550% sparsity in token filtering. Naively using existing sparse GEMM implementations in token filtering reduces efficiency rather than accelerating it.
CENTRIFUGE
To solve this problem, we propose CANTRIFUGE, a system leveraging algorithm and system co- design to unleash the full efficiency of token filtering. CENTRIFUGE features two main design points: 1) At the algorithm level, CENTRIFUGE filters the activations in the memory-efficiett atten- tion backward kernel to amplify sparsity ($4.1); 2) At the system level, CENTRIFUGE transforms the dense GEMM through automatically updating the backward computation graph to achieve maximum efficiency using existing ML library (422
sparse GEMM to dimension-reducc
4.1 AMPLIFYING SPARSITY IN MAINSTREAM ATTENTION IMPLEMENTATION
To properly amplify the sparsity, we need to comprehensively analyze the impact of gradients back- propagation on the desired sparsity. Based on the architecture of decodcr-only model, we can cate- gorize the backward process into inter-token and intra-toker computations. The inter-token compu- tations happen in self attention block when computing the so ttmaa ttaaz(QKT) QKT)V, which incorporates computation between tokens. The other computations, including the feed-forward network (FFN), projection of Q, K, V, and merge of multi-head attention, are all intra-token computations, which only change the hidden dimensions. In token filtering, desired sparsity exists in sequence di- mension and only influenced by inter-token computations (i.e., the attention block). Thus we only need process the attention backward kernel to amplify the sparsity.
Based on observation (3..I), it is essential to address the issue of inadequate sparsity while en- suring compatibility with mainstream memory-efficient attention implementations to deliver true efficiency improvements. To accomplish this, we analyze the characteristics of the three outputs of memory-efficient attention backward kernel when using token filtering: aQ (query's gradient), ŌK (key's gradient), and @V (value's gradient).
eQ have zero values at positions of filtered tokens long as the input gradients of the attention backward kernel is sparse. Specifically,, eQ have the same row sparsity with the gradients of Altn based on the attention forward computation Alln = Ioo QKK COTQQK//VV))) /vdjV. Thus, aQ is not factor of causing inadequate sparsity in token filtering.
ÖK and @V have non-zero values at positions of filtered tokens since they have the same row sparsity with Q which is a dense matrix, instead of depending on the input gradients of the attention backward kernel. Thus, we need to filter 0K and 0V accordingly to retain sparsity.
The above findings indicate that only 0K and @V impact sparsity and require filtering to amplify the efficiency. In contrast, aQ remains to be zero for filtered tokens and does not affect sparsity, thus no filtering is needed. Mathematically. we find a conflict: calculating aQ requires the full K, V acti- vations, whereas calculating ØK, ØV requires only part of the K, V activations (i.e., the remaining part after filtering). Strawman approach fails because they cannot satisfy these requirements simul- taneously. Specifically, when combining the strawman approach with memory-efficient attention,
5

---

Published as a conference paper at ICLR 2026
the strawman method needs to filter the activations of K, V in advance to obtain efficiency gain. Unfortunately, this confficts with calculating aQ, which requires the full K, activations. Thus, there is an interference between processing öK, 0V and aQ.
We propose a novel attention backward kemnel that separately processes outputs depending on whether using filtered non-filtered activations to prevent interference, and the detailed compu- tations are demonstrated in Equation (3) and Equation (4) We can divide the computation into two parts: (1) computing 0K and @v using activations of non-filtered tokens (Equation (3; (2)
computing eQ using activations from all tokens (i.e., including the filtered tokens) (Equation (4)
Essentially, the new kernel reorganizes the FlashAttention backward workflow efficiently com- pute 0K, @V while preserving data integrity for aQ.
$ = QRT P=crp($ -LSE) op aÔVT as = Po (OP -D) " =oss
(3)
ØK = ass' Q @V prag
$ QKT p= erp($r-ppSp) OP aóèt
0$ = P -8P - D)
(4)
8Q - 0SK + OSK
Given an activation A E Rxxssh the A Rx xh and À [2xx represent the activations of remaining and filtered tokens, where b is batch size, 51 + 82 is sequence length, and / is the hidden dimensions. We use the same notations with FlashAttention backward algorithm (Dao et al, 2022), where D is the sum of 8o, LSE is the logsumexp. The kemel output aQ, öK av have the same shape of R xh where 81 is the number of remaining tokens. The proposed new attention backward computation is fully compatible with existing memory-efficiet attentions, amplifies the sparsity in attention backward kernel, and thereby enables the sparsity propagate from the output layer all front layers.
4.2 DIMENSION CHIENSIONRREDUEED DENSE GEMM
While the sparsity is properly amplified from algorithm level, another problem at system level is that existing sparse GEMM implementations cannot effectively support token filtering. Our experiments show that existing sparse GEMM is effective only when the data is highly sparse (e.g., filtering >95% tokens) ($3.2.. However, the typical token filtering ratio is 30% 50% under which existing sparse GEMM implementations even have worse efficiency than dense GEMM.
To address this problem, we propose transforming sparse GEMM into dimension-reducee dense GEMM by leveraging the characteristics of sparsity in token filtering scenarios. We first carefully analyze the characteristics of sparse computation in the backward process. Figure 3 shows the generalized backward process in the computation graph of existing machine learning libraries (e.g., PyTorch), where G E axahi is the gradients matrix, b is the batch size, s is the sequence length, h1 is the hidden size, W Rb1xh, xha is the parameter matrix, and X haxaaa is the input matrix.
The gradients of the filtered tokens are set to zero (i.e., G is row-wise sparse).
Model Weight
Next Node
Gw X (dense)
Gx G-W
or Gx G g (sparse)
Saved Variables (X. W.g)
Backward Node
G (sparse)
The generalized backward computation in existing MI libraries (e.g., PyTorch). Each node in the computation graph necessarily has the input gradients and output for the next nodes. The model weights and saved variables could be optional and vary in different nodes.
Figure 3:
We can categorize the sparse GEMM into two types: 1) Gradients for next nodes: G Gaparse W, G Gaparse C g, which are sparse and passed to the next nodes. g scaler vector that performs element-wise computation the gradients; 2) And the gradients of model parameters: G.
Gapaee" X,

---

Published as a conference paper at ICLR 2026
which will be dense and updated to the model weights. Based these analyses, we have the following findings:
The gradients passed to the next nodes inherit the sparsity of the input gradients and the sparsity follows the same pattern the input gradients (i..,, the gradients of filtered tokens are zeros), Thus if we shrink the sequence dimension at the initial gradients (i.e., removing the zeros), all the afterward gradients will be automatically reduced.
The sequence dimension vanishes in the gradients of model parameters, which reasonable since the parameters are independent of the sequence length, and we can directly shrink the sequence dimension to reduce the computational cost.
Instead of directly optimizing the sparse matrix computations, we leverage the above observations and propose to globally reduce the sequence dimension of the gradients and the saved variables in the backward computation graph to accelerate the computation. By shrinking the sequence dimen- sion, we can transform the sparse GEMM to dimension-reduced dense computations with optimized performance. Compared with directly optimizing the sparse GEMM, transforming dense GEMM is more effective since the dense GEMM has been well optimized in existing ML libraries.
However, implementing the dimension-reducce GEMM is non-trivial due to the dynamics of the computation graph's dynamic structure and node usage, which vary across implementations and inputs, making static updating rules impractical. To address this issue, propose an automatic workflow to amend the computational graph. The key insight is that, even though the graph is highly dynamic, it still can be deterministic when using the same implementation and inputs. Partic- ularly, the model implementation and inputs remain the same during the whole training (i.e., runtime stability). Thus, we can mimic the input and traverse the graph using the same implementation to dynamically determine the node information. The automatic workflow contains two steps: 1) gener- ating the skeleton code for processing each type of nodes and their attributes; 2) leveraging special markers (e.g., prime numbers) to identify the target updating dimensions and dynamically generating detailed node processing rules. Due to the space limitation, we put the detailed design of automatic graph updating workflow in appendix B
4.3 OVERALL SYSTEM IMPLEMENTATION
CENTRIFUGE improves the efficiency of token filtering through an algorithm and system co-design: new attention backward kernel amplify the sparsity from algorithm level, and transforming sparse GEMM to dense GEMM to optimize efficiency from system level. We implement the new attention kernel through leveraging FlashAttention on filtered Q K, V activations and further computing the remaining aQ using another cuDNN memory-efficient attention kernel that sup- ports attention bias. For non-attention kernels, we directly remove the activations based on the system design. Eventually, for systems that already utilize token filtering, only one line of code centr i fuge. ops backkard..iitee (loss, filter .maaer) is needed to get the full ef- ficiency of token filtering. Due to space limitation, put the implementation details in appendix C
EVALUATION
In this section, comprehensively evaluate CENTRIFUGE in terms of both utility and efficiency. The results demonstrate that CENTRIFUGE fully preserves the utility benefits token filtering and can improve model performance by up to 26.6% compared to regular training. When filtering 50% of tokens, CENTRIFUGE reduces backward and end-to-end training costs by 40.0% and 17.9%~ respectively.
5.1 EXPERIMENTAL SETUP
Testbed setup. We evaluate CENTRIFUGE Ubuntu servers, each equipped with 8 NVIDIA RTX 3090 GPUs (24GB), 40 CPU cores, 256GB of memory, PyTorch 2.8.0 and CUDA 12.8. We implement gradient accumulation to facilitate large batch training and use mixed precision with BF16 to reduce memory consumption and accelerate training.
Datasets, models, and tasks. Our experiments focus on fine-tuning pre-trained foundation mod- els for downstream tasks. We utilize small but powerful LLMs in the experiments: TinyLlama- 1.1B (Zhang et al. 2024), Qwen2.5-1.5B, Qwen2.5-7B (Yang et al. 2024), Qwen3-14B, Qwen3-
7

---

Published as a conference paper at ICLR 2026
Only Filtering Loss Stramman Approach Centrifuge
Only Fillering Loss
Centrifuge
Training Luws
1.4
Training Loss
1.8-
1.6
1.0
1.4
0.0
0.2
0.4
0.6
0.8
1.0
of Tokens
le6
of Tokens
le6
(a) Training loss while using reference model trained
(b) Training loss while using publicly available model blend of manually curated math-related tokens. from Lin al. (2024) reference.
Figure 4: The training convergence on mpth-web-math dataset while using different reference mod- els. CENTRIFUGE maintains the utility of token filtering method while the strawman approach fails to converge due to the unintentionally harmed öQ.
Training
Evaluation Tasks (using TinyLlama- 1.IB as foundation model)
Method
Dataset
GSMSK
MATH
SVAMP
ASDiv
MAWPS
TAB
MQA
MMLU
SAT
Average
Finetuning
NA
2.3
2.4
9.9
18.1
20.2
8.8
22.1
17.9
21.9
13.7
Regular
OWM
3.6
4.2
19.1
31.5
36.2
14.7
10.3
21.7
18.8
17.8
Finetuning
(Full)
CENTRIFUGE
OWM
11.8
6.4
35.5
47.4
62.8
22.4
15.3
18.7
25.0
27.3
(Ours)
(Filter 50%)
Table 1: Model performance on different tasks. Compared with regular training, CENTRIFUGE
significantly improves the model performance by up to 26.6% single task and 9.5% on average.
32B Team (2025).
Llama3.2-33 3.B-BB Llama3.1-8B (Dubey al. 2024b), and ALIA-40B Gonzalez-
Agiire et al. (2025) Following previous work (Lin et al. 2024, we first train a reference model on
small yet high-qualit datasets, and then use the loss from the reference model to filter tokens during the training of the target model. We focus on mathematical reasoning as the main task, employing
blend of synthetic and manually curated math-related tokens (Yu et al. 2024 Yue al. 2024
Mitra et al. 2024) Amini et al. 2019a Wang et al. 2024) as high-quality data to train the reference model. For large-scale datasets to train the target model, we use open-web-math (OWM) (Paster itt al. 2024). We follow (Lin et al.) 2024) and use the same architecture for both the reference and target models. To evaluate the utility of the models, we use the following tasks: GSM8K (Cobbe 2021). MATH (Lightman et al. 2024), SVAMP (Patel et al., 2021). ASDiv (Miao et al., MAWPS KKoncl-Keeziorrs et al. 2016), TabMWP (Lu et al. 2023), MathQA (Amini 2019b), SATMath (McAleer 2023). and MMLU (Hendrycks et al. 2021
et al.
2020
al.
Baselines and hyperparameters. We compare CENTRIFUGE with the following baselines: (1) regular training, which involves training the model without token filtering; and (2) token filtering that only filter loss (Lin et al., 2024), i.e., the backward filtering that only perform filtering during loss computation. In convergence evaluation, we aggregate different samples into context length of 2048, set the batch size to one million tokens, and use a learning rate of 5 x 10 10-5 with cosine decay. In system efficiency evaluation, report throughput time consumption averaged on ten iterations after two iterations of warming up.
5.2 CONVERGENCE EVALUATION
Figure illustrates the convergence comparison between token filtering that only processes loss, the strawman approach, and CENTRIFUGE on the open-web-math dataset while using different ref- erence models. The results demonstrate that CENTRIFUGE maintains the same utility with token filtering that only filters the loss. Our training system leverages mainstream memory tention implementation and the strawman approach has failed to converge since it unintentionally harms the s gradients during the backward computation
at-
have discussed in $4.1
To further demonstrate that CENTRIFUGE achieves better utility than regular training, we compare the performance on math-related tasks and the results presented in Table [ The results demon- strate that CENTRIFUGE improves model performance by 9.5% on average. For single tasks, the performance of CANTRIFUGE surpasses that of regular training by up to 26.6% (MAWPS task).
8

---

Published as a conference paper at ICLR 2026
Throughput (tokenssssccood)
50k
Regular Training
SSN
Centrifuge
1.6
TinyLlamaa-LL (ZeRO1 DF:4)
Qwen2.5-1.5U (ZeROI DP~1)
40k
1.20x
1.21x
1.4
1amu33238 (LoRA Singlne
Speedup
30k-
1,12x
1.2
20k-
1.0
10k
0.6
0
1K
2K
4K
20%
30%
40%
50%
60%
Context Length
Filtering Ratio
(a) Throughput of regular training and CENTRIFUGE
(b) The end-to-end speedup of CENTRIFUGE com- pared with regular training on three models with dif- ferent filtering ratios. The speedup of CENTRIFUGE
TinyLlama model with different context lengths.
CENTRIFUGE shows superior efficiency, with its effi-
ciency advantage growing at longer contexts.
linearly increases with the filtering ratio.
Figure 5: Efficiency evaluation different context length and filtering ratio.
Training
Time Consumption of Processing Million tokens (seconds)
Model
Method
Forward
Filter Operator
Backward
Total
TanyL
Regular Training (DP=4)
24.75
53.63
79.63
(1.1B, 4K)
CENTRIFUGE
24.60
2.98
32.17 () 40.0%)
60.36 (} 24.2%)
Qwen2.5
Regular Training (DP=4)
16.03
/
35.97
52.86
(1,5B, 2K)
CENTRIFUGE
15.40
3,33
21.26 (L 40.9%)
41.22 (} 22.0%)
Llama3.2
Regular Training (LoRA)
34.27
/
42.13
76.80
(3B, 2K)
CENTRIFUGE
34.17
4.23
23.96 (C 43.1%)
63.08 (L 17.9%)
Qwen2.5
Regular Training (TP=4)
362.54
/
811.26
1173.81
(7B, 2K)
CENTRIFUGE
378.61
12.84
442.25
(L 45.5%)
833.72 (L 28.9%)
1
Regular Training (TP=8)
519.50
/
1038.63
1558.14
(8B, 2K)
CENTRIFUGE
520.27
14.35
529.39
(1 49.0%)
1064.02 (4 31.7%)
Qwen3
Regular Training (TP=8)
810.35
/
1953.37
2764.03
(14B, 2K)
CENTRIFUGE
809.79
16.94
977.27 (4. 49.9%)
1804.22 (4 31.7%)
Qwen3
Regular Training (TP=8)
95.41
/
191.7
287.16
(32B, 4K)
CENTRIFUGE
95.42
9.83
117.49 (4 38.7%)
222.81 (L 22.4%)
ALIA
Regular Training (TP=8)
116.48
/
234,949
351.52
(40B, 4K)
CENTRIFUGE
116.89
14.54
133.04 (L 43.4%)
264.58 (L 24.7%)
Table 2: The detailed time consumption on four models using different training methods, including distributed training using DP and TP, and parameter-efficien training using LoRA on single GPU. When filtering 50% tokens, CENTRIFUGE reduces backward time and end-to-end training time by 40.0% and 17.9%~<34.79 respectively. The results for Qwen3-32B and ALIA-40B are obtained on testbed comprising eight H20-96GB GPUs interconnected with NVLink, a 96-core CPU, and 1.2TB of memory.
5.3 SYSTEM EFFICIENCY EVALUATION
While evaluating the efficiency of CENTRIFUGE, we divide the training process into three stages: forward, filtering operator (i.e., updating the graph), and backward. Table presents detailed time comparison for these three stages under different scales of models and training methods. In distributed training using ZeRO1 (Rajbhandari et al. 2020) data parallel (DP) and tensor par- allel (Narayanan et al., 2021) (TP), CENTRIFUGE significantly improves efficiency of backward Particularly, CENTRIFUGE achieves more speedup in TP since the communication size in TP is also linearly decreased after activation filtering. Thus, CENTRIFUGE extremely useful in distributed training not only for de- creasing the computational cost but also reducing the communication size. Similarly, CENTRIFUGE can also reduce the communication size of pipeline parallel (PP) and mixture-of-expert (MoE) par- allel, which is extensively discussed in appendix D We also evaluate CENTRIFUGE on parameter- efficient training using low ow-k adapter (Hu et 2022) (LoRA) on a single GPU. The results show that CENTRIFUGE also significantly improves the efficiency by 43.1 43.1% in backward computation and 17.9% in end-to-end LoRA training, while setting rank size to 64 and finetuning the attention mod- ular (i.e., projections of q. k, and o).
computation by 40%%~49.944 and end-to-end training by 22%

---

Published as a conference paper at ICLR 2026
Overhead of the filtering operator. We report the overhead of updating the graph in Table Our current implementation updates the graph layer by layer. Consequently, the associated compu- tational cost increases with the model size, which explains why the larger models in Table Tableetend to exhibit higher overhead than the smaller models. However, our experimental results indicate that the cost of the filtering process is substantially lower than the reduction in training time, resulting in clear improvement in end-to-end efficiency. In particular, the overhead of updating the graph can be further reduced by overlapping the filtering process with computation using two micro-batches, i.e., while one updates the graph, the other executes the forward computation. Additionally, the offline preparation cost of our system is single forward pass, for getting the overall graph structure, which is minimal compared to the entire training process.
Impact of the context length. The context length critical factor influencing the training ef- ficiency of LLMs, as the training complexity increases quadratically with the context length. We evaluate the efficiency of CENTRIFUGE using different context lengths on the TinyLlama model. Figure ((a) illustrates the throughput of regular training and CENTRIFUGE across various context lengths ranging from 1K to 4K. The throughput of CENTRIFUGE consistently exceeds that of reg- ular training, with the efficiency improvement becoming more pronounced as the context length increases. At a context length of 4K, CENTRIFUGE achieves 1.40x higher throughput. These re- sults demonstrate that CENTRIFUGE can effectively enhance efficiency in computationally intensive scenarios, such as training of long context and large models.
Impact of the filtering ratio. To further investigate the efficiency improvement of CENTRIFUGE, we evaluate the end-to-end speedup compared to regular training under different filtering ratios, with the results shown in Figure 5(b) The speedup of CENTRIFUGE increases linearly with the filtering ratio, demonstrating the effectiveness CENTRIFUGE 's system design and indicating potential per- formance gains in scenarios with higher filtering ratios (e.g., long-context training).
6 DISCUSSION
CENTRIFUGE with more models. The core designs of CENTRIFUGE are generalizable across diverse model architectures and training frameworks. In the implementation, we adopt an approach that minimizes dependence specific models and frameworks sppes-iccallll..... by directly updating nodes in the computation graph. When different models and frameworks are implemented using the same backend (e.g., PyTorcc),, the underlying computation graph nodes are consistent. Com- bined with our design for autonomously processing graph nodes, which is detailed demonstrated in Section B our system can seamlessly support models and frameworks within the same backend.
CENTRIFUGE with MoE. We provide detailed analysis in Section D that CENTRIFUGE can bene- fit different parallelism strategies (e.g., .. PP, SP, and MoE) by reducing the communication size. We would like to discuss more on MoE since it inherently faces load balancing challenges and integrat- ing CENTRIFUGE with MoE could lead to an interesting research problem. While token filtering would improve overall efficiency, its impact on load balancing remains unclear, as the number of filtered tokens may vary across experts. Therefore, effectively addressing the load balancing issue may be essential to maximize the efficiency gains achieved through token filtering. We identify this as a valuable direction in the future.
CENTRIFUGE with forward filtering. Regarding forward filtering, we consider it more suitable for long-sequence training. Existing studies have demonstrated that forward filtering can compro- mise utility since the context for modeling each token is reduced (Zhong et al. 2023 Yao et al.)
2022 Lin et al., 2024). and this problem could be more severe with short sequences. However, for longer sequences (e.g.. 128K), recent studies (DeepScek-A]] 2025) have demonstrated impressive performance in context compression during training. Therefore, integrating sengeeeuennce forward filtering solutions with CENTRIFUGE represents promising direction for improving the efficiency of long- sequence training in the future.
7 CONCLUSION
In this paper, we propose CENTRIFUGE, a system that unlocks the full efficiency of token filtering in LLM training. CENTRIFUGE maintains sparsity by further filtering the activations and trans- forms sparse GEMM into dense GEMM to optimize efficiency in existing ML libraries. Extensive evaluations demonstrate that CENTRIFUGE effectively achieves its design targets.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We thank the anonymous reviewers for their valuable feedback and suggestions. This work is sup- ported by NSFC 62402407, Hong Kong RGC TRS T41-603/20R, and Turing AI Computing Cloud (TACC) Xu et al. (2025) Junxue Zhang and Kai Chen are the corresponding authors.
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Gre- gory S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, lan J. Good- fellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Józefowiez, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mané, Rajat Monga, Sherry Moore, Derek Gor- don Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul A. Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda B. Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. Tensorflow: Large-scale machine learning heterogeneous distributed systems. CoRR, abs/1603.044p) 2016.
Aida Amini, Saadia Gabriel, Shanchuan Lin, Rik Koncel-Kedziorste Yejin Choi, and Hannaneh Hajishirzi. Mathqa: Towards interpretable math word problem solving with operation-bant based for- malisms. (1), pp. 2357-2367. Association for Computational Linguistics, 2019a.
Aida Amini, Saadia Gabriel, Shanchuan Lin, Rik Sancel-Kedziorski Yejin Choi, and Hannaneh Hajishirzi. Mathqa: Towards interpretable math word problem solving with operation-bact based for- malisms. NAACL-HLT (1), pp. 2357-2367. Association for Computational Linguisties, 2019b.
Foozhan Ataiefard, Walid Ahmed, Habib Hajimolahoseini, Saina Asani, Farnoosh Javadi, Moham- mad Hassanpour, Omar Mohamed Awad, Austin Wen, Kangling Liu, and Yang Liu. Skipvit: Speeding up vision transformers with token-level skip connection. CoRR, abs/2401.15293, 2024.
Tianyi Bai, Ling Yang, Zhen Hao Wong, Jiahui Peng. Xinlin Zhuang, Chi Zhang, Lijun Wu, Jiantao Qiu, Wentao Zhang, Binhang Yuan, and Conghui He. Multi-agent collaborative data selection for efficient LLM pretraining. CoRR, abs/2410.08102, 2024.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. CoRR, abs/21 10. 14168, 2021.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In NeurlPS, 2022.
DeepSeek-AL. AI. Deepseck-v3.2-1 Boosting contextntext efficienccy with deepseek sparse attention, 2025.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. Advances in neural information processing systems, 36: 10110008800131 15, 2023.
Ning Ding, Yujia Qin, Guang Yang, Fuchao Wei, Zonghan Yang. Yusheng Su, Shengding Hu, Yulin Chen, Chi-Min Chan, Weize Chen, Jing Yi, Weilin Zhao, Xiaozii Wang, Zhiyuan Liu, Hai-Tao Zheng, Jianfei Chen, Yang Liu, Jie Tang, Juanzi Li, and Maosong Sun. Parameter-efficien fine- tuning of large-scale pre-trained language models. Nat. Mac. Intell. 5(3):220-23,, 2023.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv.2407.21783 2024a.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv e-prints, pp. arXiv-2407, 2024b.
11

---

Published as a conference paper at ICLR 2026
Simin Fan, Matteo Pagliardini, and Martin Jaggi. Doge: domain reweighting with generalization estimation. In Proceedings of the 41st International Conference on Machine Learning. ICML 2.. JMLR.org, 2024.
Aitor Gonzalez-Agirre, Mare Pamies, Joan Llop, Irene Baucells, Severino Da Dalt, Daniel Tamayo, José Javier Saiz, Ferran Espuña, Jaume Prats, Javier Aula-Blasco Mario Mina, Adrián Rubio, Alexander Shvets, Anna Sallés, Iñaki Lacunza, lñigo Pikabea, Jorge Palomar, Júlia Falcão, Lucía Tormo, Luis Vasquez-Reina, Montserrat Marimon, Valle Ruíz-Fermández, and Marta Villegas. Salamandra technical report, 2025. URL htttps://aaxiv.orrrrrrr/ 250...00000//b0/000000000044
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In ICLR. OpenReview.net, 2021.
Le Hou, Richard Yuanzhe Pang, Tianyi Zhou, Yuexin Wu, Xinying Song, Xiaodan Song, and Denny Zhou. Token dropping for efficient BERT pretraining. In ACL (1), pp. 3774-3784, Association for Computational Linguistics, 2022.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang. and Weizhu Chen. Lora: Low-rank adaptation of large language models. In ICLR. OpenRe- view.net, 2022.
Rabeeh Karimi Mahabadi, James Henderson, and Sebastian Ruder. Compacter: Efficient low-rank hypercomplex adapter layers. Advances in neural information processing systems, 34:1022-1035, 2021.
Rik
Subhro Roy, Aida Amini, Nate Kushman, and Hannaneh Hajishirzi. MAWPS: A math word problem repository. In HLT-NAACL, pp. 1152-1157. The Association for Computational Linguistics, 2016.
Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mo- hammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. In MLSys, misys.org, 2023.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint 1xXiv:2101.000 2021.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, lly Sutskever, and Karl Cobbe. Let's verify step by step. In ICLR. OpenReview.net, 2024.
Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Ruochen Xu, Chen Lin, Yujuu Yang, Jian Jiao, Nan Duan, Weizhu Chen, et al. Not all tokens are what you need for pretraining. The Thirty- eighth Annual Conference on Neural Information Processing Systems, 2024.
Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng. Longxu Dou, Tianyu Pang. Jing Jiang, and Min Lin. Regmix: Data mixture as regression for language model pre-training. In The Thirteenth International Conference Learning Representations, 2025. URL https:
Sebastian Loeschcke, Mads Toftrup, Michael Kastoryano, Serge Belongie, and Vésteinn Snæb- jamaarson. Loqt: Low-rank adapters for quantized pretraining. Advances in Neural Information Processing Systems, 37:1 17:1112222-1555 115308. 2024.
Pan Lu, Liang Qiu, Kai-Wei Chang, Ying Nian Wu, Song-Chun Zhu, Tanmay Rajpurohit, Peter Clark, and Ashwin Kalyan. Dynamic prompt learning via policy gradient for semi-structured mathematical reasoning. In ICLR. OpenReview.net, 2023.
Stephen McAleer, Sat multiple choice math may 23. https: httpt::tttttttttts http:s////oggngggfcc..oo//tttttt mcaleste/sal mult. 1
pelppplel.l...... ple choice mathh may 23, 2023.
Shen- Yun Miao, Chao-Chun Liang, and Keh- Yih Su. A diverse corpus for evaluating and devel- oping english math word problem solvers. In ACL, pp. 975-984, Association for Computational Linguisties, 2020.
12

---

Published as a conference paper at ICLR 2026
Arindam Mitra, Hamed Khanpour, Corby Rosset, and Ahmed Awadallah. Orca-math: Unlocking the potential of slms in grade school math. CoRR, 0en/2402.148800, 2024.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vi- jay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on GPU clusters using megatron Im. In SC. pp. 58. ACM, 2021.
Keiran Paster, Marco Dos Santos, Zhangir Azerbayev, and Jimmy Ba. Openwebmath: An open dataset of high-quality mathematical web text. In ICLR. OpenReview.net, 2024.
Arkil Patel, Satwik Bhattamishra, and Navin Goyal. Are NLP models really able to solve simple math word problems? In NAACI-HHIT, pp. 2080-2094. Association for Computational Linguis- tics, 2021.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Perfor- mance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Am- mar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-moc: Advancing mixture-of-expeee inference and training to power next-generation ai scale. In International conference on machine PMLR, 2022.
learning, pp.
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System opti- mizations enable training deep learning models with over 100 billion parameters. In KDD, pp. 3505-3506. ACM, 2020.
Qwen Team. Qwen3 technical report, 2025.
Tristan Thrush, Christopher Potts, and Tatsunori Hashimoto. Improving pretraining data using per- plexity correlations. In The Thirteenth International Conference Learning Representations, 2025. URL httpppp://peeneerreeeeeeeee/eeee http://wweeeoouuuuooooouuuuooo/o
Mojtaba Valipour, Mehdi Rezagholizadeh, Ivan Kobyzev, and Ali Ghodsi. DyLoRA: Parameter- efficient tuning of pre-trained models using dynamic search-free low-rank adaptation. In Andreas Vlachos and Isabelle Augenstein (eds.), Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics, pp. 3274 3287, Dubrovnik, Croatia, May 2023. Association for Computational Linguistics. doi: 10. 18653/11/0023.callmman..... URL /attps://aclie www.aaaanhoaaooooooo wow.ooogoooooooooooooooooooooooo ogy.org wwt.lorgg2222..em 20w.:222..222e.. 2397
Peiyi Wang, Lei Li, Zhihong Shao, Runxin Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce lims step-by-step without human annotations. In ACL
(1), pp. 9426-9439. Association for Computational Linguisties, 2024.
Guillaume Wenzek, Marie-Anne Lachaux. Alexis Conneau, Vishrav Chaudhary, Francisco Guzmán, Armand Joulin, and Edouard Grave CCNet: Extracting high quality monolingual datasets from web crawl data. In Proceedings of the Twelfth Language Resources and Evaluation Conference, pp. 4003-40440 4012, Marseille, France, May 2020. European Language Resources Association. ISBN 979-10095546-34-4. htrLl//ggee///llaatttoooooooo///
1.494/
Alexander Wettig, Aatmik Gupta, Saumya Malik, and Danqi Chen. Qurating: selecting high-quality data for training language models. In Proceedings of the 41st International Conference on Ma- chine Learning, ICML'24. JMLR.org. 2024.
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoe V Le, Tengyu Ma. and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL htttts::///aaeeeeeeeeeetttt formm? i d= dd=XxuBBUeHHha
Kaigiang Xu, Decang Sun, Hao Wang, Zhenghang Ren, Xinchen Wan, Xudong Liao, Zilong Wang, Junxue Zhang, and Kai Chen. Design and operation of shared machine learning clusters on campus. In Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume I, pp. 295-310, 2025.
13

---

Published as a conference paper at ICLR 2026
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang. Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue, Na Ni, Pei Zhang. Peng Wang. Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yugiong Liu, Zeyu Cui. Zhenru Zhang, and Zhihao Fan. Qwen2 technical report. arXiv preprint wxXiv:2477.10677 2024.
Zhewei Yao, Xiaoxia Wu, Conglong Li, Connor Holmes, Minjia Zhang, Cheng Li, and Yuxiong He. Random-ltd: Random and layerwise token dropping brings efficient training for large-scale transformers. CoRR, abs/2211. 1.1666,, 2022.
Jiasheng Ye, Peiju Liu, Tianxiang Sun, Jun Zhan, Yunhua Zhou, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. In The Thirteenth Inter- national Conference on Learning Representations, 2025. URL
Forum 3CF
Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T. Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. In ICLR. OpenReview.net, 2024.
Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. Mammoth: Building math generalist models through hybrid instruction tuning. In ICLR. Open- Review.net, 2024.
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model, 2024.
Qihuang Zhong, Liang Ding, Juhua Liu, Xuebo Liu, Min Zhang, Bo Du, and Dacheng Tao. Re- visiting token dropping strategy in efficient BERT pretraining. In ACL (7), pp. 10391-10405. Association for Computational Linguistics, 2023.
14

---

Published as a conference paper at ICLR 2026
A THE USE OF LLMS IN WRITING
We used L.M,, namely DEEPSERK-R R1 1, polish the writing of this manuscript. No other generative AI functionality is used in the writing of this submission.
B
IMPLEMENTATION DETAIL OF DIMENSION-REDDCED GEMM
B.I DYNAMIC GRAPII COMPLICATES THE TRANSFORMATION
Based on observations from backward computations, sparse matrix operations can be reformulated into dimension-reduced dense computations. To implement this proposed approach, it is crucial to first understand the functionality of existing automatic differentiation (autograd) libraries. For example, PyTorch, one of the most widely used frameworks, employs a dynamic computational graph that constructed incrementally as operations are executed. Other machine learning libraries (e.g., TensorFlow (Abadi et al.] (0166) and systems (e-g.. MegatronI.M (Narayanan et al. [22T) and DeepSpeed (Rasley et al. (020 also use a similar graph-based approach or are mostly built on PyTorch. The backward graph is built during forward propagation and is utilized only once to compute gradients (i.e., discarded after the backward pass in each iteration).
Backward token filtering occurs after the forward computation, at which point the computation graph has already been constructed. Therefore, we need to update the computation graph node by node (e.g., updating the sizes and variables) prior to performing the backward computation. However, modifying the computational graph poses significant challenges due to the following reasons:
Dynamic graph structure. The computational graph is dynamically built. Different implementa- tions of the same algorithm can have significantly different backward computation graph. Even the input and output can impact the graph, e.g., FlashAttention (Dao et al. 2022) only accept model weights in 16-bits and naive attention implementation is the only choice if need ex- plicitly output the attention values. The dynamic graph structure makes it impractical to design static updating rules based on the model (e.g.. designing static rules for updating self-attention and FFN layers).
Dynamic usage of the graph nodes. The same type of node can have different inputs and outputs in different models or in the same model. For example, the same multiplication node has different outputs when multiplying with scalar, vector, or matrix. The dynamic usage of the graph static updating rules based on the node types.
nodes makes it impractical
Numerous types of nodes. Different models typically have different computations and thus use different type of nodes. Different nodes require different updating logic. PyTorch, for instance, has over 300 node types, significantly increasing the complexity of system implementation.
In summary, effectively accelerating the token filtering requires us to transform the sparse GEMM to dimension-reduuced dense GEMM that requires updating the computation graph, which is chal- lenging due to the dynamic of graph structure, the dynamic usage of nodes, and the numerous types of nodes.
B.2 AUTOMATIC GRAPH UPDATING WORKFLOW
To address this issue, we propose an automatic workflow to amend the computational graph. The key insight that, even though the graph is highly dynamic, it still can be deterministic when using the same implementation and inputs. Particularly, the model implementation and inputs remain the same during the whole training (i.e., runtime stability). Thus, we can mimic the input and traverse the graph using the same implementation to dynamically determine the node information. The automatic workflow contains two steps: 1) generating the skeleton code for processing each type of nodes and their attributes; 2) leveraging special markers to generate detailed node processing rules.
Specifically, we first perform a coarse-grained graph traversing using synthetic data (i.e., mimicking the actual inputs) to collect all the node types and find the target attributes (c.g., sizes and variables) that need to be updated. Table shows examples of the node attributes that need to be amended. We select the list of target attributes and corresponding data types based on Torchgen, which could be easily updated if more attributes needs to be processed. We generate skeleton codes for processing attributes of all the nodes and the implementation detail is presented in C.11
15

---

Published as a conference paper at ICLR 2026
Node Attributes InputMetadata SavedVariables
Data Type Int Array
Operations
Update size
Tensor
Reduce Dimensions
Matrix Sizes
Int Array
Update size
Matrix # of elements
Int
Update value
Table 3: Examples of the nodes' attributes that need to be updated during amending the graph.
Next Nodes
query 32, 77, 64)
key
ScaledDeetroddutisssbbae entionBackkwardl
32, 27,
valuuc
32, 27,
TransooeRackwaard dimn:: diml:
logaummxxp
ViewBackkwadd self sym_ size: (19, 77, (2) 64)
Attn. peuj weight (2048, 2048)
philoxx
philox offist
outpur
ViewBackward0
AccumulateGrad
32, 77,
sell sym sioc: (79, 2048)
TRackwardo
self (1463, 2048)
MrBBackward
self sym (7463, 2048) nut2 sym size: (2048, 2048)
ma12
(1463, 2048)
Nodes
Special marks: Bsz=19, Seq=77, Ssq=
Figure 6: An example of marking batch size and sequence length with prime numbers. Leveraging the special marks is cample-but-effeeee way precisely find the shrinking dimensions of various nodes in the computational graph.
After obtaining the skeleton code for updating the node attributes, we still need to determine the updating logic. Specifically, we need to determine which dimensions should be reduced and the sizes after the reduction. We only focus on the sequence dimension since we need filter out unimportant tokens. However, the index of the sequence dimension varies among different nodes and may also mixed with the batch size. Moreover, the same type of node can have different dimensions at different positions in the same graph, making it impractical to use static updating logic. To address this issue, we design a simple-but-effeee
method by marking the batch size and sequence length with special numbers to precisely find the shrinking dimensions of various nodes in the computational graph. Figure G shows an example of marking the batch size and sequence length with prime numbers. In the skeleton code, and sequence dimension. Directly using the greedy algorithm can produce wrong results as the batch size or sequence length may be identified with other dimensions (e.g.. the sequence length and hidden dimension could be both 2048). The special markers avoid the ambiguity of the dimension and enable the greedy algorithm to precisely find the shrinking dimensions and determine the size after the reduction. The system will cache the output of greedy algorithm for online training ($C.1)
greedy algorithm find the batch size
C IMPLEMENTATION
We implement CENTRIFUGE in PyTorch, one of the most widely used frameworks, and use its C++ extension to create the backward filtering operator. CENTRIFUGE improves the efficiency of token filtering through two designs: filtering the activations and transforming sparse GEMM to dense GEMM. We implement these two designs in single operator by directly reducing the quence dimension, as the filtered activations (i.e., those set to zero) will be subsequently removed in the transformation from sparse GEMM dense GEMM. Thus, we can directly remove the acti- vations instead of setting them to zero in advance. To address the challenges posed by the dynamic
'C++ extensions in Py Porch allow users to create custom operators outside the PyTorch backend, providing flexibility and reducing boilerplate code. Once defined, these extensions can he organized into native PyTorch functions for upstream contributions.
16

---

Published as a conference paper at ICLR 2026
(e.g. Llama)
Target Node Attributes
Operator Skeleton
Bsz/Seq Marker
Model-
customized
Operator
Offline Preparation
Globally Reudce Dimensions
Forward
Backward Graph
Update
Next
Model
Fiftered Graph
Efficient Online Training with Token Filtering
Figure 7: Implementation and usage of CENTRIFUGE.
computation graph and to support various LLM architectures, we implement CENTRIFUGE two phases: the offline and online stages, as illustrated in Figure 7 In the offline stage, CENTRIFUGE employs an automatic workflow to generate a model-customized operator for updating the graph. In the online training ($C.I), the operator filters the activations and transforms the sparse GEMM into cinnension-reducee dense GEMM to accelerate the training process ($C.2). CENTRIFUGE can also be implemented in other frameworks, such as TensorFlow (Abadi al. (2066), by following similar procedures to update the computation graph during backpropagation.
C.1 OFFLINE GENERATING MODEL-CUSTOMMIEER OPERATOR
Given the model, run forward and backward computations using synthetic data (i.e., simulating the training samples) to obtain all the node types. Due to runtime stability (i.e., the graph remains stable during training), the node information from synthetic data is identical that during train- ing. We then parse the node attributes from Torchgen and generate the operator's skeleton code for processing each node. Code [ shows examples of the generated skeleton code for processing GEMM and FlashAttention nodes. The generated code is skeleton and cannot be used directly be- cause the operator employs a greedy algorithm to determine the batch size and sequence dimension based the inputs. However, the actual batch size and sequence length might be the same other dimensions (e.g., hidden states of 2048), which can mislcad the operator into reducing the wrong dimensions.
start
generstion +/
if (fn->name () =s "#msockward0")
MaRackwardO op_ dynamt wttaac/caatstaaaaaaaarrlddoddo
auto urpacked_selI op_[n u
if [aapaccke__ss f.defined())
op sp__n=>sei__ gieee__ileeerepeeeeeee.......
tuspacked_ falsel;
auto unpacked_mat2 op__fa-aaat.. (ungscke6_mat22defiiee fn- graph_tilter--prreeee
www.rraaaassssasalleeeooaaaa (urpaclee_tt mot2, falsely
graph_filter sptpeeesppppessssseissoooe htg.:assooop.a..e...e........ee.
(ff >nane wwc.edeeerrdellleellaaaaaaoooeee
op_fn
auto urpacked_quie -
on__nn__ graps_Eiiter-pprcees.cee
(uapacked false);
auto urpacked_key op_ -o__fn->eey__uupcckf
[unpscke@_key.deefieedll
fn->kyy___ = graph__ilter-ppreeee www.ipreeeess_vvialllleleeeeeee (unpzcked false);
auto urpacked_valise -p__n__xall_____npppa..
junpaoket__allue.deeiine
ep_in ovalue_ grapi_filter wprccess_variaalle (uaesaablleeaaaccedddeellll false);
unpac wtg._aatoooooppttttt..m..m ott.--oottt...nkkkooooogggeeppp
sanpackad_.oupptt.ddfinee ep op_in->output__ gtappp.illeeee-ssrsssssss..csse
otg.ssssserrbbleemmpccee..oooos (unpribiipeppacckeeoooouuu true);
..... more attrieutes onitted
fraah_filter-pprrrees..
httaa-pprroee_ppfeet_____aaaa fn-
17

---

Published as a conference paper at ICLR 2026
graph_ Srephh_flleer-ypoceess..s (np____)max__))
// more nodes
< 68 end of code generation $8 A/
Code Generated skeleton code for processing GEMM and FlashAttention nodes. The automatic code generation enables CENTRIFUGE to support various nodes with low implementation costs (e.g., PyTorch has more than 300 types of nodes).
To solve this issue, as demonstrated in B.2 we compile the generated skeleton code and run the operator using inputs with special markers for both batch size and sequence length. These special markers (i.e., unique from other dimensions) enable the greedy algorithm to precisely identify the correct dimensions for reduction. The operator will the output of the greedy algorithm and load it during online training, which is guaranteed to be correct due to runtime stability.
After executing the above workflow, we obtain a model-customized operator that contains all the node information and corresponding dimension updating logic for specific model. We have pre- pared scripts that allow users to easily execute the workflow and generate operators for their own models. The reset system implementation includes updating the node attributes, e.g.. changing the InputMetadata to pass verification and update the saved variables, which is quite straightforward as long the attributes and reducing dimensions are correctly identified.
C.2 ONLINE TRAINING USING CENTRIFUGE
import centrifugs
for step, betch in enumerate dataaatt logits modda (batch "input, (ds")) logits
loss causal (batch[" (aaucch"innnt.ffi" inputt.ids", logits) loss, filter.mack token_filter.loss
batch( "input ids"), logits,
zef ross-bbbba (o=s"_ drop.rate=0.4.
contrifuge, ops www.ppccssaee.e.oeeppopoppccpppp backward filter (loss, filter (ackward
opt pptimizer.ssepp optimizer.zero_gerr
Code 2: Using CENTRIFUGE in the online training.
Using the operator only requires adding a few lines (i.e., five lines) of code. 2ode shows an exam- ple of using CENTRIFUGE in online training. Starting from regular training, first need to change the loss computation to loken-filtered loss, i.e., only considering the loss on selected tokens. Then, call the CENTRIFUGE operator using the loss and filter mask to update the graph. The filter mask a tensor consisting of zeros and ones to indicate which tokens are filtered. For systems that already utilize token filtering, only one line of code cent r fuge ops www.ward........................ filter (loss, fi teteer (tr_massk))) is needed to get the full efficiency of token filtering.
C.3 HOSTING
The system is open-sourced and hosted on GitHub. The link httpp:://githuu.cmm//.... Chai/Centrifuge
D CENTRIFUGE REDUCES COMMUNICATION OVERHEADS IN DISTRIBUTED TRAINING
CENTRIFUGE optimizes computational efficiency in backward token filtering by transforming sparse GEMM into dense GEMM, which significantly reduces communication overheads in distributed LLM training. Specifically, CENTRIFUGE updates the entire computation graph, where the gradi- ents and activations are reduced along the sequence dimension. Existing parallelism strategies that enable LLM training on distributed systems typically transfer gradients between different nodes or GPUs during backward computation. Therefore, reducing the sequence dimension of gradients can linearly decrease communication overheads in distributed training. We discuss the advantages of CENTRIFUGE across different parallelism strategies follows.
18

---

Published as a conference paper at ICLR 2026
Tensor Parallel (TP) (Narayanan al. 2021). In TP, the transformer models are typically par- titioned along the multiple heads and hidden dimensions. All-reduce of gradients on inputs are required twice (i.e., inputs of FFN and attention block) in each layer's backward computation. CENTRIFUGE can linearly reduce the amount of data transferred in each all-reduce operation. Our evaluation results training Ilama3. I-8BB when using TP=8 show that CENTRIFUGE im- proves end-to-end training efficiency by 31.7% when filtering 50% tokens.
Sequence Parallel (SP) (Korthikanti al. 2023). SP is designed to be combined with TP to further reduce the memory usage caused by the redundant activations of dropout and layer nor- malization. In SP, the sequence dimension is partitioned
multiple devices through all-gather during computing dropout and layer normalization and recovered to partition on hidden dimen- sions through reduce-scatter. CENTRIFUGE can reduce the communication overhead in the cor- responding all-gather and reduce-scatter operations.
Pipeline Parallel (PP) (Narayanan et al. 2021). The advantages of using CANTRIFUGE in PP is straightforward as PP sequentially transfers the gradients in the graph which are linearly reduced
by CENTRIFUGE.
Mixture-of-ENae Expeerss (MoE) (Rajbhandari et al., 2022). The integration of CENTRIFUGE in MoE helps in optimizing communication between experts during the backward passes, thereby mini- mizing the data transferred across devices and enhancing throughput. Specifically, CENTRIFUGE ensures that only gradients of important tokens routed to the corresponding experts. Con- currently, we also find that integrating our system with MoE leads to new research problem concerning load balancing, which is discussed in the paper.
19


Reviews:
[{"review_id": "qN5qIWWCCT", "reviewer": "Reviewer_6NpQ", "summary": "This paper proposes CENTRIFUGE, a system that unlocks the full efficiency of token filtering in LLM training through algorithm-system co-design. At the algorithm level, it filters activations in the attention backward kernel to amplify sparsity while maintaining compatibility with memory-efficient attention implementations. At the system level, it transforms sparse GEMM into dimension-reduced dense GEMM through automatic graph updating. Experiments on four models (TinyLlama-1.1B to Llama3.1-8B) show up to 49% backward time reduction and 31.7% end-to-end speedup when filtering 50% tokens, while achieving 26.6% utility improvement over standard training.", "questions": "1. Table 2: Why does the filter operator overhead vary significantly (2.98s for TinyLlama vs 14.35s for Llama3.1)? Is this linear in model size or are there other factors?\n2. Why only evaluate Qwen2.5-1.5B but not Qwen2.5-7B, especially when Llama3.1-8B is included?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper clearly identifies why existing token filtering methods fail to achieve efficiency gains (inadequate sparsity propagation and sparsity range mismatch with ML libraries), with solid empirical evidence.\n2. The paper evaluates across multiple dimensions: different model sizes, various training scenarios, context lengths, and filtering ratios, demonstrating consistent gains and broad applicability.", "weaknesses": "1. All experiments use models ≤8B parameters. Given that efficiency gains are most critical for larger models (70B+), the absence of such experiments is a significant limitation. The 8B model uses TP=8, but modern large-scale training uses more complex parallelism (TP+PP+DP). How CENTRIFUGE scales to 70B+ models with TP=8, PP=4 setups remains unknown.\n2. The automatic graph updating approach relies on \"runtime stability\" and special prime number markers, which may be fragile. The paper doesn't discuss failure cases, what happens when graph structure changes unexpectedly, or computational overhead of the offline preparation phase.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "V2dvARrcq0", "reviewer": "Reviewer_P85p", "summary": "This work tackles the problem of improving training efficiency for LLM via token filtering. Prior work showed token filtering can boost model utility/quality but fails to yield meaningful speedups in practice. The authors identify two key obstacles: 1. Insufficient sparsity propagation: existing token filtering (applied only at the loss layer) produces sparse gradients at the output, but upstream layers still process all tokens, negating potential speed gains. Naively filtering intermediate activations breaks standard attention backpropagation and harms gradient correctness. 2. Unsupported sparsity regime: Token filtering typically drops ~30-50% of tokens which leads to a \"moderate\" sparsity level at which standard sparse deep learning libraries (optimized for >=95% sparsity) become inefficient or even slower. To address these issues, the paper proposes an integrated solution called CENTRIFUGE, with contributions at both the algorithm and system levels:\n\n* Algorithm level: introduces a modified attention backward pass that filters out activation for unimportant tokens during backprop, amplifying sparsity in the gradient computation. This new backward kernel avoids interfering with the attention softmax by processing each output’s gradient separately, which preserves the correct gradient flow and the model quality gains of token filtering. This design is also compatible with high-performance attention implementations (e.g. FlashAttention)\n\n* System level: develops an automated workflow to convert sparse operations into smaller dense operations that can run efficiently on existing libraries. CENTRIFUGE exploits the structure of token sparsity (i.e. whole rows or columns dropped) to shrink the matrix multiplications, achieving the effect of sparse computation by adjusting tensor dimensions at runtime. The authors devise a robust method to dynamically alter the computation graph (despite PyTorch’s dynamic nature) by detecting stable graph patterns and updating tensor shapes before each backward pass. This allows the \"moderate\" sparsity to be leveraged for real speedups using standard GPU kernels instead of generic sparse GEMM routines which has high overheads.\n\nExperiments on four LLMs (1.1B to 8B parameters) fine-tuned on a large math reasoning dataset show that this method achieves significant speedup with on-par or improved model performance. For example, filtering 50% tokens cut down backward pass time by up to 49% and end-to-end training time by up to 31.7%. Test performance is improved by up to 26.6% on individual tasks (+9.5% on average) comparing to no-filtering baseline.", "questions": "Echoing points listed in Weaknesses section:\n1. How well do the authors expect CENTRIFUGE to scale to much larger models (e.g., 70B or 100B parameters) and training on massive corpora?\n2. Have the authors considered evaluating CENTRIFUGE on different domains or tasks beyond math reasoning, such as coding?\n3. Could some form of forward token filtering be incorporated in future to further boost efficiency, and what might be the impact on utility?\n\nIn addition, it would be also interesting to know:\n4. In the experiments, a fixed fraction of tokens (e.g. 50%) is filtered each step. Have the authors considered making the filtering threshold adaptive, for instance, filtering more aggressively as training progresses or varying by batch?\n5. Can CENTRIFUGE be combined with other training efficiency methods (e.g., MoE layers, gradient checkpointing (I think so), or data selection at the example level)? \n6. Considering the practical importance, do the authors plan to open-source CENTRIFUGE or integrate it into common libraries like DeepSpeed/FSDP, Megatron, HF/Lightning trainer?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The authors identify a crucial bottleneck (lack of real efficiency gains from token filtering) and provide a novel solution with algorithm and system co-design, effectively addressing both the ML aspect (ensuring gradients remain correct and useful) and the systems aspect (making use of sparsity in existing hardware/software). By further filtering attention activations in the backward pass in a safe way, they increase sparsity where it matters, and by transforming sparse ops to dense via graph manipulation, they avoid the usual pitfalls of sparse computation overhead. Together, these yield a unique solution that had not been explored before in literature on token filtering.\n2. The experimental evidence is compelling and comprehensive. The achieved speedups are significant (nearly half the backprop time eliminated, and ~30% overall training time reduction at 50% token drop); and importantly, the approach does not harm model quality but in fact improves it. The experiment coverage is also comprehensive: multiple models (ranging 1B to 8B parameters) and different training regimes (single GPU, multi-GPU with DP and TP, and fine-tuning with LoRA) are all tested. This diversity builds confidence that the method is robust and general.\n3. The work has high practical significance. Two facts of CENTRIFUGE lower the barrier to adoption: 1) it can be integrated with only one extra line of code for users already doing token filtering, and 2) it's compatible with common distributed training (DP+TP) and in fact further improve throughput by reducing gradient communication payload (Appendix D).", "weaknesses": "1. The experiments, although extensive for the provided settings, are limited to relatively small-scale dense LLMs and a specific domain (mathematical reasoning). It remains an open question how well CENTRIFUGE would scale to much larger LLMs (e.g. 70B+ params) or to training on more diverse, general-domain data. Larger LLM might introduce new bottlenecks or slightly different execution characteristics (communication load, optimizer overheads) that were not encountered at 3–8B scale. It might also be helpful to evaluate how the effect of CENTRIFUGE might change in MoE. \n2. Forward pass remains unoptimized: Centrifuge focuses on filtering tokens during backprop and does not skip any computation in the forward pass. This design is intentional to preserve model utility, but it means the maximum possible speedup is inherently limited. It would be valuable for the authors to clarify this limitation in the paper. Currently, the paper implies this in Section 2.1 and via citations, but a more explicit discussion of forward-pass cost might be helpful to set expectations for the reader.\n3. The system-level solution, involving dynamic graph rewriting in PyTorch, is quite complex. While the authors do an admirable job automating it, one might worry about maintainability. If underlying libraries or model architectures change, these techniques might need adjustment. A discussion on how robust the implementation is to different model architectures (beyond the ones tested) and framework versions would be useful. This isn’t a flaw in the core idea per se, but acknowledging the complexity and how it’s managed would improve transparency.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kLlMYBqLZc", "reviewer": "Reviewer_VSJV", "summary": "Token filtering aims to improve training efficiency by focusing on only the most important tokens. However, current methods only introduce marginal training speedup because they do not fully utilize the token sparsity. This paper introduces CENTRIFUGE, a novel system designed to unlock the full efficiency of token filtering in large language model training. At the algorithm level, CENTRIFUGE proposes filtering activations in the attention backward kernel to amplify the sparsity. At the system level, CENTRIFUGE transforms sparse GEMM into dimension-reduced dense GEMM to maximize the computional efficiency. Experiments show that CENTRIFUGE maintains a similar training loss with the traditional approach of only filtering loss, while reducing training time by 31.7%.", "questions": "None", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The paper provides a clear diagnosis of why token filtering has not delivered superior training efficiency. The proposed method is intutive and effectively overcomes current limitations.\n2.  Experiments on fine-tuning foundation models show that CENTRIFUGE preserves the benefits of token filtering and reduces end-to-end training time by 31.7%.", "weaknesses": "1. Experiments are only conducted on fine-tuning pre-trained foundation models for downstream tasks. However, since the main advantage of the proposed method lies in its training efficiency, it should be validated on computationally intensive pretraining tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "d89dJh3OxN", "reviewer": "Reviewer_5NLK", "summary": "This paper presents CENTRIFUGE, a system that combines algorithm and system-level optimizations to improve the efficiency of token filtering in LLM training. The authors identify two key limitations of existing token filtering approaches: (1) inadequate sparsity propagation beyond the output layer, and (2) mismatch between token filtering sparsity (30-50%) and what sparse GEMM libraries can efficiently handle (>95%). CENTRIFUGE addresses these through filtering activations in attention backward kernels and transforming sparse GEMM into dimension-reduced dense GEMM, achieving up to 49% backward time reduction.", "questions": "1. How does this interact with other optimizations like FlashAttention-2 or PagedAttention? In particular, how exactly does this integrate with FlashAttention's tiling strategy?\n\n2. What is the overhead of the offline preparation phase for new models?\n\n3. Can you provide theoretical analysis of the gradient approximation quality?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper clearly identifies why existing token filtering methods fail to deliver efficiency gains despite their utility benefits, with strong empirical evidence (Figure 2, Section 3).\n\n2. Testing across four models (1.1B-8B parameters) and multiple training scenarios (DP, TP, LoRA) demonstrates broad applicability. The utility evaluation on 9 tasks shows preserved benefits.\n\n3. The one-line integration for existing systems and consideration of real-world distributed training scenarios (TP, PP, MoE) enhance practical value.\n\n4. The paper is generally well-written with good use of figures to illustrate concepts.", "weaknesses": "1. Limited algorithmic novelty. The core contribution (filtering K, V activations while preserving Q) is relatively straightforward once the problem is identified. meanwhile, theoretical analysis of why this approach is optimal and exploration of alternative filtering strategies are missing. The connection to existing memory-efficient attention implementations could also be more rigorously analyzed\n\n2. Insufficient Analysis of Strawman Approach. Section 3.1 claims the strawman approach \"cannot work\" but provides limited evidence, In particular, figure 4(a) shows non-convergence but lacks quantitative analysis of gradient corruption, ablation studies on different filtering strategies, or detailed explanation of the \"interference\" mechanism.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)