Plateau Dropping: a simple SFT rule for better agentic post-training

1. Introduction

Marin is a project aimed at building performant open models (open weights, open data, open recipes, etc.). They are currently pretraining a 535B-A23B MoE model (view the progress here!). Prior to this, they finished pretraining a smaller 67B-A2B model (codenamed “Snowball”). This model then went through an SFT stage focused on general capabilities. Our goal was to spend a few weeks turning this checkpoint into an “agentic expert”. We were specifically interested in improving its performance on long-horizon SWE and terminal tasks. We started out with RL but could not improve the model’s performance by more than a few points (within noise for our benchmarks). Our hypothesis was that the base model itself was not performant enough at agentic tasks, and could greatly benefit from an SFT warmup stage. This pipeline (general SFT → domain-specialized SFT → RL) is not specific to our models: many frontier post-training recipes train domain-specialized experts with SFT and RL and then combine them (GLM-4.5, DeepSeek-V3.2, MiMo-V2-Flash). The resulting expert can also serve as a teacher in multi-teacher on-policy distillation (MOPD) setups (MiMo-V2-Flash). Since we eventually want to run RL on this warmed-up model (likely with some GRPO variant), we were more interested in preserving and improving the model coverage during this specialized SFT phase, rather than solely benchmaxxing pass@1. Indeed, a recent work (TailSFT) shows that starting GRPO from models with better initial coverage leads to better post-RL checkpoints (measured by pass@1), even if the initial pass@1 is not optimal. TailSFT argues that vanilla SFT keeps raising the likelihood of already “fit” / “learned” sequences, and this in turn hurts coverage. They address this by dropping, from each batch, the sequences whose loss has fallen the most relative to the base model. On OLMo-3 7B, they improve pass@16 in 15 of 18 math and code settings. However, when we try this method on our long-horizon agentic workload, we do not observe any benefits (Section 4.3).

Inspired by the ideas presented in TailSFT, and by the fact that GRPO doesn’t learn anything from all pass / all fail groups, we propose a simple SFT approach called Plateau Dropping, where we track a token’s loss history during SFT and drop it if its loss hasn’t changed enough in the last few times it was seen. If its loss later changes substantially again, it is trained on again. The key intuition is to stop spending updates on tokens whose loss has stopped changing: either the model has already learned them, or further updates on them are not helping. In practice, most of the tokens this rule drops are “learned” ones, which already have a very low loss. Dropping only the small set of “stuck” tokens, whose loss has stopped changing at a high value, recovers all of the pass@1 gain but only about half of the pass@8 gain; the rest of the coverage gain appears only when both kinds are dropped together (Section 4.5). We find that this simple rule works pretty well in agentic settings, yielding a +4.5pp gain at pass@8 on SWE-bench Verified, and smaller gains on terminal benchmarks, with no significant drop in pass@1.

Plateau Dropping is closest to methods that select training data by how it behaves. Rho-1 (Lin et al., 2024) targets (continual) pre-training and trains only on tokens whose loss is high relative to a reference model. In contrast, Plateau Dropping needs no reference model, and its selection criterion follows the current model’s own loss history. Methods like TailSFT, Selective Backprop (Jiang et al., 2019), and loss truncation (Kang and Hashimoto, 2020) drop or prioritise whole examples, while Plateau Dropping selects individual tokens instead. Focal loss (Lin et al., 2017) down-weights well-classified examples, based on a snapshot of the model’s confidence; Plateau Dropping instead looks at how a token’s loss changes over time.

2. Methodology

We use \(\mathcal{D} = \{y_n\}_{n=1}^N\) to denote our dataset of agentic trajectories. A trajectory \(y \in \mathcal{D}\) is a sequence of tokens \((y_i)_{i=1}^L\). We use the standard NLL loss for supervised fine-tuning: \(\ell_i(\theta) = -\log p_\theta (y_i \mid y_{<i})\) for token \(i\) of a single trajectory \(y\). Since these trajectories are agentic in nature, they also contain tool outputs. We zero out the loss on these tokens, as well as on the prompt. Thus, the final SFT loss we use is:

