1 Introduction↩︎

Large language models (LLMs) continue to improve in automatically verifiable domains like competitive programming. Currently, the dominant strategy for state-of-the-art performance relies on repeated sampling: generating numerous potential solutions to pass a verifier: the pass@k regime [2]. However, the major limitation of this approach lies in its significant GPU computational costs for scaling inference. Attempts to mitigate this through reinforcement learning post-training have primarily focused on improving single-sample accuracy (pass@1) at the cost of pass@k in high sampling regimes for \(k >> 100\). This reduction in diversity undermines its effectiveness when scaling inference at test time, because performance gains grow logarithmically for linear computational costs (Appendix 10).

Our approach targets tasks where solutions are hierarchically decomposable into independently implementable modules and where correctness is automatically verifiable. This structural assumption holds broadly: competitive programming problems naturally decompose into helper functions; formal proofs factor into lemmas [3]; hardware modules compose into chip designs [4]; and scientific codebases separate data processing, modeling, and evaluation [5]. Whenever verification is cheap (unit tests, proof checkers, simulators) but generation is expensive (LLM inference), modular recombination offers a path to scale search without proportionally scaling GPU cost.

We propose DecompRL: a reinforcement learning framework that fundamentally shifts the scaling bottleneck from expensive GPU inference to cheap CPU evaluation. Inspired by modular inference strategies [6], DecompRL trains models to decompose challenges into a hierarchy of independently solvable sub-functions that can be re-combined to form a complete solution. For example, for a sorting task, instead of generating a single monolithic solution, the model identifies \(N\) sub-problems such as partitioning, merging, and base-case sorting. By generating \(K\) implementations for each of the \(N\) modules in parallel, we can recombine them into up to \(K^N\) unique candidate solutions. This approach creates a combinatorial explosion of potential programs for a fixed inference budget of only \(K \times N\) tokens, allowing us to scale search on cheap CPU clusters without increasing the GPU footprint.

By framing hierarchical modular inference as an RL objective, DecompRL prioritizes exploration (Section 4.1), maximizes the utility of recombinations (Section 4.2) and successfully discovers sparse rewards (Section 4.3). We do not position DecompRL as a replacement for standard RL post-training, which remains more sample-efficient when the starting policy is already strong. Rather, DecompRL is an exploration tool: it discovers correct solutions to problems that standard high-pass@\(k\) sampling cannot reach, making it a natural complement for policy distillation pipelines and for overcoming the cold-start problem on hard tasks where the base policy has near-zero solve rates.

2 Method↩︎

We first introduce hierarchical generation in the context of problem solving (Section 2.1, 2.2) as described in Parsel [6] and then show how we adapt this to train decomposition and implementation policies (Section 2.3).

2.1 Hierarchical inference↩︎

We consider the task of generating code in a modular fashion, where model users have tighter control over the generation process than in the standard setup of whole-code generation from an autoregressive language model. Standard whole-code generation models the joint distribution of a piece of code of multiple parts \(I_1, \ldots, I_n\) by means of its autoregressive decomposition and can optionally include a plan \(D\) sampled before \(I_1\). Suppressing conditioning on the problem \(x\), we have: \[\pi(I_1, \ldots, I_n) = \pi(I_1 |D)\prod_{i=1}^n \pi(I_i \mid I_{i-1}, \ldots, I_1, D).\]

In this work, we instead consider a hierarchical decomposition: \[\pi(D, I_1, \ldots, I_n) = \pi(D) \prod_{i=1}^n \pi(I_i \mid D)\] which models the implementations of the pieces \(I_i\) as conditionally independent of each other given \(D\), and captures all interaction in the decomposition plan \(D\).

In practice where we consider code generation tasks in Python, we represent decompositions \(D\) by a set of functions specified by their signatures and natural language specifications in the form of docstrings. The parts \(I_i\) are then implementations of those functions. We give an example of a decomposition in Figure 1 and of a full hierarchical trajectory \((D, I_1, I_2)\) in the Appendix 20. We use prompted and trained language models for the decomposition and implementation policies \(\pi(D)\) and \(\pi(I_i \mid D)\) (see Appendix 17 for full prompts).

2.2 Recombination↩︎

Hierarchical inference comes with an important benefit: while the conditional independence assumption makes individual samples less precise, it can be exploited to generate samples cheaply by recombination. Here and in the following, let \(D\) be a decomposition of size \(n\) and \(I_i^j, 1\leq i \leq n, 1 \leq j \leq k\), be \(k\) implementations per function \(I_i\) described by \(D\). From these \(nk\) implementation samples, we can create \(k^n\) complete trajectories by considering all combinations \((D, I_1^{j_1}, \ldots, I_n^{j_n})\) for \(1 \leq j_1, \ldots, j_n \leq k\). In other words, the number of complete trajectories scales polynomially in \(k\) while the inference cost scales linearly.

This is particularly useful in domains where rewards are only observed for complete trajectories, which is the case for verifiable reasoning domains like mathematics and code. In such setups, verification is a cheap CPU operation while generation is an expensive GPU operation involving large language models. In this work, we consider the task of competitive programming where the model must produce a code solution to a problem in natural language. The reward is \(1\) if the code passes all private tests (unknown to the model at inference time) and \(0\) otherwise.

2.3 Reinforcement learning: policy gradient estimators↩︎

As customary with large language models, we consider policy gradient reinforcement learning algorithms [7]. We write \(\tau = (D, I_1, \ldots, I_n)\) for a trajectory, \(r(\tau)\) for the reward it receives and \(\pi_{\theta}\) as the policy given by a large language model. According to the policy gradient theorem [8], an unbiased estimate of the gradient \(g\) of the expected reward is given by \[g = \nabla_\theta \mathbb{E}_{\tau \sim \pi_\theta}[r(\tau)] = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{a \in \tau} (r(\tau) - b) \nabla_\theta \log \pi_\theta(a) \right],\] for any baseline \(b\) which does not depend on the action \(a \in \{D, I_1, \ldots I_n\}\).

We can now view recombination as a way to produce Monte-Carlo estimates of the policy gradient for hierarchical inference more efficiently. Given \(k\) implementations \(I_i^j\) of the functions \(I_i\) in the decomposition \(D\), the standard estimator for \(g\) used in sequential reinforcement learning is \[\hat{g}_\text{standard} = \frac{1}{k} \sum_{j \leq k} \sum_{a \in \{D, I_1^j, \ldots, I_n^j\}} (r(D, I_1^j, \ldots, I_n^j) - b) \nabla_\theta \log \pi_\theta(a),\] which is optimal in the case of linear (autoregressive) rollouts where \(I_k\) depends on \(I_j\) for \(k > j\).

