-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathindex.html
More file actions
417 lines (404 loc) · 25.1 KB
/
Copy pathindex.html
File metadata and controls
417 lines (404 loc) · 25.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Sparrow</title>
<meta name="description" content="Project page for Sparrow: Sparse Rollout for Stable and Efficient Long-context RL of Large Language Models.">
<meta property="og:title" content="Sparrow">
<meta property="og:description" content="Sparse Rollout for Stable and Efficient Long-context RL of Large Language Models.">
<meta property="og:type" content="website">
<meta property="og:image" content="radom/headpic.png">
<meta name="twitter:card" content="summary_large_image">
<meta name="twitter:title" content="Sparrow">
<meta name="twitter:description" content="Sparse Rollout for Stable and Efficient Long-context RL of Large Language Models.">
<meta name="twitter:image" content="radom/headpic.png">
<link rel="icon" type="image/png" href="assets/figures/figure112.png">
<link rel="preconnect" href="https://fonts.googleapis.com">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;450;500;600;700;800&family=Newsreader:ital,opsz,wght@0,6..72,400;0,6..72,500;0,6..72,600;1,6..72,400&family=JetBrains+Mono:wght@400;500;600&display=swap" rel="stylesheet">
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/academicons/1.9.4/css/academicons.min.css">
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.2/css/all.min.css">
<link rel="stylesheet" href="styles.css">
<script defer src="script.js"></script>
</head>
<body>
<header class="hero">
<nav class="topbar" aria-label="Page navigation">
<a class="brand" href="#top" aria-label="Sparrow home">
<img src="assets/figures/figure112.png" alt="" aria-hidden="true">
<span>Sparrow</span>
</a>
<div class="navlinks">
<a href="#desiderata">Desiderata</a>
<a href="#insights">Insights</a>
<a href="#control-study-1">Control Study 1</a>
<a href="#control-study-2">Control Study 2</a>
<a href="#empirical-generalization">Empirical Generalization</a>
<a href="#distillsparse">DistillSparse</a>
</div>
</nav>
<section id="top" class="hero-inner">
<p class="eyebrow">Sparse Rollout for Stable and Efficient Long-context RL for Large Language Models</p>
<h2 class="hero-title">
<img src="assets/figures/figure112.png" alt="" aria-hidden="true">
<span>
Sparrow: <u>Spar</u>se <u>Roll</u>out for Stable and Efficient Long-context RL of Large Language Models
</span>
</h2>
<p class="authors">Yang Zhou, Ranajoy Sadhukhan, Zhaofeng Sun, Zhuoming Chen, Souvik Kundu, Saket Dingliwal, Sai Muralidhar Jayanthi, Aram Galstyan, Haizhong Zheng, Beidi Chen</p>
<p class="affiliation">Carnegie Mellon University, Independent Researcher, Intel, Amazon</p>
<div class="button-row" aria-label="Project links">
<a class="pill" href="https://arxiv.org/abs/2606.08446" aria-label="arXiv paper">
<span class="icon arxiv-icon">
<i class="ai ai-arxiv"></i>
</span>
arXiv
</a>
<a class="pill" href="https://github.com/Infini-AI-Lab/Sparrow" aria-label="GitHub repository">
<span class="icon">
<i class="fab fa-github"></i>
</span>
GitHub
</a>
<a class="pill" href="#citation">
<span class="icon">{ }</span>
Citation
</a>
</div>
<div class="hero-figure">
<div class="figure-card large hero-image-card hero-video-card">
<video autoplay muted loop playsinline controls aria-label="Sparse rollout video demonstration.">
<source src="assets/figures/2603.08640.mp4" type="video/mp4">
</video>
</div>
<div class="hero-stats" aria-label="Paper highlights">
<div>
<strong>0.86 ≥ τ</strong>
<span>Tail Per-token Mismatch Threshold τ Identified for Qwen3 Family thinking models</span>
</div>
<div>
<strong>2.2x, 2.4x, 2.0x</strong>
<span>Speedup Achieved using Sparsity Scheduling Qwen3 1.7B, 4B, 8B thinking models' Stable RL rollout</span>
</div>
<div>
<strong>14B</strong>
<span>threshold generalization validated to significantly larger model size</span>
</div>
</div>
</div>
<div id="overview" class="hero-overview">
<p class="lead">
Despite being <span class="text-highlight">powerful</span>, reinforcement learning with verifiable rewards (RLVR) triggers <span class="text-highlight">extremely long
COT</span>, thus highly computationally expensive. RLVR per-step training cost is dominated by
<span class="text-highlight">long-context generation in rollout</span>, making sparse attention a promising technique to accelerate vanilla dense rollout
generation. However, in practice, mastering the <span class="text-highlight">tradeoff</span> between sparse rollout RL stability and efficiency
is difficult - either sparsity too aggressive and training collapse or too lenient and insufficient speedup.
</p>
<p class="lead">
To study the optimal tradeoff, we first observe that most tokens generated by sparse rollout are perfectly
aligned with dense policy even under high sparsity. We then hypothesize that once we constrain the
<span class="text-highlight">tail distribution</span> of the sparse to dense actor-policy mismatch by a threshold, we are able to train with sparse
rollout stably. We validate that our hypothesis holds by introducing a <span class="text-highlight">dynamic sparsity scheduling</span>
method for keeping the tail distribution constant through generation and study how the threshold scales
with model size.
</p>
<p class="lead">
Surprisingly, across a range of model sizes in Qwen3 thinking family, we find that keeping
<span class="text-highlight">5-percentile mismatch threshold</span> above 0.86 generally works and use a cost-model analysis to
find the a sparsity scheduling for maximum speedup under mismatch threshold, thus achieving
<span class="text-highlight">2.2x, 2.4x, and 2.0x</span> in rollout when training Qwen3-1.7B, Qwen3-4B, Qwen3-8B. Empirically, we show the
identified threshold generalizes to <span class="text-highlight">much larger model size</span> (Qwen3-14B) and other RL domain (Coding)
and enables stable training. Additionally, we propose and show that <span class="text-highlight">DistillSparse</span>, a lightweight
LoRA-based distillation method that further actively aligns sparse rollouts with the dense policy and
enables higher speedup with more aggressive sparsity.
</p>
</div>
</section>
</header>
<main>
<section id="desiderata" class="section visual-band">
<div class="section-heading">
<span class="section-icon section-icon-image">
<img src="radom/Telescope.png" alt="" aria-hidden="true">
</span>
<h2>Previous Works Limitation and Desirata</h2>
</div>
<div class="desiderata-layout">
<div class="section-copy">
<p>
Sparse attention (especially dynamic token selection) has been widely studied. However, prior work falls short of building a concrete understanding of the optimal tradeoff between stability and efficiency in the RL setting.
</p>
<ul class="lead">
<li>
<span class="text-highlight">Sparse attention inference accuracy ↑ ≠ RL stability ↑.</span>
For downstream-task inference, sparse rollout RL instability is mainly due to per-token distribution mismatch
between the sparse actor and dense policy rather than insufficient rollout rewards.
</li>
<li>
<span class="text-highlight">Suboptimal convergence under severe actor-policy mismatch.</span>
Prior work (TIS, Jackpot) addresses actor-policy distribution mismatch, but often studies milder scenarios,
such as staleness, where actor-policy KL divergence is one order of magnitude smaller than in sparse rollout.
When directly applied, these techniques require clipping or masking significant training signals to maintain
stability, leading to poor training convergence.
</li>
<li>
<span class="text-highlight">Poor efficiency.</span>
Sparse rollout can be trivially recovered to achieve stable training by applying elementwise Top-k or using a
huge KV budget, but these approaches do not achieve large efficiency gains.
</li>
</ul>
<p>
Ideally, we desire RL with sparse rollout to:
</p>
<ul class="lead">
<li>
Enable <span class="text-highlight">stable</span> dense-policy training.
</li>
<li>
<span class="text-highlight">Match dense performance</span> regardless of model size and generation length.
</li>
<li>
Achieve strong <span class="text-highlight">efficiency</span> benefits over dense rollout.
</li>
</ul>
</div>
<figure class="figure-card side-figure">
<img src="assets/figures/figure102.png" alt="Sparse rollout stability and efficiency desiderata.">
<figcaption>Inference Accuracy using sparse attention (lower average reward than dense rollout) isn't the main problem to sparse rollout collapse. Here on top of an aggressive sparsity which triggers collapse, we sample 2N times and select the top-N rewards during sparse rollout, achieving higher average reward than dense (lower), and the training still crashes (higher).</figcaption>
</figure>
</div>
</section>
<section id="insights" class="section split full-width-section">
<div>
<div class="section-heading">
<span class="section-icon section-icon-image">
<img src="radom/Idea.png" alt="" aria-hidden="true">
</span>
<h2>Insights to Sparse to Dense Actor-Policy Mismatch</h2>
</div>
<div class="wide-grid single-figure">
<figure class="figure-card">
<img src="assets/figures/figure101.png" alt="Sparse rollout reward collapse under low KV budget.">
<figcaption>Arbitrary sparse rollout can miss the stable low-cost region.</figcaption>
</figure>
</div>
<p>
The key observation is that <span class="text-highlight">sparse rollout collapse is not driven by a uniform degradation across all tokens.</span>
Even under aggressive sparsity, most generated tokens remain nearly distribution-aligned with the dense policy.
The unstable signal instead appears in the small fraction of tokens where sparse and dense behavior diverge.
</p>
<p>
To see this, we collect tokens from sparse-attention generations with 20K prompt length and 2K generation length,
then measure each token's distribution mismatch against the dense model. The resulting distribution is highly
skewed: most tokens are close to perfectly aligned, even when the sparsity budget is aggressive enough to cause
RL collapse. (In fact, in the paper we show that the tail distribution can be well-modeled by a Beta Distribution)
<span class="text-highlight">The high skewness makes average mismatch a weak stability indicator.</span> For example, when training Qwen3-1.7B
with a 37K generation budget, a KV budget of 4096 remains stable while 2560 collapses, but their average
per-token sparse-dense L1 distances are still very close and very sensitive to measurement precision: 0.977 and 0.968.
</p>
<p>
We therefore evaluate sparse-dense mismatch with <span class="text-highlight">lower-tail statistics rather than the average</span>. In particular,
we use the lower 5-percentile to measure the worst-aligned tokens while ignoring the many "perfect" tokens that
are unlikely to cause RL training collapse. This leads to the following hypothesis.
</p>
<p class="hypothesis-box">
<strong>Hypothesis:</strong> if the tail distribution mismatch between the sparse actor and dense policy stays above a threshold
throughout rollout, sparse rollout for dense policy RL will be stable.
</p>
<p>
For our study, we focus on block-sparse attention as a study, but we believe that the principles discovered should be applicable to other types of context compression methods.
</p>
</div>
</section>
<section id="control-study-1" class="section visual-band">
<div class="section-heading">
<span class="section-icon section-icon-image">
<img src="radom/Observation.png" alt="" aria-hidden="true">
</span>
<h2>Control Study 1: How Mismatch Threshold Varies to Increasing Model Sizes</h2>
</div>
<div class="wide-grid">
<figure class="figure-card">
<img src="assets/figures/fig-008.png" alt="Qwen3-1.7B acceptance rate versus sequence length.">
</figure>
<figure class="figure-card">
<img src="assets/figures/fig-009.png" alt="Qwen3-4B acceptance rate versus sequence length.">
</figure>
<figure class="figure-card">
<img src="assets/figures/fig-010.png" alt="Qwen3-8B acceptance rate versus sequence length.">
</figure>
<figure class="figure-card">
<img src="assets/figures/fig-011.png" alt="Qwen3-1.7B AIME2026 training curves.">
</figure>
<figure class="figure-card">
<img src="assets/figures/fig-012.png" alt="Qwen3-4B AIME2026 training curves.">
</figure>
<figure class="figure-card">
<img src="assets/figures/fig-013.png" alt="Qwen3-8B AIME2025 training curves.">
</figure>
</div>
<p class="full-width-text">
There is one main challenge to conduct systematic study of sparse dense mismatch and training stability.
Under any fixed sparsity budget, the distribution mismatch always deteriorates as the generation length increases.
</p>
<p class="full-width-text">
We introduce the technique of <span class="text-highlight">sparsity scheduling</span>, which uses a more lenient sparsity budget as generation length increases
to keep the per-token mismatch approximately constant throughout generation, more details in the paper.
With sparsity scheduling and ability of hold the sparse dense Actor-Policy mismatch constant throughout the trajectory, we study the relationship between the target mismatch threshold and training stability.
Surprisingly, we make the following finding.
</p>
<p class="hypothesis-box">
<strong>Takeaway</strong>: across a range of model sizes in Qwen3 thinking family, we find that keeping
5-percentile mismatch threshold above 0.86 generally leads to stable RL training, supporting our hypothesis.
</p>
<p>
It supports that our hypothesis holds.
</p>
</section>
<section id="control-study-2" class="section">
<div class="section-heading">
<span class="section-icon section-icon-image">
<img src="radom/GPU.png" alt="" aria-hidden="true">
</span>
<h2>Control Study 2: Finding the Lowest Cost given the Mismatch Threshold</h2>
</div>
<div class="results-layout paired-figures">
<figure class="figure-card">
<img src="assets/figures/figure106.png" alt="Qwen3-14B sparse rollout and dense rollout AIME2026 results.">
<figcaption>From our cost model study, we identified that for block-sparse attention (block-size 16 and above), smaller page size consistently ahead of the tradeoff of distribution alignment to dense and cost.</figcaption>
</figure>
<figure class="figure-card">
<img src="assets/figures/figure103.png" alt="DistillSparse training improves AIME2025 accuracy.">
<figcaption>Empirically, we achieve speedup for all different sized models.</figcaption>
</figure>
<p class="full-width-text">
With the mismatch threshold identified, we then look for the lowest cost achievable while meeting the mismatch threshold.
To make sure our analysis is general and can transfer to different hardware, we use a cost model for our analysis.
</p>
<details class="formula cost-model full-width-text">
<summary>Full explanation of the cost model</summary>
<p>
We follow Sadhukhan et al. (2025)
and model rollout cost from model size and hardware memory bandwidth. For repeated sampling <code>N</code> times,
let <code>P</code> be the number of model parameters, <code>L<sub>in</sub></code> the input prompt length,
<code>L<sub>out</sub></code> the output length, <code>D</code> the Key/Value dimension, <code>r</code> the GQA ratio,
and <code>I</code> one over the GPU SRAM memory bandwidth.
</p>
<div class="cost-equation">
<p class="formula-label">Dense attention</p>
<pre><code>C<sub>comp</sub> = 2 * P * N * L<sub>out</sub> + r * (2 * L<sub>in</sub> + L<sub>out</sub>) * L<sub>out</sub> * N * D
C<sub>mem</sub> = 2 * L<sub>in</sub> * L<sub>out</sub> * D + N * L<sub>out</sub><sup>2</sup> * D
C<sub>dense</sub> = C<sub>comp</sub> + I * C<sub>mem</sub></code></pre>
</div>
<p>
For block-sparse attention, we assume Top-k kernels incur only minimal overhead for page size
<code>≥ 16</code>. With KV budget <code>B</code> and page size <code>pagesize</code>, the sparse cost is:
</p>
<div class="cost-equation">
<p class="formula-label">Block-sparse attention</p>
<pre><code>C<sub>sparse,no scoring</sub> = 2 * N * P * L<sub>out</sub> + 2 * r * N * D * B * L<sub>out</sub> + 2 * I * N * D * B * L<sub>out</sub>
C<sub>scoring</sub> = 2 * N * L<sub>in</sub> * D * L<sub>out</sub> + (r * N * D * L<sub>out</sub><sup>2</sup>) / (2 * pagesize)
+ 2 * I * L<sub>in</sub> * D * L<sub>out</sub> + (I * N * D * L<sub>out</sub><sup>2</sup>) / (2 * pagesize)
C<sub>sparse</sub> = C<sub>sparse,no scoring</sub> + C<sub>scoring</sub></code></pre>
</div>
</details>
<p class="full-width-text">
Using this cost model, we find that for <span class="text-highlight">block-sparse attention with page size 16 and above</span>,
smaller page sizes consistently dominate larger page sizes in the tradeoff between
<span class="text-highlight">dense-policy distribution alignment</span> and cost. In the paper, we provide full details
for other generation-length regimes, and the same conclusion holds across them. Although the cost model does
not explicitly account for <span class="text-highlight">Top-k kernel overhead</span>, our page-size measurements show that when
the Top-k kernels are well implemented, as in Vortex, sparse decoding cost is not highly sensitive to page size.
We therefore use <span class="text-highlight">page size 16</span> for sparsity scheduling, since it gives the strongest
cost-alignment tradeoff while remaining practical for efficient sparse rollout.
</p>
</div>
</section>
<section id="empirical-generalization" class="section visual-band">
<div class="section-heading">
<span class="section-icon section-icon-image">
<img src="radom/iss.webp" alt="" aria-hidden="true">
</span>
<h2>Generalization to Larger Model and Other RL domains (Coding RL)</h2>
</div>
<div class="wide-grid single-figure">
<figure class="figure-card">
<img src="assets/figures/figure107.png" alt="Sparse rollout reward collapse under low KV budget.">
<figcaption>Arbitrary sparse rollout can miss the stable low-cost region.</figcaption>
</figure>
</div>
<p class="full-width-text">
We further test our hypothesis in settings where exhaustive grid search is impractical due to limited compute.
By holding the <span class="text-highlight">tail sparse-dense mismatch threshold</span> at 0.86, we are able to
stably train the <span class="text-highlight">Qwen3-14B</span> model for a full epoch on Polaris, reaching performance on par
with dense rollout. This setting is especially challenging because dense training would normally require roughly
<span class="text-highlight">8 days (190 hours)</span> on a 4-node, 32-GPU H200 cluster.
</p>
<p class="full-width-text">
Beyond math reasoning RL, we
also verify that the same threshold transfers to <span class="text-highlight">coding RL</span>. Specifically, we train the
Qwen3-1.7B thinking model for a full epoch on TACO while maintaining the 0.86 threshold, and observe that both
average reward and downstream performance remain on par with dense rollout. More details presented in paper.
</p>
</section>
<section id="distillsparse" class="section">
<div class="section-heading">
<span class="section-icon section-icon-image">
<img src="radom/cosmonautllama.png" alt="" aria-hidden="true">
</span>
<h2>DistillSparse: A Technique Pushes for More Aggressive Sparsity and Higher Speedup</h2>
</div>
<div class="results-layout paired-figures balanced-figures">
<figure class="figure-card">
<img src="assets/figures/figure110.png" alt="DistillSparse brings us higher speedup while reaching the same sparse-dense mismatch threshold.">
<figcaption>DistillSparse brings us higher speedup while reaching the same sparse-dense mismatch threshold.</figcaption>
</figure>
<figure class="figure-card">
<img src="assets/figures/figure109.png" alt="DistillSparse training improves AIME2025 accuracy.">
<figcaption>LoRA generally improves the sparse attention to better aligned with the dense bringing original 0.80 sparsity setting to get mismatch level near 0.86, improving speedup from 2.2x to 2.5x.</figcaption>
</figure>
</div>
<p class="full-width-text">
After identifying a stable sparse-dense mismatch threshold, we ask whether sparse attention can be pushed to
deliver even higher rollout speedup. A useful observation is that sparse-rollout dense-policy training already
contains the ingredients needed for <span class="text-highlight">on-policy distillation</span>: trajectories are generated with
sparse attention, while dense log probabilities are computed by the dense policy for training. This naturally
provides supervision for making the sparse actor closer to the dense policy.
However, main challenge is to perform this alignment without contaminating the dense policy and without adding substantial training overhead.
</p>
<p class="full-width-text">
We propose a <span class="text-highlight">LoRA-based sparse distillation</span> design. DistillSparse adds an auxiliary sparse
distillation objective that actively aligns sparse rollouts with the dense policy while updating only the LoRA
parameters. Starting from the original <span class="text-highlight">0.86 mismatch-threshold</span> setting, we find that after
training on 20K examples, the learned LoRA generally improves sparse attention enough to make a more aggressive
<span class="text-highlight">0.80 sparsity setting</span> approach the original 0.86 mismatch level, enabling higher speedup.
Empirically, this LoRA-only procedure introduces minimal overhead while improving rollout efficiency; full
details are provided in the paper.
</p>
</section>
<section id="citation" class="section citation-section">
<div class="section-heading">
<span class="section-icon">{ }</span>
<h2>Citation</h2>
</div>
<p>If you find our study helpful to your understanding, consider citing us.</p>
<div class="codebox">
<button class="copy-button" type="button" data-copy-target="bibtex">Copy</button>
<pre id="bibtex"><code>@misc{zhou2026sparrowsparserolloutstable,
title={Sparrow: Sparse Rollout for Stable and Efficient Long-context RL of Large Language Models},
author={Yang Zhou and Ranajoy Sadhukhan and Zhaofeng Sun and Zhuoming Chen and Souvik Kundu and Saket Dingliwal and Sai Muralidhar Jayanthi and Aram Galstyan and Haizhong Zheng and Beidi Chen},
year={2026},
eprint={2606.08446},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2606.08446},
}</code></pre>
</div>
</section>
</main>
<footer class="footer">
<p>Website template inspired by the Infini AI Lab project-page style.</p>
</footer>
</body>
</html>