← Back to blog

Training Effective and Efficient RLMs


Recursive Language Models (RLMs) are an inference-time paradigm that allows models to scale to arbitrarily long contexts. Early this year, I started a project to train tiny RLMs to improve their accuracy and efficiency. During this experience, I achieved some pretty interesting results and faced a few engineering challenges that I think can be broadly relevant to anyone trying to train multi-turn agents.

In short, the main result is that, for the most stable run, I managed to train an RLM to more than double its performance while reducing the number of subcall tokens spent by more than 4x.

Below, I’ll briefly introduce RLMs (I strongly recommend you read the paper), then walk through the setup of my runs and a few tips for training, and finally present the results of the experiments.

I hope you’ll find this useful!

Background

So, an RLM is an inference paradigm where the context needed by a model to solve a task is held in a variable inside a REPL, outside of the neural network. The model can interact with the context symbolically, inspecting part of it, modifying it and potentially calling itself recursively on parts of it.

I think RLMs are very neat and come with some interesting properties that technically enable them to work over arbitrarily long inputs and outputs.

Schematic of the RLM loop, with a model interacting with its context through a REPL and calling sub-calls or sub-RLMs over parts of it.
Schematic of the RLM loop, with a model interacting with its context through a REPL and calling sub-calls or sub-RLMs over parts of it.

While the architecture enables the model + harness to achieve impressive results, RLMs still suffer from a few failure cases:

  • Long Tails: models sometimes make hundreds or thousands of unnecessary subcalls, leading to high costs and runtimes.
  • Failure to stop: even when finding a correct answer, an RLM may occasionally continue trying to verify or recompute it.
  • Complex protocol: while coding is nowadays in-distribution even for small models, the full RLM harness requires following a demanding interaction protocol.

I saw this in practice when testing the harness myself with small models: on the eval set, Qwen 3 4B (not trained) had a median of only 1.3k LLM-subcall tokens but a mean of 139k and a p95 of 679k. This seemed very inefficient, yet it felt somewhat similar to the early issue of balooning CoT length (this is a literary device known as foreshadowing).

When I started working on this, I thought the idea of training RLMs with RL was this super novel and brilliant and no one else would do it. I was wrong, but I think it can still be interesting to share my experience training them, because this experience broadly applies to various agent training setups.

Experimental Setup and Infra Stuff

When preparing to train the RLMs, I had 3 questions I wanted to work on:

  • Can we train a model to be an effective and efficient RLM? That is, can we improve accuracy while reducing the amount of tokens spent?
  • Does higher recursion depth improve the cost-accuracy trade-off? Meaning, does letting an RLM call child RLMs improve the accuracy or reduce cost for some tasks?
  • Is there a low-regret regime for training with LoRA? And how do different LoRA configurations affect stability, accuracy, cost, and policy behavior?

Answering these questions meant building a training setup where I could control the RLM harness and measure both correctness and cost. I built this as an environment using Prime-RL, which I found quite delightful to use and hack around with. Here, an environment is a container composed of three elements: a task, a harness, and a verifier.

On the data side, I created a dataset with three types of long-context tasks:

  • Deep research: For this, I took the tasks from the DeepMind FRAMES benchmark but drowned the useful Wikipedia pages in dozens of randomly sampled C4 pages.
  • Coding: Here I scraped a bunch of random Python codebases and asked questions generated by analyzing the AST (i.e., where does variable X change value, how many top-level functions are there, etc.).
  • Aggregation: These were long sequences of text records drawn from classification datasets, with questions asking about statistics over them.

This gave me 1,564 training examples and 452 evaluation examples, with the median context at 437k Qwen tokens and the largest at 1.9M. I created the dataset trying to make it as verifiable as possible, but I still chose to apply LLM as a judge when static checks failed, as deep research and coding tasks could have slight changes in formatting that should still be considered correct.

Once I had the data, Prime-RL had already solved a first big systems problem, as it already came with async training. This means that the inference and trainer are decoupled, so rollouts can keep flowing while the trainer works on the previous batch, with an upper limit on how stale rollouts can get.

Since an RLM can bounce between inference and a Python REPL for many turns, and then get stuck waiting for a broken subcall or messed up code, I replaced the single environment server with a pool of isolated workers and let only one rollout occupy each worker at a time, so a single broken job could not block future requests. The same thing applies to reward: especially when using an LLM as a judge, I’d advise splitting it across workers and adding a queue in front of it.

Overall, you might need to implement quite a few budgets and timeouts to cope with the unpredictability of training. I recommend strict budgets for the number of subcalls and turns, the number of tokens delegated to subcalls, and the amount of time per code execution. While this sounds limiting and you would ideally hope the model would learn to be reasonable just from training, in practice, training without these limits is super unstable and basically intractable.