For hierarchical inference, however, we made the modeling assumption that implementations are conditionally independent and all interactions are captured by their dependence on the decomposition \(D\). In this case, we can evaluate all combinations of implementations generating a matrix \((I^{j}_i)_{j,i} \; \forall i,j\) for all functions. We can obtain the improved estimator: \[\hat{g}_\text{hierarchical} = \frac{1}{k^n} \sum_{j_1, \ldots, j_n \leq k} \sum_{a \in \{D, I_1^{j_1}, \ldots, I_n^{j_n}\}} (r(D, I_1^{j_1}, \ldots, I_n^{j_n}) - b) \nabla_\theta \log \pi_\theta(a),\] of which \(\hat{g}_\text{standard}\) represents the diagonal terms \((I^{j}_i)_{i,i} \; \forall i\). Clearly, both estimators are unbiased, but we have \(\mathop{\mathrm{Var}}(\hat{g}_\text{hierarchical}) \leq \mathop{\mathrm{Var}}(\hat{g}_\text{standard})\), as we show in Theorem 2.

2.4 Reinforcement learning objectives↩︎

We do not expect hierarchical inference to improve upon whole code generation when compared on a single-attempt (pass@1) basis. This is because hierarchical inference comes with a loss in precision (conditional independence instead of joint modeling of implementations) and with only one sample we do not gain anything from the recombination. Instead, we propose to focus on the task of efficient inference scaling, i.e, comparing pass@k [2] metrics for large values of \(k >> 100\).

Given \(n\) rollouts \(\tau_j\) with rewards \(r = (r(\tau_1), \ldots, r(\tau_n))\) and a multi-sample objective function \(f(r)\), [9] optimize \({E}_{\tau_i \sim \pi_\theta} [f(r)]\). Here, \(f\) is typically a function of a variable number of arguments, permutation-invariant and monotonic in each argument such as the maximum function. They show this can be achieved directly by using gradients for any baseline \(b\) independent of \(\tau_i\) such as the leave-one-out baseline \(b = f(r_{-i}) := f(r_1, \ldots, \hat{r}_i, \ldots, r_n)\) which omits \(r_i\), giving \[g = \mathbb{E}_{\tau_i \sim \pi_\theta} \left[ \sum_i \left(f(r) - f(r_{-i})\right) \nabla_\theta \log \pi_\theta(\tau_i) \right].\]

In the case of hierarchical inference where the solutions are recombined before evaluation, we have a large number of correlated rewards and optimize for large pass@\(k\). Unlike the original paper [9], we can’t take \(f\) as the maximum over all samples since it would saturate and lead to extremely sparse advantages: a positive advantage would only be observed if \(\tau_i\) is the only solution among \(k\) attempts that solves the task (see Figure 10 in the Appendix). To adapt multi-sample objective training to this use case, we propose to use soft objective functions that interpolate between the hard maximum (pass@k training) and the average (standard training). We opt for the \(\mathop{\mathrm{\mathrm{logmeanexp}}}_\beta\) function defined by \[\label{eq:lme} \mathop{\mathrm{\mathrm{logmeanexp}}}_\beta(r_1, \ldots, r_n) = \beta \log \left( \frac{1}{n} \sum_{i=1}^{n} e^{r_i / \beta} \right).\tag{1}\]

This function smoothly interpolates between the average (\(\beta \to \infty\)) and the maximum (\(\beta \to 0\)), and for a well-chosen \(\beta\) closely approximates a log-uniform mixture of pass@\(k\) objectives (see Figure 3), balancing exploration and exploitation (see Appendix 11).

Concretely, in the setup of hierarchical inference, we apply this form of reward aggregation as follows. We generate \(d\) decompositions \(D_i\) of size \(n_i\), \(1 \leq i \leq d\), respectively and \(k\) implementations for each function in each \(D_i\). For each decomposition, we evaluate \(m \leq k^{\max{(n_1,..,n_d})}\) combinations. Write \(r\) for all \(dm\) rewards obtained in this way, \(r_{ij}\) for a single reward, \(i \leq d, j \leq m\) and \(A(r_{ij})\) for the set of actions involved in \(r_{ij}\).

DecompRL is a cooperative framework of two policies: decomposition \(\pi_\theta(D)\) and implementation \(\pi_\theta(I \mid D)\). For the decomposition policy, we use the policy gradient \[\label{eq:objective} g = \mathbb{E}_{r_{ij}} \left[ \sum_{i=1}^{n} \left(\mathop{\mathrm{\mathrm{logmeanexp}}}_\beta(r) - \mathop{\mathrm{\mathrm{logmeanexp}}}_\beta(\{ r_{i'j} \mid i' \neq i) \})\right) \nabla_\theta \log \pi_\theta(D_i) \right].\tag{2}\]

For the implementation policy, assume \(d=1\) and write \(r_j\) for \(r_{1,j}\). Then we use the policy gradient \[g = \mathbb{E}_{r_j} \left[ \sum_{l \leq n, j_l \leq k} \left(\mathop{\mathrm{\mathrm{logmeanexp}}}_\beta(r) - \mathop{\mathrm{\mathrm{logmeanexp}}}_\beta(\{ r_j \mid I_l^{j_l} \notin A(r_j) \})\right) \nabla_\theta \log \pi_\theta(I_l^{j_l} \mid D) \right].\]

In other words, the advantage is computed from the multi-sample objective using all observed rewards and baselined by the multi-sample objective computed from all rewards that an action did not participate in. Because, by this construction, the baseline is independent of the action, it does not bias the gradient estimate. The leave-one-out baseline reduces the overall gradient variance (see Appendix 13 for the computation adapted to multi-sample objectives). Empirically, the reward tensor over recombinations is well-approximated by a rank-1 (Appendix 18), confirming that credit concentrates along individual function axes and that the leave-one-out baseline provides reliable advantage estimates from a sparse subset of \(k^n\) combinations.

3 Experimental setup↩︎

3.1 Overview↩︎

Datasets and models: We perform experiments with three different large language models: Qwen 2.5 7B [10], Llama 3.1 8B Instruct [11] and the Code World Model 32B [12]. We train all models with online reinforcement learning on a mix of \(15,000\) competitive programming problems including the CodeContest [13] and TACO [14] training sets. We evaluate on the DeepMind Code Contest validation set [13] (117 problems) and LiveCodeBench range \(2024/08/01\) to \(2025/02/01\) (279 problems) [1]. All evaluations are done with temperature \(1.0\) and top-p \(1.0\) to promote answer diversity.

Algorithms: Similar to [15], [16], we optimize policy gradients using Proximal Policy Optimization (PPO) [17] without a critic model using the advantage described in section 2.3. We use an asynchronous distributed RL framework [18], [19] where we divide \(80\) H100s into a set of workers (producing samples) and trainers (updating the current policy). We use \(72\) workers and \(8\) trainers. All code generations are evaluated on an external CPU cluster.

