Bhavya Kailkhura

Recent advances in depth-recurrent language models show that recurrence can decouple train-time compute and parameter count from test-time compute. In this work, we study how to retrofit existing pretrained non-recurrent language models into depth-recurrent models. We find that using a curriculum of recurrences to increase the effective depth of the model over the course of training preserves performance while reducing total computational cost. In our experiments, on mathematics, we observe that retrofitting pretrained models to recurrent ones results in better performance at a given compute budget than simply post-training the original non-recurrent language model.

Large language models have shown remarkable ability in serial code generation, but they still struggle with parallel code for which training data is comparatively scarce. A common remedy is to use coding agents that interact with external tools, but tool calls can be costly and sometimes impractical, e.g., for partially written code. We propose *Parallel-Code World Models (PCWMs)*, reasoning LLMs that aim to predict tool outcomes directly from parallel source code. To train PCWMs, we design a novel exploration and data generation pipeline that samples diverse parallel-coding problems and candidate implementations across multiple domains, then executes them via tools to record data races and performance profiles. From these, we synthesize reasoning traces that aim to causally connect source code to observed tool outcomes. Fine-tuning on the resulting data yields noticeable gains, with a 7B-parameter world model improving from 64.3\% to 72.8\% accuracy in race-outcome prediction, while an 8B-parameter model improves in a performance profiling task from 49.3\% to 58.6\% accuracy. Furthermore, when open-weight models were tasked with fixing data races, world-model feedback improved their race-fixing rates relative to self-feedback by 2.7\%-9.1\% using our 7B-parameter world model and by 6.1\%-11.1\% using our 14B-parameter world model. Our results suggest that reasoning world models may have the potential to serve alongside external tool calls in parallel-coding agents.

The reverse Kullback-Leibler (KL) divergence between the trained policy and a reference policy is commonly used as a regularizer when training large language models (LLMs) with reinforcement learning (RL). Since computing the exact sequence-level KL divergence is intractable, practical algorithms use sample-based estimators computed from on-policy rollouts, and incorporate them in the reward or the loss. Despite its ubiquity, the interaction between estimator choice, position of the regularization and downstream performance is poorly understood. Recent work also identifies that certain implementation choices can result in biased gradients. We further analyze these practices and examine the gradients of several estimator configurations, showing how the choices shape gradient bias. We substantiate these observations with empirical evidence by RL fine-tuning LLMs with different KL configurations, and evaluating their performance on both in- and out-of-distribution tasks. In on-policy settings, we find that biased-gradient estimator configurations can cause training instabilities, while unbiased-gradient configurations lead to better performance on in-domain as well as out-of-domain tasks. We also observe that KL regularization can help stabilize asynchronous training. Overall, our findings provide useful takeaways for using KL-regularized objectives during RL post-training of LLMs.