where \(w_i=1\) only for the agent’s tokens, and \(=0\) for the prompt and the tool responses.

Detecting plateaus. Let \(\ell_i^{(j)}\) denote token \(i\)’s loss at its \(j^{\mathrm{th}}\) visit. For \(j \ge 2\), let \(\delta_i(j)\) be the relative change in this loss since the previous visit:

We say token \(i\) has plateaued at visit \(j \ge 3\) if \(\delta_i(j) < \tau \text{ and } \delta_i(j - 1) < \tau\). This can be written as \(\alpha_i(j) = \mathbf{1}\{\delta_i(j) < \tau \,\wedge\, \delta_i(j - 1) < \tau\}\). This definition can be extended to consider the past \(k\) visits as well, but we only consider \(k=2\) here. We drop the plateaued tokens from the SFT loss:

where \(\kappa_i = 1 - \alpha_i(v_i)\) and \(v_i\) is token \(i\)’s current visit.

In our experiments, we use \(\tau = 0.1\) and \(\varepsilon = 0.05\). \(\mathcal{L}_{\mathrm{PD}}\) reduces to \(\mathcal{L}_{\mathrm{SFT}}\) for \(\tau=0\). Note that the history is always updated for all tokens, even the ones that get dropped, so that if a dropped token’s loss moves again, it is trained again.

Figure 1 shows how the fraction of dropped tokens changes as the training proceeds.

Figure 1. Plateau Dropping, 10 epochs on the SWE data: the share of trained tokens dropped as learned, dropped as stuck, and still being learned at the training steps shown. Dropped as learned: the loss has stopped moving, below −log 0.9; dropped as stuck: it stopped moving at a higher loss. Bars: the logged per-batch shares, averaged over the 5 steps ending at each step shown. Figure 2 (Appendix C) colours every token of three randomly chosen assistant turns by its final state.

3. Experimental Setup

• Model. We use Marin’s 67B MoE model with about 2B parameters active per token: open-athena/Grug-67B-A2B-Datakit-SFT-262K-2026.09.21. This is an earlier version of the Snowball checkpoint released recently. Note that this model has already gone through one general SFT phase after pre-training.
• Terminal data. 6,301 total successful trajectories (45M tokens, 1,610 tasks): 5,324 from the public allenai/tmax-sft dataset and 977 of the base model’s own successes on synthetic terminal tasks from Endless Terminals and Recursive Synthesis of Tasks (622 and 355 trajectories, from 331 and 219 tasks). Context 32k.
• SWE data. 9,229 successful trajectories (7,005 tasks) from the public nvidia/Open-SWE-Traces dataset, restricted to one shell command per step to match our agent (319M tokens). Context 64k.
• Training. The learning rates are given in Section 4. We use a cosine schedule over the run that decays to 10% of the peak learning rate, with 3% linear warmup, AdamH, no weight decay, bf16 compute with fp32 master weights and optimizer states, expert capacity factor 2.0 (no dropped token-to-expert assignments), router biases frozen. Vanilla SFT and Plateau Dropping runs of the same length share the data order. We train with Marin and its JAX training library Levanter.
• Harness. We use the mini-swe-agent harness everywhere: all training trajectories are in its format, and all evaluations run in it.

4. Results

We focus only on agentic capabilities in this work, evaluating SWE and terminal tasks. For SWE, we evaluate on a fixed random subset of 100 tasks from SWE-bench Verified. For terminal tasks, we use 88 of the 89 tasks in Terminal-Bench 2.0 and 97 of the 100 tasks in OpenThoughts-TBLite; infrastructure errors kept us from building the sandboxes for the remaining tasks. Terminal tasks are run with 3x the default time limits, and when an agent hits the time limit, its final state is scored. Sandbox and API errors are counted as failures. We report pass@1 and pass@8 for all models. We sample at temperature 1 with no top-p / top-k / penalties.