Baselines: We compare with Group Relative Policy Optimization (GRPO) and algorithms optimized for high pass@\(k\) performance specifically: pass@\(k\) training [20] (set \(k = 8\)), Soft Policy Optimization (SPO [21]) and our own \(\mathop{\mathrm{\mathrm{logmeanexp}}}\) leave-one-out objective (see Equation 2 ) with \(\beta = 0.3\). All methods use \(16\) samples per prompt. We also include a \(\mathop{\mathrm{\mathrm{logmeanexp}}}\) objective with \(48\) samples per prompt and \(\beta = 0.1\) to mimic the number of forwards per decomposition present in DecompRL training (at most \(6\) functions and \(8\) implementations per problem). See hyperparameter and baseline tuning details in Appendix 15 and 16.

3.2 DecompRL as Multi-Agent Reinforcement Learning↩︎

DecompRL is a form of cooperative multi-agent RL where a decomposition and an implementation policy maximize the number of correct code combinations. Multi-agent reinforcement learning poses challenges not present in single-agent reinforcement learning. Agent policies are interdependent and non-stationary, which raises the questions of how to perform credit assignment after jointly observed rewards and how to optimize toward a social optimum (in the cooperative case) [22]. To address the moving target problem, we use sequential training and a counterfactual action value estimation [23] with our leave-one-out baseline. Similar to Expectation-Maximization algorithms [24], our training has two stages. The decomposition and implementation policies are represented by two different copies of our model. First, we train the decomposition policy for \(30k\) steps keeping the implementation policy fixed. Second we train the implementation policy also for \(30k\) steps keeping the decomposition fixed. For each problem, we use \(8\) decomposition samples. For a decomposition with \(n\) functions (\(n_{max} = 6\)), we implement each \(8\) times yielding \(8^n\) code combinations. We randomly sample \(512\) to evaluate out of \(8^n\) to balance exploration between breadth-first (exploring new problems) and depth-first search (scaling evaluations per problem). Although we have two copies of the model, only one is trained at a time so this has similar cost as the reference policy used in traditional RL algorithms.

4 Results↩︎

DecompRL produces large numbers of diverse samples that lead to diverse solutions (Section 4.1), learns modularity (Section 4.2) and solves new problems (Section 4.3)

4.1 DecompRL leads to diverse solutions↩︎

A decomposition with \(n\) functions and \(k\) implementations per function costs \(nk+1\) model forwards and yields \(k^n\) (correlated) code combinations for evaluation. For hierarchical inference to make a difference, we need the final solutions to have enough diversity to offset their generation cost compared to \(nk+1\) independent whole-code samples. With the Llama 3.1 8B and Qwen 2.5 7B, we show training the decomposition policy increases success rates per decomposition and number of evaluations.

In Figure 5 (a), we investigate the diversity of recombined programs by fixing a maximum number \(m\) of code evaluations per decomposition and observing the relationship between \(m\) and the overall success rate of a DecompRL model. It can be seen that the success rate from less than 50 model forwards per problem shows no sign of plateauing even when scaling \(m\) up to 1000 evaluations. Crucially, this remains true throughout reinforcement learning training as shown in Figure 4 (b) where the solve rate continues to increase given a fixed evaluation budget of \(m = 512\) code samples.

4.2 Recombination matters↩︎

The benefit of recombining partial solutions is most pronounced when the decomposition size \(n\) is large where \(n\) is the number of functions picked by the model. In our experiments, we see that models with similar success rates using 10 decompositions behave very differently when inference is scaled to 100 decompositions. As we scale the inference budget per problem, models that create larger decompositions improve faster than those with smaller ones. The success rate with 10 decomposition attempts does not predict the success rate with 100 attempts, but the size of the decomposition does (see Figures 5 (b) and 5 (c)).

4.3 Solving new problems↩︎

DecompRL through recombination and diverse evaluations solves new problems that are harder (Figure 1) and only solvable at high inference budgets (Figure 4 and Table 1). Using different models with DecompRL training, we beat the best existing methods for reinforcement learning when using high inference budgets (up to 3 million tokens per problem). We evaluate \(m = 4096\) combinations per decomposition, which is roughly 80 times more than \(nk+1\) language model calls would have produced with standard whole-code generation. On the other hand, \(m = 4096\) is small compared to the theoretical maximum \(k^{n_{\mathrm{max}}} = 8^6\) of combinations that hierarchical inference can produce. Additional CPU execution resources could further increase the appeal of the method according to the scaling behavior in Figure 4. DecompRL can help overcome the limits of our SFT checkpoint and find solutions not available with regular inference. We see it as a useful way to gather offline data to improve the starting policy.

5 Analysis↩︎

5.1 Ablation Study on DecompRL↩︎

In standard RL with LLMs, one action leads to one reward and with GRPO we compute the advantage over a group \(G\) of rewards where \(G << 100\) due to inference limitations. In DecompRL, we instead have two policies performing \(k \times n\) actions which lead to \(k^n\) rewards. The advantage of each action is therefore estimated by many indirect and correlated rewards. In order to prevent over-estimating certain actions, we introduce in equation 2 : 1) a leave one out instead of mean baseline, 2) transforming the rewards per action using a logmeanexp function (Equation 1 ) 3) training sequentially a set of two policies. Figure 6 shows how each component prevents instabilities. We run ablations on the Code World Model 32B.

Ablation on logmeanexp. We compare our max interpolation (logmeanexp, \(\beta = 0.03\)) with a mean interpolation (softmax, \(\beta = 0.5\)). As we have many noisy rewards optimizing towards the mean eventually leads to reward collapse. Lower \(\beta\) in \(\mathop{\mathrm{\mathrm{logmeanexp}}}\) also leads to faster decomposition size decay (see Figure 4 (c)).

Ablation on sequential training. Instead of alternating between training the implementation and decomposition policies, we study two other strategies: training only the implementation policy, and jointly training both policies (meaning no frozen model). Without clear credit assignment and with an over representation of implementation samples, we overfit the implementation policy and underfit the decomposition policy. Joint training leads to entropy collapse (Figure 6) whereas sequential training allows us to keep scaling the number of functions with increased gains in performance (Figure 5).

Ablation on leave one out baseline. Removing the leave one out baseline in the advantage calculation for both policies leads to much lower rewards and earlier plateau.

5.2 Computational cost: GPU to CPU bottleneck shift↩︎

A key advantage of DecompRL is that it shifts the dominant cost from GPU inference to CPU evaluation. Table 2 reports wall-clock breakdowns for Qwen 2.5 7B on our training setup (see details in Appendix 16).