Once you have something that kinda works, you’ll realize that having the model spin up tons of subcalls comes with its own set of issues. The most intuitive thing you can do regarding subcalls is to use the same model you use for root turns for the subcalls. This is nice because you already have your inference server available, so you can just point it in that direction and get the completions, but this comes with a few problems:

  • First off, early in training, the model is generally inefficient and it’s gonna spin up as many subcalls as you allow it to have, often all in the same turn. This will force you to find a balance in how many GPUs you need to allocate to inference versus training, which is a finicky trade-off that changes a lot during training.
  • The second issue I found is that, if you’re training a smaller model, the subcalls can often be quite useless. I was working with Qwen 3 4B which was often weak as a leaf model and struggled with tiny atomic tasks, especially when given a poor excerpt.
  • The third issue is how many of the subcalls you’d like to train on. Ideally you would want the quality of your leaf calls to improve during training, but if you’re in a compute-constrained environment, this is super expensive. In an early prototype I had put a hard cap at an arbitrary 40 subcalls per rollout, and even with the few rollouts I did per step, training on all subcalls used by the model was so slow it made it impossible to progress. I tried various metrics to be able to at least train a subset of the leaf calls, including random sampling or trying to figure out which were most useful for the final result, but it was incredibly finicky and basically didn’t work out at all, so I decided to abandon it.

The setup I ended up using was to use a cheap yet capable model over an API for the subcalls and only train the root RLM. My awesome friends at @seldon_tech sponsored me with tons of Gemini credits, so I landed on Gemini 3.1 Flash-Lite. For the final runs, I used a shared budget of 50 subcalls per rollout and masked those calls from training.

After having made all these decisions and set everything up, I chose to go with a fairly standard GRPO setup that fit comfortably on my node. I used GRPO-style group advantages with Prime-RL’s IPO-masked loss, a group size of 4, and a batch size of 64. I trained models for 150 steps, which isn’t a lot but generally allows you to get a solid idea of where training is going. I also removed all-zero and all-one groups from batches, as they provide no group-relative correctness signal.

For the reward, I took heavy inspiration from what is used for efficient CoT training (1, 2). In GRPO-style training for each prompt the trainer samples a group G of m rollouts, and each rollout receives a final grading that is either zero or one. The group solve rate is then:

s(G)=1m∑i=1mcis(G)=\frac{1}{m}\sum_{i=1}^{m}c_i

Prior work showed that you should avoid penalizing the model for tasks that it can still barely solve, thus I introduced two adaptive terms to normalize the penalty according to the task difficulty. First I introduced a coefficient that scales with the group solve rate:

β(s)={0,s≤s0if task mostly unsolved,βmax⁡(s−s01−s0)γ,s>s0if task mostly solved.\beta(s)= \begin{cases} 0, & s \leq s_0 \quad \text{if task mostly unsolved},\\ \beta_{\max}\left(\dfrac{s-s_0}{1-s_0}\right)^{\gamma}, & s>s_0 \quad \text{if task mostly solved}. \end{cases}

Then I added a component to make the penalty scale with tokens spent by all rollouts, with the goal to scale the penalty based on the minimum and maximum amount of tokens spent by the various rollouts in a group:

T~i={Ti−Tmin⁡Tmax⁡−Tmin⁡,Tmax⁡>Tmin⁡,0,otherwise.\widetilde{T}_i= \begin{cases} \dfrac{T_i-T_{\min}}{T_{\max}-T_{\min}}, & T_{\max}>T_{\min},\\ 0, & \text{otherwise}. \end{cases}

Another issue I noticed in early experiments was that the model for some reason became very lazy; sometimes it would compute the correct answer yet still waste turns until it was forced to finalize by the harness, or it would just spam the answer without using the correct format until the parser at the last step caught it. To avoid this, and put further efficiency pressure, I added another penalty that subtracted 0.25 from the final score if the model only finalized at the last step (and got the answer right). Putting all of these elements together the reward looked like this:

Ri=clip⁡[−0.5,1](ci−pi−β(s)T~i)R_i=\operatorname{clip}_{[-0.5,1]}\left(c_i-p_i-\beta(s)\widetilde{T}_i\right)

To sum it all up: in each group of four, cost pressure stays at zero until at least two rollouts are correct, then increases with the solve rate and in proportion to the number of tokens spent by each rollout, with the laziness penalty encouraging models to terminate as soon as they can. Considering my setup with four rollouts, reward ends up looking somewhat like this:

Correct rolloutss(G)s(G)β(s)\beta(s)Penalty at T~=0.5\widetilde{T}=0.5Penalty at T~=1.0\widetilde{T}=1.0
0–10.00–0.250.000.0000.000
20.500.050.0250.050
30.750.100.0500.100
41.000.150.0750.150

The penalty turned out to be quite low in practice, but this worked well for me anyway and tweaking parameters is not a particularly easy nor fun experience.

Finally, once the infra was set up I came up with a few experiments to poke at the questions I introduced:

  • I swept nine LoRA configurations across three ranks (4, 16, and 64) and three learning rates (5 × 10⁻⁷, 10⁻⁵, and 10⁻⁴) training the model with depth 1 recursion.
  • I ran a full-parameter fine-tune to compare against LoRA.
  • I enabled depth-2 recursion for the LoRA configuration with the highest pass@1 on the full evaluation set.

Results and Discussion

Training ended up working pretty well. Below is a table reporting pass@k of the best checkpoints on the eval set, correctness, mean number of subcalls, and mean number of LLM-subcall tokens.

Modelp@1p@5p@10Avg. correctMean subcallsMean tokens
Base Qwen3 4B22.555.066.822.58.68139k
LoRA r4, lr 10⁻⁴49.773.981.049.72.9231.6k
LoRA r16, lr 10⁻⁴47.066.571.747.01.552.2k
LoRA r64, lr 10⁻⁵55.375.180.855.33.9319.3k
Depth-2 r64, lr 10⁻⁵53.771.076.153.72.9920.7k
Full FT, lr 10⁻⁵64.874.677.064.80.43433

The figure below shows pass@k vs mean LLM-subcall tokens per rollout of the models above.

Mean LLM subcall tokens per rollout (thousands, log scale).
Mean LLM subcall tokens per rollout (thousands, log scale).

I also tested the checkpoints on some long-context benchmark tasks as in the RLM paper, below are the results for those (percentages apart from OOLONG-Pairs which uses F1).

ModelBrowseComp-Plus (N=150)LongBench v2 Code QA (N=50)OOLONG-Pairs (N=20)OOLONG TREC-coarse (N=50)
Base30.738.00.016.0
Depth-236.038.00.328.0
LoRA r434.044.00.330.0
LoRA r6442.740.02.934.0
Full FT18.024.010.514.0

The main result is that training made the model much better while using far fewer LLM-subcall tokens. The base model reached 22.5% pass@1 while spending 139k subcall tokens on average. The best LoRA run reached 55.3% at 19.3k subcall tokens, while full fine-tuning got to 64.8% and only 433 subcall tokens. So the model learnt to be much more efficient with its subcalls. Nice!

This has the caveat that there wasn’t really a single best LoRA setup, with adapters learning surprisingly different behaviours. The full fine-tuning of all parameters was also interesting as it clearly somewhat overfit, even regressing on performance from the long-context benchmarks and underperforming in higher pass@k.

Furthermore, not every run survived. Rank 64 at learning rate 1e-4 diverged into extremely long rollouts and I stopped it after 72 of 150 steps, while 5e-7 mostly underfit across all ranks. In practice, learning rate mattered as much as LoRA rank.

The behavior during training was probably the most interesting part. Early on, the model would make broad, expensive subcalls for almost everything. It then overcorrected and mostly stopped using them while mostly solving some easier tasks with just code, before eventually recovering a more selective strategy with short, targeted calls. You can see the behavior in the charts below, which plot correctness, LLM usage, and all-zero groups (prompts that are too hard at a certain step) over time.

Training history for the LoRA rank-16, 10−4 run (similarly applies to all training runs).
Training history for the LoRA rank-16, 10−4 run (similarly applies to all training runs).

I know correctness looks mostly flat after a while, but it’s mostly because a prompt that was “too hard” (i.e., led to an all-zero group) was not discarded but just sidelined for a few steps and later reinserted into the training mix. So, effectively, during training, the RLM was confronted multiple times with hard tasks.

The shaped reward also helped the model learn to follow the harness protocol. For rank 4, the rate of rollouts missing a final answer fell from 37% at the start of training to 0.16% near the end.

Deeper recursion was a bit underwhelming. The model used it in only about 0.1% of eval rollouts, and depth 2 did not improve the overall accuracy-cost trade-off. This might be a data or exploration problem: certain tasks and modalities may be a better fit, or the model may need an incentive to try child RLMs more often, but I didn’t test either.

Conclusion

I am quite happy with the project. I think RLMs are a really cool innovation and that training can make them extremely effective and efficient. There were definitely a few limitations, from my use of a separate model for subcalls to not having a baseline trained without the cost penalty to the rank 64 run at high learning rate completely collapsing during training.

I explicitly tried to include some fuckups in the writeup. I think there’s a general tendency towards not speaking of failures or limitations that makes it harder both to understand design choices and to replicate results or build upon previous work. Don’t do it! Share your fuckups! It’s cool!

I’m really interested in seeing larger-scale training of RLMs, with different reward shaping and new datasets. I’d also love to see RLMs trained to work natively across modalities.

The dataset I used and the checkpoints are available here.

Part of the setup described above was adapted in this paper by my @nzuma0, in which he also introduced a lot of other cool stuff. Check it out! (I’m also a coauthor).

You can technically read a full version of what’s described above with all of the details and sauce in my master thesis but I would advise against it. Due to the nature of a master thesis it’s extremely long, boring and sloppy. If you’re 100% sure of wanting to check it out you can send me a DM and I’ll send it over.

I’d also love to chat with anyone interested in these same topics!

Special thanks to @seldon_tech for providing me with an infinite supply of tokens (and a job) and to @PrimeIntellect for the infra that allowed me to train everything.