Note: the SWE-bench Verified tasks served by Harbor’s registry (swebench-verified@1.0) run git clean -fd before the tests, which deletes every file the agent created that is not tracked by git, so a correct fix that adds a new file is marked as a failure. We instead use task scripts that delete only the untracked files touched by the task’s test patch, which is what Harbor’s SWE-bench adapter now does as well (PR #1094).

4.1 SWE

We use a learning rate of 3e-4 for SFT and 1e-4 for PD. We find these to be optimal for both settings (searched over 1e-4, 3e-4, 5e-4). We run both SFT and PD for 3 and 10 epochs and find that 10 epochs works much better for both. Table 1 thus compares pass@1 and pass@8 for SFT and PD at 10 epochs, each with its best learning rate. All results are averaged over 2 training runs with different data orders. PD is level with SFT in pass@1 and outperforms it by 4.5 points in pass@8.

Note that it is uncommon in prior work to train for as long as 10 epochs (agentic SFT recipes train for 2 to 7: 2 in R2E-Gym, Nemotron-Terminal and TMax, up to 3 in SWE-smith, up to 5 in SWE-Gym and 7 in OpenThoughts-Agent), but we find that it can work best, especially when using SFT to specialize in a domain (here, agentic tasks) with limited data.

Table 1. SWE-bench Verified (random 100 tasks), 10 epochs (760 steps), pass@1 and pass@8 in %, 8 seeds, averaged over 2 training runs with different data orders.

4.2 Terminal

We use a learning rate of 1e-4 for both SFT and PD. For terminal data, unlike SWE data, this outperforms both 3e-4 and 5e-5. We run both SFT and PD for 3 and 10 epochs and find that 10 epochs works much better for both. Table 2 thus compares pass@1 and pass@8 for SFT and PD at 10 epochs, with a learning rate of 1e-4. We observe that the gains obtained by PD over vanilla SFT are much more modest for terminal tasks. PD outperforms SFT on TB-Lite by 2.2 points in pass@1 and 3.1 points in pass@8. On TB2, PD gives no consistent gains over SFT in either metric.

Table 2. TB-Lite (all 97 tasks) and Terminal-Bench 2.0 (all 88 tasks), 10 epochs (220 steps), pass@1 and pass@8 in %, 8 seeds.

4.3 Baselines: TailSFT

We try out TailSFT and its token-level variant on SWE data and find that neither beats vanilla SFT at pass@8 in any setting. We train all methods on the same data for 3 epochs with a learning rate of 1e-4. We try out TailSFT with 50% of the trajectories dropped (ramped over the run) and 25% dropped (fixed). These are the settings recommended by the authors for math and code data respectively. We also try a token-level variant of TailSFT which drops tokens instead of entire trajectories. Table 3 shows the results: every variant we tried lags behind vanilla SFT at pass@8.

TailSFT proposes a diagnostic called the coverage ratio (Definition 4.1 in their paper) to predict whether it will help in a given setting: the coverage that vanilla SFT loses relative to the base model divided by the coverage it gains. For vanilla SFT after 3 epochs, this ratio is above 1 (3.1 at k=8), the regime in which the paper reports gains in 10 of 11 settings, but no TailSFT variant we tried is ahead of vanilla SFT at pass@8. Appendix D gives the details of the calculation.

Table 3. SWE-bench Verified (random 100 tasks), 3 epochs (228 steps), learning rate 1e-4, pass@1 and pass@8 in %, 8 seeds.

4.4 Learnable tasks

Table 4. Number of tasks never solved / solved by some of the 8 seeds / always solved, at 10 epochs, with the models of Tables 1 and 2 (for SWE-bench Verified, the first of the two data orders).

Group-based RL methods like GRPO learn only from tasks the model solves some of the time. On SWE-bench Verified, vanilla SFT solves more tasks reliably (26 vs 21 always solved), which gives it its higher pass@1, while Plateau Dropping keeps more tasks within reach (57 vs 48 in the “learnable band”); on the terminal benchmarks however, the two are about the same. Note that even vanilla SFT gains overall coverage (pass@8) over the base model in every setting we tried. The “sharpening” effect is relative here. We believe this is expected with relatively weak base models like the one we have used in this study. We are continuing to test more model families and will update results soon!

4.5 What is Plateau Dropping helping with?

Plateau Dropping drops two kinds of tokens: “learned” ones, which make up most of what it drops, and a small set of “stuck” ones. In our convention, learned tokens are those whose loss plateaus at a value \(< -\log 0.9\), while “stuck” tokens are all the other plateaued tokens, whose loss is still higher. To see which kind matters more, we drop only one of them, at the same learning rate (Table 5). We find that dropping only the learned tokens does about as well as vanilla SFT, while dropping only the stuck tokens matches Plateau Dropping at pass@1 and recovers about half of its pass@8 gain over vanilla SFT.

Table 5. SWE-bench Verified (random 100 tasks), 10 epochs (760 steps), learning rate 1e-4, pass@1 and pass@8 in %, 8 seeds; one training run each, and the two “only” runs use a different data order from the other two; “Dropped” is the share of trained tokens dropped in each batch at the end of training.

To see how much each method sharpens the model, we measure the entropy of its next-token distribution (Table 6) at three sets of positions:

• stuck tokens: tokens of the SWE training data that Plateau Dropping was dropping as stuck at the end of training, i.e. whose loss had stopped changing while still high;
• other training tokens: for each stuck token, a token from the same training trajectory where possible, with the same next token, in the same kind of span (reasoning or action), that Plateau Dropping was not dropping at the end of training;
• uncertain rollout tokens: reasoning tokens from successful SWE-bench Verified rollouts of the Table 1 models where the base model’s most likely next token has probability below 0.5; every model is scored on the same positions.

At its best learning rate (3e-4), vanilla SFT cuts the base model’s entropy by 58-73% across these sets, while Plateau Dropping cuts it by 29-41%, at the same pass@1 and a higher pass@8. Being less “sharp” is not enough on its own, though: vanilla SFT at learning rate 1e-4 ends up slightly “softer” than Plateau Dropping, but scores 6.6 points lower at pass@1 and 7.5 points lower at pass@8.

Table 6. Mean next-token entropy (nats, over the top 100 tokens) after 10 epochs on the SWE data, and SWE-bench Verified pass@1 and pass@8 in %. Vanilla SFT at 3e-4 and Plateau Dropping are averaged over the 2 data orders; vanilla SFT at 1e-4 is a single run.

An example. On django__django-11477, the base model and vanilla SFT never succeed (0 of 8 and 0 of 16), and Plateau Dropping succeeds twice in 16 attempts (across 2 data orders). In every attempt of both models, the reasoning ties the failure to None values in the resolved kwargs (11 of 16 Plateau Dropping and 10 of 16 vanilla SFT attempts name match.groupdict() explicitly). They differ in where they fix it. Every vanilla SFT attempt that edits the source (9 of 16) patches translate_url in django/urls/base.py, after setting the alternative aside:

Let me also check if maybe the fix should be more general — e.g., in resolve / ResolverMatch — but no, the issue is clearly in translate_url’s use of reverse.

The hidden tests also check resolve() itself, so this fix fails. Plateau Dropping takes the same route in 7 of its attempts, which fail too, but its two successes fix the resolver:

Actually, let me reconsider one more time. I keep going back and forth, so let me make a firm decision based on the strongest reasoning: […] Root cause: match.groupdict() includes None for non-participating groups. This is a resolver-level concern. Fix: Exclude None group values from kwargs in the resolver’s match() methods.

Note: both successes come from one of the two data orders, so this is an illustration rather than a strong statistical result.

Our reading of these measurements is that Plateau Dropping reaches vanilla SFT’s accuracy at a lower learning rate, and that this checkpoint is substantially less concentrated than vanilla SFT at its tuned learning rate. That leaves alternative reasoning paths reachable when sampling, and some of them produce additional solves as illustrated above.

5. Next steps

More model families. This work was done in collaboration with folks at Marin, with the aim of creating an agentic expert (a domain-specialized model that can also later be used in MOPD setups). This is why all experiments in this work have been performed using Marin’s 67B-A2B MoE model. We’re currently running more experiments with other popular open-source models to test the generalizability of this approach.

RL. The entire premise of this work has been that starting RL from a model that has more coverage can lead to better post-RL performance. TailSFT indeed demonstrates this for non-agentic tasks. We are now starting RL experiments that compare the post-RL performance of vanilla SFT and Plateau Dropping.

Data scale. The dataset sizes we use for SFT in this work are pretty small by frontier standards and by non-agentic SFT standards (although they are typical for academic agentic SFT; see Appendix B). In particular, our terminal data has ~7x fewer tokens than our SWE SFT data. We’re currently in the process of scaling the number of trajectories and tasks to see how well our approach works at larger token budgets.

Evaluation. We limit our agentic evals to SWE and terminal tasks. However, there are other practically relevant “agentic” domains like web search (BrowseComp), tool use (τ-bench) and computer use (OSWorld), which we are yet to evaluate on.

References

1. Malladi, S. et al. (2026). TailSFT: Filtered Fine-Tuning Improves Post-Training Performance. arXiv:2608.25756.
2. Dang, X., Baek, C., Wen, K., Kolter, Z. and Raghunathan, A. (2025). Weight Ensembling Improves Reasoning in Language Models. arXiv:2504.10478.
3. Chen, F., Raventos, A., Cheng, N., Ganguli, S. and Druckmann, S. (2025). Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning. arXiv:2502.07154.
4. GLM-4.5 Team (2025). GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Models. arXiv:2508.06471.
5. DeepSeek-AI (2025). DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models. arXiv:2512.02556.
6. Xiaomi LLM-Core Team (2026). MiMo-V2-Flash Technical Report. arXiv:2601.02780.
7. Lin, Z. et al. (2024). Rho-1: Not All Tokens Are What You Need. arXiv:2404.07965.
8. Jiang, A. H. et al. (2019). Accelerating Deep Learning by Focusing on the Biggest Losers. arXiv:1910.00762.
9. Kang, D. and Hashimoto, T. (2020). Improved Natural Language Generation via Loss Truncation. ACL 2020. arXiv:2004.14589.
10. Lin, T.-Y. et al. (2017). Focal Loss for Dense Object Detection. arXiv:1708.02002.
11. Pan, J. et al. (2025). Training Software Engineering Agents and Verifiers with SWE-Gym. arXiv:2412.21139.
12. Jain, N. et al. (2025). R2E-Gym: Procedural Environments and Hybrid Verifiers for Scaling Open-Weights SWE Agents. arXiv:2504.07164.
13. Yang, J. et al. (2025). SWE-smith: Scaling Data for Software Engineering Agents. arXiv:2504.21798.
14. Wang, H. et al. (2025). SWE-Dev: Building Software Engineering Agents with Training and Inference Scaling. arXiv:2506.07636.
15. Nebius (2025). SWE-rebench OpenHands trajectories. Hugging Face.
16. NVIDIA (2026). Nemotron-Terminal-Corpus. Hugging Face; arXiv:2602.21193.
17. OpenThoughts (2025). OpenThoughts-Agent-v1-SFT. Hugging Face.
18. Ivison, H. et al. (2026). Tmax: A Simple Recipe for Terminal Agents. arXiv:2606.23321.
19. Gandhi, K., Garg, S., Goodman, N. and Papailiopoulos, D. (2026). Endless Terminals: Scaling RL Environments for Terminal Agents. arXiv:2601.16443.
20. Li, Z. et al. (2026). Recursive Synthesis for Long-Horizon Terminal Tasks. arXiv:2608.05466.
21. Wei, J. et al. (2025). BrowseComp: A Simple Yet Challenging Benchmark for Browsing Agents. arXiv:2504.12516.
22. Yao, S. et al. (2024). τ-bench: A Benchmark for Tool-Agent-User Interaction in Real-World Domains. arXiv:2406.12045.
23. Xie, T. et al. (2024). OSWorld: Benchmarking Multimodal Agents for Open-Ended Tasks in Real Computer Environments. arXiv:2404.07972.

Acknowledgements

I thank Bespoke Labs for supporting my coding agent plan, Alex Dimakis for his advice, and the Jülich Supercomputing Centre (JSC) and the Texas Advanced Computing Center (TACC) for providing compute for this work.

Appendix A: training/infra footguns

We ran into several issues while training.

1. The master weights were bf16. The script that imported the base checkpoint in Levanter wrote the weights in bf16, and the trainer kept whatever precision it loaded, so the optimizer was initially applying its fp32 updates to bf16 weights and rounded the result back to bf16 at every step! A bf16 number has an 8-bit mantissa, so any update smaller than about 0.4% of a weight’s magnitude rounds to zero, which is most Adam updates at a learning rate of 1e-4. Even after 1520 steps, 92-95% of the weights were bit-identical to the base model because of this; and the training loss stopped at 0.55. The fix was to just keep fp32 master weights (the standard mixed-precision recipe: bf16 compute, fp32 weights and optimizer state). In the fixed runs 65% of the weights changed after just 228 steps, 83% after 760 and 89% after 1520. Every one of our initial experiments was trained this way: vanilla SFT, Plateau Dropping, TailSFT and its variants, the ablations, etc.

2. Dropping token-to-expert assignments. We were initially training the model with an expert capacity factor of 1.0, so whenever a batch routed more tokens to an expert than its even share, the excess assignments were just dropped. On this data 14-19% of assignments were dropped, which cost 0.067 to 0.082 nats of loss. With the capacity factor raised to 2.0 the overflow is 0.16% at the start of a run and 0 by the end, and the trainer’s loss matches an independent inference implementation to 0.0002.

3. Epochs are overstated. We first counted passes from token totals and ignored row padding, so what we called “10 epochs” on the terminal data are actually 8.6 “real” passes and “10 epochs” on the SWE data are 7.4. This post keeps the nominal names: our packer fills fixed-length rows (64k tokens on the SWE data, 32k on the terminal data) with whole trajectories, which leaves about 26% of each SWE row and 16% of each terminal row as padding, and the data loader shuffles the packed rows. One real pass over the SWE data is 103 steps, so 228 and 760 steps (“3 and 10 epochs”) are 2.2 and 7.4 passes; one pass over the terminal data is about 26 steps, so 66 and 220 steps are 2.6 and 8.6 passes. The plateau test counts visits, not epochs, so is unaffected by this.

4. Gradient clipping was silently penalising Plateau Dropping. The trainer clips the global gradient norm at 1.0 (the usual setting) and the loss Plateau Dropping optimises is a mean over the kept tokens only. Once it has dropped half the tokens its gradient norm is 2-3x vanilla SFT’s, so the clip, which touches 3-5% of vanilla SFT’s steps, hit 50% of the steps of the 10-epoch Plateau Dropping run on the SWE data and 82-100% of the steps of the terminal runs. Nothing in the training logs flagged it: the loss curves of the clipped and unclipped runs are indistinguishable, and only the evaluations differ, by 2-6 points at 10 epochs. The TailSFT paper also clips at 1.0 with the loss averaged over the tokens of the kept sequences (Malladi et al., 2026, Appendix D.1), but in our sequence-level TailSFT runs the clip engaged no more often than for vanilla SFT (see below), so we do not expect its results to be affected in the same way. At 10 epochs on the SWE data the two settings compare as follows (8 seeds, all 100 tasks; the loss curves of the two Plateau Dropping runs are indistinguishable, 0.347 against 0.346 over all tokens at the end):

Clip 4.0 beats clip 1.0 by 2.1 points at pass@1 and 6.0 points at pass@8, paired over tasks. The TailSFT runs were clipped no more than vanilla SFT was and stand as they are; the drop-share 0.25 and per-token TailSFT runs were trained at clip 4.0. The 10-epoch Plateau Dropping runs were redone at clip 4.0, and only those appear in Tables 1 and 2; the 3-epoch Plateau Dropping run in Table 3 never reached the clip at 1.0.

Appendix B: dataset sizes in agentic SFT

We compare our dataset sizes to a few publicly available datasets. In trajectory count, they are close to the sets several SWE agents were fine-tuned on (SWE-smith, R2E-Gym, SWE-Dev) and to the terminal set we draw from, and one to two orders of magnitude below the largest public releases (Open-SWE-Traces, Nemotron-Terminal-Corpus).

Table 7. Size of public agentic SFT datasets, as stated in each paper or dataset card (“–”: not reported).

Appendix C: annotated trajectories

We annotate some example trajectories using the still learning / dropped as learned / dropped as stuck categories of Figure 1.

Excerpt 1: 40% dropped: learned, 5% dropped: stuck, 55% still learning

<hdr>tool</hdr>
<tool_response name="bash">(no output)

(exit_code=0)</tool_response><eot>
<hdr>assistant</hdr>
<think>Now I can see the `registerKeepAliveHook` function. Let me understand the issue.

In the `registerKeepAliveHook` function, when a hook is registered with `onActivated` or `onDeactivated`, it creates a `wrappedHook` (with `__wdc` = with deactivation check) that:
1. Checks if the target instance is in a deactivated branch
2. If not, calls `hook()` directly

Then it calls `injectHook(type, wrappedHook, target)`.

The `injectHook` function wraps the hook with error handling (`__weh`), so `callWithAsyncErrorHandling` is used. So the flow is:

1. `injectHook(type, wrappedHook, target)` creates `wrappedHook.__weh` which calls `callWithAsyncErrorHandling(wrappedHook, target, type, args)`.
2. When the hook fires, `wrappedHook.__weh` is called (via the array of hooks in `target[type]`).
3. `callWithAsyncErrorHandling(wrappedHook, ...)` calls `wrappedHook()`, which checks the deactivation branch, then calls the original `hook()`.

Now, the issue is about async functions. Let me look at the `callWithAsyncErrorHandling` function in `errorHandling.ts`.</think>Now let me look at the error handling file which is central to this issue.<tool_call>
{"name": "bash", "arguments": {"command": "cat /testbed/packages/runtime-core/src/errorHandling.ts"}}
</tool_call><eot>

Excerpt 2: 69% dropped: learned, 3% dropped: stuck, 28% still learning

JSON dict: '{"a": true, "b": null}'
JSON list: '["1", true, null]'
Plain str: 'my param'
Int infer: 1234
Disabled: '{"a": true}'


(exit_code=0)</tool_response><eot>
<hdr>assistant</hdr>
<think>Confirmed. Currently:
- Python dict works: `{'a': True, 'b': None}`
- JSON dict fails: returns `'{"a": true, "b": null}'` (string)
- JSON list fails: returns `'["1", true, null]'` (string)

Alright, let's implement the fix. Add `import json` and modify `_infer_param_type` to fall back to `json.loads`.

First let's check the import section.</think>Confirmed the issue. Now let me implement the fix. I'll add the `json` import and modify `_infer_param_type` to fall back to `json.loads` when `literal_eval` fails.<tool_call>
{"name": "bash", "arguments": {"command": "cd /testbed && sed -n '19,26p' toolium/utils/dataset.py"}}
</tool_call><eot>

Excerpt 3: 42% dropped: learned, 7% dropped: stuck, 51% still learning

<hdr>tool</hdr>
<tool_response name="bash">(no output)

(exit_code=0)</tool_response><eot>
<hdr>assistant</hdr>
<think>Interesting — the one-argument form of `where` in `ops.py` already implements a non-zero-like behavior using `jnp.where(condition.array, size=new_axis.size, fill_value=fill_value)`. This is a great pattern to follow for `nonzero`.

Let's take a look at the entire `unique` function and other related parts to understand the pattern being used. I'll check the elided middle portion of ops.py.</think>The `where` function already shows the pattern for `jnp.where(condition, size=..., fill_value=...)`. Let me see the elided middle portion of `ops.py` to fully understand the `unique` function and its patterns.<tool_call>
{"name": "bash", "arguments": {"command": "cd /testbed && sed -n '140,330p' src/haliax/ops.py"}}
</tool_call><eot>

Figure 2. Three assistant turns from the SWE training data, every token coloured by its state at the end of the 10-epoch Plateau Dropping run: dropped as learned (green), dropped as stuck (amber), still learning (blue); grey text is untrained context (tool output). Turns were drawn at random (seed 0) among those with 120 to 320 trained tokens that contain both reasoning and a tool call; they are not hand-picked. Special tokens (<think>, <tool_call>, <eot>) are shown by name.

Appendix D: TailSFT coverage ratio

We compute the coverage ratio of Section 4.3 for vanilla SFT after 3 epochs. Over the 21 tasks the base model solves on only 1-2 of its 8 seeds, vanilla SFT loses more pass@8 coverage than it gains. Note that the ratio is computed only on these tasks, which the base model can already solve some of the time; overall, vanilla SFT still improves pass@8 by 12 points. Its coverage ratio (at k=8) is therefore 3.1, and 6 of the 42 tasks that the base model solves on some but not all of its seeds fall to 0/8 after SFT. With the paper’s own protocol (k=16, estimated from pass@1, over the 11 tasks the base model solves on 1 of its 8 seeds), the ratio comes out to be 12.8. The ratio falls below 1 only after longer training (0.06 for vanilla SFT at 10 epochs and learning rate 3e-4), a setting in which we did not run TailSFT. One caveat: the paper uses k=16, while we have only 8 samples per task, so our k=16 ratio is extrapolated from pass@1, and with only 11 to 21 tasks in the set, both ratios are noisy, though both stay above 1.

Appendix E: memory overhead of Plateau Dropping

Unlike vanilla SFT, Plateau Dropping needs to remember the loss history of every trained token. For each token, we store the last two losses it was seen with (float32), its visit count (int16) and whether it was dropped at its last visit (bool), i.e. 11 bytes per token. We key this history by (trajectory, position in the trajectory) rather than by position in the packed row, so it follows the trajectory wherever it ends up in a packed row. For simplicity, we store it as a dense table with a slot for every position up to the maximum context length of every trajectory, so on our SWE data (9,229 trajectories, 64k context) the table has 605M slots and takes 6.65 GB, of which only 30% (the trained tokens) are ever written. For comparison, the model and optimizer state (fp32 master weights and optimizer moments) take about 805 GB, so the history adds less than 1% to the checkpoint.

Implementation. We keep the table in host RAM as numpy arrays, outside JAX, with a full copy on every training process. Before each step, every process looks up the history of the tokens in the batch and passes it to the jitted train step as extra inputs, alongside the batch itself. Inside the step, the drop decision uses the per-token losses from the same forward pass that computes the training loss, so there is no extra forward pass. After the step, the per-token losses are gathered back to the host and every process updates its copy of the table, for the dropped tokens as well. The table is saved as a separate file next to the model checkpoint. All of this costs about 1.3% in step time (9.20 s per step against 9.08 s for vanilla SFT).

Scaling. Our implementation is far from optimal: the table grows with the number of trajectories times the maximum context length, and is replicated on every process. Storing only the trained positions would cut it to about 2 GB on our SWE data, storing the losses in fp16 and the count in uint8 would bring it down to 6 bytes per token, and sharding the table across processes would remove the replication. None of these were necessary at our data scale, but they would be at larger token budgets.

Citation

If you find this post useful, please cite it as:

@misc{varma2026plateaudropping,
  title  = {Plateau Dropping: a simple SFT rule for better agentic post-training},
  author = {Varma, Harshit and Lee, Donghyun and Dimakis, Alexandros G.},
  year   = {2026},
  month  = oct,
  url    = {https://hrshtv.github.io/blog/plateau-dropping/}
}