At 512 evaluations per problem, DecompRL generates only \(\sim\)​4k tokens (one decomposition plus \(n \times 8\) short function implementations) versus \(\sim\)​198k for standard sampling, a \(\sim\)​50\(\times\) reduction in GPU cost. The evaluation cost is identical in both cases since both evaluate the same number of candidate solutions. This makes the total wall-clock time dominated by CPU execution rather than GPU generation.

When parallelized, the wall-clock saving from DecompRL scales with how GPU-bound the system is: at 8 GPUs with 128 remote execution threads, where generation dominates, DecompRL delivers a 3.6\(\times\) speedup; at 8 GPUs with 64 threads, a 2.3\(\times\) speedup; while at 72 GPUs with 128 threads, where remote execution already dominates, the saving shrinks to 1.3\(\times\). Since our cluster is heavily GPU-optimized, we argue scaling the number of CPUs (which is often cheaper than scaling GPUs) could unlock even larger gains for DecompRL since the marginal cost of evaluating more recombinations is essentially free on the GPU side.

6 Limitations and Discussion↩︎

6.0.0.1 Format tax.

Hierarchical inference without RL training underperforms standard inference, with pass@1 dropping by 3 percentage points for Qwen 2.5 7B and performing similarly below the standard baseline for CWM 32B (Table 4). DecompRL with Llama 3.1 8B Instruct model degrades performance from \(26.6 \%\) to \(14.5 \%\) on the CodeContest validation with the same token budget (10 attempts with DecompRL, corresponding to pass@\(316\) for standard inference). On the “easy" split of LiveCodeBench, DecompRL performs worse than the starting policy using Qwen 2.5 7B (see Figure 1). More generally, for easy problems and low token budgets, the format tax of decomposition outweighs the gains from additional recombined evaluations, since a strong base model can already solve these problems in a single monolithic attempt. Looking at the CodeContest training set as an example of competitive programming data, we notice decomposition and function generation are off-policy tasks for the model. Only 65 % of the provided solutions in the CodeContests training set contain functions. We expect data in the training stage of the Llama 3.1 8B Instruct model to follow a similar distribution and potentially not encourage modular coding. The model has also been heavily fine-tuned through RLHF [25], human annotations and synthetic data for the specific task of generating full code answers when prompted with problems in the standard inference setup. DecompRL unlike other reinforcement learning algorithms has to overcome this off-policyness during training. RL training closes this gap, with DecompRL surpassing even standard inference for both models (Table 4).

6.0.0.2 Size collapse.

Hierarchical inference comes with the cost of reduced precision at the single sample level. Decomposition training leads to a decreasing decomposition size (see Figure 4 (c)). This can allow for more relevant functions but in the limit if we have a single function, hierarchical inference becomes equivalent to classical whole-code generation. We see this as a form of reward hacking during training which can be delayed by using a lower \(\beta\) in the \(\mathop{\mathrm{\mathrm{logmeanexp}}}\) objective and alternating between implementation and decomposition training. Simpler approaches such as bonuses based on the decomposition size led to reward hacking.

6.0.0.3 Sampling time.

During training, DecompRL scales the number of evaluations per problem to up to \(512\) in our experiments compared to the default \(16\). This allows for more exploration during training (see Figure 5 (a)) but can also lead to: slower worker GPUs and introduce off-policyness in online RL. We see DecompRL as a policy distillation method for gathering new solutions on the training set not available with regular RL. If we have a good enough starting policy using regular RL remains more efficient as a post training strategy.

7 Related works↩︎

7.0.0.1 Scaling inference methods.

Intermediate steps to help models generate correct code solutions has been studied at the “plan”, “algorithm” and “function” level [26]–[30]. We adapt the decomposition and function approach from [6]. To promote sampling diversity, optimizing over code generation trajectories has been studied mainly with repeated sampling [13], self-refinement chains with execution feedback [31], [32], and tree search using multiple children solutions per chain [33]–[35]. Tree based generation systematically organizes and potentially directs linear generation, but generally does not create a combinatorial increase in solutions to evaluate (each evaluated leaf has to have been generated).

7.0.0.2 Recombination of partial solutions.

Methods that decompose a problem into separate parts which can be solved independently and then recombined have been studied for solving mathematical and coding tasks. In mathematical theorem proving, tree search in formal theorem proving systems decomposes the problem into sub-goals that can be solved and checked independently by the verifier [3], [36]. In code generation and symbolic regression, library learning creates solutions to sub-problems for later reuse [37]–[39]. These methods study re-combination at the inference level with a fixed policy for determining blocks whereas we propose reinforcement learning to improve the policy. Programs can be analyzed hierarchically [40] or generated hierarchically [6]. Via clustering in the embedding space (CodeChain [41]), prompting different agents MapCoder [42] and iterating on trees of functions based on shared consensus (Divide and Conquer [43]), inference methods have enforced modularity in code generation trajectories.

7.0.0.3 Self-play reinforcement learning.

In formal mathematical proving, a conjecturer model asserts statements for a prover to prove, either as a self-guided exploration task with intrinsic motivation [44] or as a live augmentation technique during a standard reinforcement learning run for theorem proving [45]. In the domain of code, self-play involves a programming puzzle generator and a model that solves the generated tasks [46], [47]. In DecompRL, the decomposition model can be viewed as a self-play teacher that is conditioned on the underlying programming problem to be solved.

7.0.0.4 Expert Iteration.

Expert iteration [48] divides the problem into an “expert” (a slow tree search algorithm doing the heavy exploration) and an “apprentice” (a fast neural network). As DecompRL is optimized for harder problems and high pass@\(k\), we can view it as the “expert” policy gathering novel solutions to train a starting RL policy offline (the “apprentice"). Online expert iteration methods such as STaR [49], Quiet-STaR [50], ReST [51], and ReST-MCTS [52] use the same model as expert and apprentice to iteratively bootstrap its reasoning capabilities using successful trajectories. Similar to our hierarchical inference process, ReST-MCTS goes beyond regular inference using tree search to find solutions normally missed by a weak starting policy.

8 Conclusion↩︎

Solving programming problems by decomposing them into simpler subproblems and recombining them into full solutions allows scaling inference-time evaluation budgets with a limited number of language model calls. However, off-the-shelf instruction finetuned models do not fare well at this task. We introduce DecompRL, a reinforcement learning technique that directly trains models for hierarchical inference, greatly enhancing their performance on this task. This method encourages diversity by directly optimizing for a mixture of pass@k objectives with a novel \(\mathop{\mathrm{\mathrm{logmeanexp}}}\) advantage function. The resulting models use recombination to solve problems unreachable with regular RL training. With evaluation up to pass@1k or \(10^7\) generated tokens per problem, we show with different models (Llama 3.1 8B Instruct, Qwen 2.5 7B) that the success rate of DecompRL continues to scale whereas other methods begin saturating. With larger CPU resources, we envision the potential for significant advancements when scaling such systems to millions of evaluated answers. Moreover, the modularity of the resulting generation pipeline offers avenues for interpretability insights and an axis of control which is not available in other methods of scaling inference compute, such as in repeated sampling or scaled chain-of-thought.

A promising direction is to use DecompRL as the explorer stage of an explorer/distiller pipeline [53], [54]: DecompRL generates diverse correct solutions at high budgets, and a standard pass@1 policy is then distilled from these trajectories via supervised fine-tuning. This two-stage protocol could combine DecompRL’s exploration strength with the efficiency of monolithic generation at deployment time. We also note that problem types that are naturally parallelizable or have a hierarchical structure such as multi-step proofs, nested algorithmic routines, or layered engineering systems are likely to benefit most from the modular decomposition that DecompRL encourages.

Looking forward, many scientific and engineering tasks are more easily evaluated than solved, presenting opportunities for innovative solutions. For writing entire codebases that solve specific problems, models could achieve acceptable pass rates at scale (e.g., pass@1M), even if their initial performance (pass@1) is poor. Our findings suggest that DecompRL is a high performing search policy to solve new and harder problems and could become an ingredient of any training pipeline where recombination is possible and evaluation is cheap while generation is expensive.

More generally, our findings shed light onto the hard exploration question in LLM post-training [55]: How can models discover new behaviors that allow solving previously unsolved tasks without the need for human annotations, and more efficiently than via finetuning on large-scale rejection sampled datasets [36], [56]? While DecompRL is specific to code generation, it hints at a possible future where exploration policies are trained specifically for creating high-quality data with behaviors that the final, exploiting policy trained with a standard pass@1 objective could not have found by itself over the course of its training.

9 Proofs↩︎

9.1 Variance Reduction↩︎

In this section, we give the proof for the theorem on variance reduction by evaluating recombinations. First, we give an accessible direct proof in the case of two components that exhibits the key idea, then we extend to an arbitrary amount of components.

The generalization to \(m\) variables is given in the following theorem. The proof proceeds in the same way but requires the general Hoeffding decomposition involving marginals over all subsets of variables [57].

The bound in the theorem is tight, as evidenced by functions \(f\) which are linear in scalar random variables \(X_i\). In the other extreme, if \(f(X_1, \ldots, X_n) = \prod_i X_i\) and \(\mathbb{E}[X_i] = 0\) for all \(i\), then all non-diagonal covariance terms disappear and the left-hand side is \(\frac{1}{n^{2m-1}} \mathop{\mathrm{Var}}f\) while the right-hand side is \(\frac{1}{n} \mathop{\mathrm{Var}}f\).

Let \(X_1, \ldots X_n\) and \(Y_1, \ldots, Y_n\) be independent random variables such that \(X_i\) identically distributed and \(Y_i\) identically distributed. Let \(f\) be a deterministic function such that \(Z_i = f(X_i, Y_i)\) has finite mean \(\mu\) and finite variance \(\sigma^2\).

According to the central limit theorem, with \(S_n = \frac{Z_1 + \ldots + Z_n}{n}\) we have \[\sqrt{n} (S_n - \mu) \to \mathcal{N}(0, \sigma^2)\] in distribution.

Now consider the recombining estimator \[T_n = \frac{1}{n^2} \sum_{i,j \leq n} f(X_i, Y_j).\]

Clearly, \(\mathbb{E}T_n = \mu\). In theorem 2 we established \[\mathop{\mathrm{Var}}T_n \leq \mathop{\mathrm{Var}}S_n = \frac{1}{n} \sigma^2\] using the Hoeffding decomposition of \(f(X_i, Y_j)\).

More precisely, setting \(f_X = \mathbb{E}[f(X, Y) \mid X] - \mu\), \(f_Y = \mathbb{E}[f(X, Y) \mid Y] - \mu\), \(f_{XY} = f(X, Y) - f_X - f_Y - \mu\) and \(f'(X, Y) = \mu + f_X + f_Y + \frac{1}{\sqrt{n}} f_{XY}\), we derived (implicitly in the proof, and noting the change of notation for \(f_X\), \(f_Y\)) \[\mathop{\mathrm{Var}}T_n = \frac{1}{n} \mathop{\mathrm{Var}}f_X + \frac{1}{n} \mathop{\mathrm{Var}}f_Y + \frac{1}{n^2} \mathop{\mathrm{Var}}f_{XY} = \frac{1}{n} \mathop{\mathrm{Var}}f'.\] Now set \(S_n' = \frac{f'(X_1, Y_1) + \ldots + f'(X_n, Y_n)}{n}\), observe that \(\mathbb{E}f' = \mu\) and write \[\sigma'^2 = \mathop{\mathrm{Var}}f' = \mathop{\mathrm{Var}}f_X + \mathop{\mathrm{Var}}f_Y + \frac{1}{n} \mathop{\mathrm{Var}}f_{XY} \leq \sigma^2.\] According to the central limit theorem, \[\sqrt{n} (S'_n - \mu) \to \mathcal{N}(0, \sigma'^2)\]

We will now prove that ¹ \[\sqrt{n} \|T_n - S'_n \|_2 \to 0\] with \(n \to \infty\). Indeed, \[n \| T_n - S'_n \|_2^2 = \frac{1}{n^3} \sum_{i,j,k,l} \left(f(X_i, Y_j) - f'(X_i, Y_i)\right) \left(f(X_k, Y_l) - f'(X_k, Y_k) \right)\] Subtracting \(\mu\) from all \(f\)- and \(f'\)-terms, proceeding as in the proof of Prop. B.1 and observing that \(\mathop{\mathrm{Cov}}(f(X_i, Y_j), f'(X_i, Y_i)) = \mathop{\mathrm{Var}}f_X\), \(\mathop{\mathrm{Cov}}(f(X_i, Y_j), f'(X_j, Y_j)) = \mathop{\mathrm{Var}}f_Y\), \(\mathop{\mathrm{Cov}}(f(X_i, Y_i), f'(X_i, Y_i)) = \mathop{\mathrm{Var}}f'\), we get \[\begin{align} n \| T_n - S'_n \|_2^2 &= \frac{1}{n^3} \left( n^2 \mathop{\mathrm{Var}}f + n^3 \mathop{\mathrm{Var}}f' - n^2(n-1) \mathop{\mathrm{Var}}f_X - n^2(n-1) \mathop{\mathrm{Var}}f_Y \right) \\ &= \frac{1}{n^3} \left( n^2 \mathop{\mathrm{Var}}f + n^3 (\mathop{\mathrm{Var}}f_X + \mathop{\mathrm{Var}}f_Y + \frac{1}{n} \mathop{\mathrm{Var}}f_{XY}) - n^2(n-1) \mathop{\mathrm{Var}}f_X - n^2(n-1) \mathop{\mathrm{Var}}f_Y \right), \end{align}\] so all \(n^3\) terms within the brackets cancel and \[n \| T_n - S'_n \|_2^2 \to 0\] with \(n \to \infty\) as claimed.

Since \(\sqrt{n} (T_n - \mu) \to \sqrt{n} (S_n' - \mu)\) in \(L_2\) and \(\sqrt{n} (S_n' - \mu) \to \mathcal{N}(0, \sigma'^2)\) in distribution, \(\sqrt{n} (T_n - \mu) \to \mathcal{N}(0, \sigma'^2)\) in distribution, satisfying a CLT.

We now state the Delta Method theorem².

We have seen that in the setup of our paper with \(u\) a concave utility function implicitly defined by a reward aggregation both the recombination estimator \(T_n\) and the standard (diagonal) estimator \(S_n\) satisfy a CLT with variances \(\sigma'^2\) and \(\sigma^2\), respectively. Applying the Delta Method theorem, we can conclude that \(u(T_n)\) and \(u(S_n)\) likewise satisfy a CLT but with different variances \(\sigma'^2 u'(\mu)\) and \(\sigma^2 u'(\mu)\), respectively. Since \(\sigma'^2 = \sigma^2 - \frac{n-1}{n} \mathop{\mathrm{Var}}f_{XY}\), the recombination estimator is an improvement over the standard estimator, not just for \(u = \mathrm{id}\), but also in the presence of nonlinear utility functions (reward aggregation functions) \(u\) in the \(n \to \infty\) limit.

10 Information-Theoretic View on Inference Scaling↩︎

Consider a dataset \(D\) of problems \(x \in D\), an evaluation function \(v(x, y) \in \{0,1\}\), a probabilistic model \(\pi\) and consider the pass rate or coverage function \[c(k) = \mathbb{E}_{x \sim D, y_1, \ldots, y_k \sim \pi(x)}[\max_i v(x, y_i)].\] Empirical investigations [58]–[60] frequently report power law scaling of the logarithmic pass rate with respect to the number of samples of the form \[\log(c) \approx -ak^{-b},\] where \(a, b > 0\).

To analyze this relationship, note that \(k\) can be seen as an expected waiting time in the Bernoulli process with repeated samples from \(v(x, \pi(x)) \sim \mathrm{Bern}(p)\) where \(p = \frac{1}{k}\), and we can consider the information quantity \(h = -\log p = \log k\). Under the logarithmic power law model, we know the cumulative distribution function \[\mathop{\mathrm{\mathbb{P}}}(k \leq k_0) = c(k_0) = \exp \left( -ak^{-b} \right)\] and thus for \(h_0 = \log k_0\), \[\mathop{\mathrm{\mathbb{P}}}(h \leq h_0) = \exp \left( -ae^{-bh} \right).\] Differentiating, we get the probability density function \[p(h) = ab \exp(-bh -ae^{-bh}) = \frac{1}{\beta} \exp \left( -\frac{h-\mu}{\beta} - e^{-\frac{h-\mu}{\beta}} \right),\] i.e. a Gumbel distribution with mode parameter \(\mu = -\frac{\log a}{b}\) and scale parameter \(\beta = \frac{1}{b}\).

Note that the Gumbel distribution arises as an extreme value distribution over exponentially distributed samples, and we can interpret this by imagining several “latent dimensions of problem difficulty” \(h_i\), the largest of which determines the overall difficulty of a given problem: \(h = \max_i h_i\).

11 Properties of the Logmeanexp Objective↩︎

The \(\mathop{\mathrm{\mathrm{logmeanexp}}}_\beta\) function (Equation 1 ) has several theoretical appeals beyond serving as a smooth interpolation between the average (\(\beta \to \infty\)) and the maximum (\(\beta \to 0\)):

• For a well-chosen value of \(\beta\), it closely approximates the average between pass@1 + pass@10 + pass@100 + pass@1000. Such objectives strike a balance between favoring exploration and diversity (due to pass@1000 components) and re-ranking to favor successful solution attempts for exploiting the current knowledge (due to pass@1 components). See Appendix [app:transform] for details and Figure 3.

• It represents a soft form of optimism over the other actions. Recall that Deep-Q-Learning [61] uses \(R_{t} + \max_a Q(s_{t+1}, a)\) as the Q-value target for \(Q(s_t, a_t)\) if the action led to the new state \(s_{t+1}\). This target can be seen as an approximation of \(Q^\ast(s_t, a_t)\) for the \(Q^\ast\) the state-action values of the optimal policy \(Q^\ast\). In the language of soft reinforcement learning [21], \(\mathcal{V}_\beta = \mathop{\mathrm{\mathrm{logmeanexp}}}_\beta(r_1, \ldots, r_n)\) can be seen as the \(\beta\)-soft value function of selecting one out of \(n\) rollouts with uniform prior \(\pi_0\).

12 Reward Aggregation as Utility Functions↩︎

Consider as in Section 2.4 \(n\) rewards \(r_1, \ldots, r_n\). For zero-one rewards \(r_i \in \{0, 1\}\), and different choices of a multi-sample objective function \(f\), we write \(c = \sum_i r_i\) and factor \(f(r_1, \ldots, r_n) = u(p)\) via \(p = \frac{c}{n}\). Note that \(p\) is the mean estimator of \(\mathbb{E}[r]\) if \(r_i\) are samples from the distribution of \(p\).

12.0.0.1 Logmeanexp.

We compute \[f(r_1, \ldots, r_n) = \beta \log \frac{1}{n} \sum_i e^{r_i/\beta} = \beta \log \frac{1}{n} \left(ce^{1/\beta} + (n-c) \right) = \beta \log \left(pe^{1/\beta} + 1-p \right).\]

12.0.0.2 Pass@k with replacement.

In this case, \(f(r_1, \ldots, r_n)\) is already defined by \[f(r_1, \ldots, r_n) = 1 - \left(1 - \frac{c}{n} \right)^k = 1 - (1-p)^k.\]

12.0.0.3 Softmax-weighted average.

[9] suggest the multi-sample objective \(\sum_i \frac{e^{r_i/\beta}}{\sum_j e^{r_j/\beta}} r_i\). In this case, we obtain \[f(r_1, \ldots, r_n) = \sum_i \frac{e^{r_i/\beta}}{\sum_j e^{r_j/\beta}} r_i = \frac{ce^{1/\beta}}{c e^{1/\beta} + (n-c)} = \frac{pe^{1/\beta}}{p e^{1/\beta} + (1-p)} = 1 - \frac{1-p}{pe^{1/\beta}+1-p}.\]

Utility functions are a convenient method to understand the impact of a given choice of multi-sample objective function. We call a function \(u: [0, 1] \to [0, 1]\) a utility function if it is monotonically increasing. We call a utility function explorative if it is concave. Examples of explorative utility functions have been derived above, and are depicted in Figure 3.

For a single problem \(x\), we have \[\mathbb{E}_{r_i \sim \pi(x)}[f(r_1, \ldots, r_n)] = u(\mathbb{E}_{r\sim \pi(x)}[r]),\] and thus the same holds when taking expectations over \(x\). In other words, reinforcement learning with multi-sample objective functions changes the objective to a utility-transformed objective. On a single problem, by monotonicity of \(u\), optimal policies for the original objective are optimal for the transformed objective and vice versa. When averaging over a dataset of problems, however, utility functions induce a different allocation of optimization budget. Concretely, let \(dx\) be the distribution over problems and \(p_x = \mathbb{E}_{r\sim\pi(x)}[r]\). Then \[\mathbb{E}_x[p_x] = \int_x p_x dx,\] and \[\mathbb{E}_x[u(p_x)] = \int_x u(p_x) dx = \int_x p_x \frac{u(p_x)}{p_x} dx,\] so \(\frac{u(p_x)}{p_x} dx\) can be seen as the importance-reweighted distribution of problem difficulties according to \(u\). Explorative utility functions are therefore the utility functions that induce a monotonically decreasing reweighting on problem difficulties, as depicted in Figure 7 (a).

13 Variance Reduction with Leave-One-Out Baselines↩︎

Consider as in Appendix [app:transform] a multi-sample objective function \(f\) of zero-one rewards \(r_1, \ldots, r_n\) that factors via the mean \(p = \frac{1}{n} \sum_i r_i\) via a utility function \(u\): \[f(r_1, \ldots, r_n) = u(p(r_1, \ldots, r_n)).\]

Writing \(g_i = \nabla_\theta \log \pi_\theta(a_i)\), we are going to compare the default gradient estimator \[G_1 = \sum_i f(r_1, \ldots, r_n) g_i = \sum_i u(p) g_i\] with the baselined one \[\begin{align} G_2 &= \sum_i \left( f(r_1, \ldots, r_n) - f(r_1, \ldots, r_{i-1}, r_{i+1}, \ldots, r_n) \right) g_i \\ &= \sum_i \left(u(p) - u(p_{-i})\right) g_i, \end{align}\] where \[p_{-i} = \frac{1}{n-1} \sum_{j \neq i} r_j.\]

Because \(\mathbb{E}[u(p_{-i}) g_i] = 0\), we have \(\mathbb{E}[G_2] = \mathbb{E}[G_1] = \nabla_\theta \mathbb{E}[f(r_1, \ldots, r_n)]\) by the policy gradient theorem.

We will now compare the covariances of \(G_1\) and \(G_2\), respectively, assuming \(u \in C^2(\mathbb{R}_+, \mathbb{R}_+)\) for convenience and setting \(\mu = \mathbb{E}[r_1]\). We have \[\label{eq:p-update} p = p_{-i} + \frac{1}{n} (r_i - p_{-i})\tag{3}\] and \[u(p) = u(p_{-i}) + \frac{1}{n} u'(p_{-i}) (r_i-p_{-i}) + \mathcal{O}\left(\frac{1}{n^2}\right),\] hence \[\mathop{\mathrm{Cov}}G_2 \approx \mathop{\mathrm{Cov}}\left[ \sum_i \frac{1}{n} u'(p_{-i}) (r_i-p_{-i}) g_i \right] \approx \mathop{\mathrm{Cov}}\left[ \frac{u'(\mu)}{n} \sum_i (r_i-\mu) g_i \right] = \frac{u'(\mu)^2}{n} \mathop{\mathrm{Cov}}\left[ (r_1-\mu) g_1 \right] \in \mathcal{O}\left(\frac{1}{n}\right)\] by independence of \((r_i-\mu)g_i\) and \((r_j-\mu)g_j\) for \(i \neq j\).

For \(G_1\), on the other hand, we obtain, \[\mathop{\mathrm{Cov}}G_1 = \mathop{\mathrm{Cov}}\left[ \sum_i u(p) g_i \right] \approx \mathop{\mathrm{Cov}}\left[ \sum_i u(\mu) g_i \right] = n u(\mu)^2 \mathop{\mathrm{Cov}}\left[ g_1 \right] \in \mathcal{O}(n)\]

Since \(\mathbb{E}[G_1] = \mathbb{E}[G_2] = \nabla_\theta \mathbb{E}[u(\mu)]\) is independent of \(n\) in the limit, we obtain signal-to-noise ratios of \[\mathrm{SNR}(G_1) = \frac{\|\mathbb{E}[G_1]\|}{\mathop{\mathrm{Tr}}[\mathop{\mathrm{Cov}}G_1]} \in \mathcal{O}\left(\frac{1}{n}\right)\] and \[\mathrm{SNR}(G_2) = \frac{\|\mathbb{E}[G_2]\|}{\mathop{\mathrm{Tr}}[\mathop{\mathrm{Cov}}G_2]} \in \mathcal{O}(n),\] making \(G_1\) inconsistent. The estimator \(G_2\) with the leave-one-out baseline, on the other hand, is asymptotically deterministic.

Note that for simplicity, in the above, we considered a leave-one-out baseline that removes one out of \(n\) terms from the aggregation. In the case of DecompRL, we remove a fraction of \(\frac{1}{d}\) out of \(md\) terms in the case of decomposition actions and a fraction of \(\frac{1}{k}\) terms out of \(m\) terms in the case of implementation actions. Since Equation 3 holds with \(n\) replaced by \(d\) or \(k\), resp., and \(r_i\) replaced by the average over the removed terms, the same asymptotic analysis holds.

14 Additional Experiments on Multi-Policy Training↩︎

We perform experiments at the 8B scale using the LLama 3.1 8B Instruct model [11] to derive our multi-policy training recipe described in section 3. We train on the CodeContest training set and evaluate on its validation set. We give additional experiments regarding multi-policy training: benchmark results of combinations of separately trained policies (Figure 8), results on sequential training (Figure 7 (c)) and an overall comparison of multi-policy training methods over the course of reinforcement learning runs including alternate training (Figure 7 (b)). See Section 3.2 for additional discussion.

Analyzing the policy entropies for decomposition and implementation in separate training runs, we find that entropy decreases substantially over the course of training, hinting at better certainty about good actions (exploitation) but also harming diversity at large sampling scales (exploration). The conditioning of a fixed implementation policy on an entropy-decreasing decomposition policy leaves the implementation policy entropy unaffected (Figure 9).

15 Baselines↩︎

We compare DecompRL with four reinforcement learning algorithms:

• GRPO [15] - the state of the art for optimizing pass@\(1\) performance.

• analytical pass@\(k\) [20] - the state of the art for optimizing pass@\(k\).

• Soft Policy Optimization (SPO) [21] - optimized for diverse policies.

• our own \(\mathop{\mathrm{\mathrm{logmeanexp}}}\) utility function with regular inference which optimizes the mean between pass@1 and pass@k.

We train all baselines until the reward plateaus and tune the hyperparameters based on the best pass@10 performance on LiveCodeBench (see Figure 10). We set \(16\) samples per prompt for all baselines, \(k = 8\) for pass@\(k\) and \(\beta = 0.3\) in the logmeanexp. Since DecompRL has a higher sampling budget per problem its reward plateaus after \(1\) epoch whereas baselines plateau after \(3\) epochs on the same training set.

16 Hyperparameters and Additional Training Details↩︎

We train our models with the Adam optimizer [62] (\(\beta_1=0.9, \beta_2=0.95\)) with decoupled weight decay [63] factor \(0.1\) and gradient clipping at a norm of 1. We warm up the learning rate linearly over 200 steps to a final value of \(6 \times 10^{-8}\). We use a local batch size of 2. Trainers discard the first 10 batches at the beginning of a run as warmup.

Our distributed runs use 80 H100 GPUs, split up into 8 trainer GPUs and 72 worker GPUs. On the worker side, we use temperature \(1.0\) for generations. During decomposition training, we set the number of implementations \(k = 8\), a maximum decomposition size of \(n_{\mathrm{max}} = 6\) and conduct up to \(m = 1024\) evaluations per decomposition. We run three separate decomposition trainings with \(\beta = 0.03, 0.07, 0.1\) with Qwen 2.5 7B which according to Figure 3 corresponds to a mixture of pass@k objectives for \(1 \leq k \leq 1000\). We pick \(\beta = 0.07\) as it had the highest pass@\(10\) performance on the DeepMind Code Contest validation set. Each RL training run takes 2–3 days on 80 H100 GPUs for Qwen 2.5 7B. Each evaluation run (up to 1000 decompositions per problem with 4096 evaluations each) takes approximately 20 minutes with tensor parallelism TP=4 on 32 GPUs for Qwen 2.5 7B.

At evaluation time we use temperature of 1.0 and up to 1000 decompositions per problem with each up to \(4096\) evaluations.

16.0.0.1 Memory cost.

DecompRL maintains two copies of the model (decomposition and implementation policies), but only one is trained at a time. For Qwen 2.5 7B, model weights occupy less than 10% of GPU memory on workers (\(\sim\)​7 GiB for two model copies vs.\(\sim\)​3.5 GiB per GPU for a single model). Only worker GPUs bear the extra memory cost. We observed zero out-of-memory errors and required no CPU offloading, even with the Code World Model 32B. We argue the cost of storing two policies is similar to the one of using a reference model to compute a KL penalty in standard RL.

16.0.0.2 Code execution.

For all our code evaluations, we use a sandbox similar to the one provided for [13] using Python 3.11. We set a time limit of \(2T+1\) seconds for all test cases if \(T\) is the number of test cases. We execute code on an external CPU cluster that parallelizes unit test executions. Each problem in our training set has a mean of \(10\) unit tests, each with average size \(<10^2\) bytes which we estimated to take \(<0.1\) seconds to execute in our evaluation setup. If we have \(512\) evaluations per problem during training, this gives an upper bound of \(512 \times 0.1 = 51.2\) seconds per problem.

17 Prompts↩︎

We use language models as decomposition and implementation policies.

For decompositions, we use the following decomposition prompt with a one-shot example, where the {context} is the dataset’s problem description:

After the decomposition step, we extract the function headers and documentation strings from the language model’s generation. We implement each function in the decomposition using the prompt below. We have {decomposition} extracted from our first language model call, {current_name} the function being implemented, and {other_names} other function descriptions in the decomposition.

18 Reward Tensor Structure↩︎

We investigate the structure of reward tensors. For this, we generate solutions to CodeContests validation problems using hierarchical inference with \(d=2\) decompositions of a maximal size of \(6\) and \(k=4\) implementations per function. We evaluate all resulting combinations, i.e. up to \(4^6 = 4096\) per problem, and analyze the resulting reward tensors \(r\) of mode (number of axes) \(n\) and shape \(k \times \ldots \times k\). Assume that each implementation is either correct or incorrect, at an individual level, and that combinations pass the tests if and only if all their constituting implementations are correct. In this case, correctness of implementations can be expressed as the maxima \(m_i\) across all modes of \(r\) except the \(i\)-th, and the reward tensor would equal the outer product \(r^{(1)} = m_1 \otimes \ldots \otimes m_n\) of the \(m_i\), which is a rank-1 approximation of \(r\). More generally, the rank of a Boolean tensor using the Boolean operations (\(\vee\) for addition, and \(\wedge\) for multiplication) is known as the Boolean rank [64]. In the case of \(n=2\) it is also known as the bipartite dimension of the corresponding bipartite graph. It is NP-hard to compute [65].

From 20 decompositions with at least one positive reward, we obtain the following statistics: In 17 cases, \(r = r^{(1)}\) agrees with the rank-1 approximation. Overall, viewing \(r^{(1)}\) as a prediction of \(r\), there are 859 true positives, 483 true negatives, 34 false positives and no false negatives (the latter by definition).

19 Number of Tokens Generated↩︎

Figure 1, 4 (a), and Table 1 compare RL-trained models (SPO, GRPO, pass@8, logmeanexp) against our hierarchical inference method (DecompRL) for different token counts. We consider one attempt as one code generation per problem for standard inference and one full decomposition with 8 implementations, 4096 max evaluations for hierarchical inference. To ensure a fair comparison, we compare the pass rates per generated tokens which correspond to a different number of attempts for each RL algorithm. Table 3 summarizes the average cost of an attempt per method for the LLama 3.1 8B and Qwen 2.5 7B models.

We compare success rates for a fixed budget of tokens similar to the pass@t from [66].

20 Example DecompRL↩︎

We show below an example of a code problem extracted from LiveCodeBench and a generated decomposition with \(6\) functions from a Qwen 2.5 7B model trained with DecompRL. This decomposition leads to correct code solutions.

21 Hierarchical Inference Without RL↩︎

22 Licenses of Used Datasets↩︎

We use datasets with the following licenses:

• CodeContests [13]: CC BY 4.0 with MIT and Apache 2.0 licenses components.

• TACO [14]: Apache 2.0.

• LiveCodeBench [1]: CC BY 4.0 with MIT and Apache 2.0 licenses components.

References↩︎