Generalizing Test-time Compute-optimal Scaling as an Optimizable Graph

• Existing TTS architectures are usually static across tasks and use one model, although MATH and MMLU show different topology and model-combination preferences.MATH favors hybrid structures and mixtures of smaller models, while MMLU performs better with parallel scaling and a single 8B model.Source: §1 · p. 1
• The search jointly chooses architecture and model assignments under a compute budget, but its combinatorial graph space makes exhaustive evaluation infeasible.The paper reports 10^18 to 10^26 possible graphs at 12 nodes, depending on model diversity.Source: §1 · p. 2
• Agent-REINFORCE replaces REINFORCE’s sample–gradient–update loop with sample–feedback–update, using textual feedback and empirical insights to update a probabilistic graph distribution.An Agent initializes candidates, an Environment evaluates them, and an Archive records results for subsequent updates.Source: §1 · p. 2
• Agent-REINFORCE outperforms traditional and LLM-based baselines in search efficiency and accuracy, and identifies graphs under joint accuracy–latency objectives.Source: §1 · p. 2

• Existing TTS paradigms include iterative refinement, candidate aggregation, and hybrids, but hybrids typically use fixed trees and a single model.Compute-optimal TTS and heterogeneous ensembles are discussed as relevant directions that remain underexplored together.Source: §2 · p. 3
• LLM optimization methods use prior knowledge for initialization and textual guidance or exploration, while this work applies LLMs to probabilistic graph initialization and updates for compute-optimal TTS.The passage distinguishes black-box candidate refinement from approaches interleaving LLM guidance with gradient-based optimization.Source: §2 · p. 3

• Parallel scaling samples multiple outputs and fuses them, while sequential scaling repeatedly refines an output; hybrid methods combine these primitives.The primitives are repeated sampling, fusion, and self-refinement.Source: §3 · p. 3
• A collaboration graph is a DAG with assistant nodes that refine predecessor outputs, fuser nodes that aggregate them, and directed edges that carry information toward a sink prediction.Inference executes nodes in topological order according to their roles and assigned models.Source: §3 · pp. 3–4
• The budget is normalized against the cost of a single-node inference using the smallest model, so a graph’s budget expresses equivalent smallest-model inference units.Source: §3 · p. 4
• The optimization objective is to find a task-specific graph that maximizes training-set performance while satisfying the compute budget; the selected graph is evaluated on test data.Feasible graphs are defined by the budget constraint.Source: §3 · p. 4

• The search is challenging because graph choices grow combinatorially with budget and task requirements call for tailored designs, motivating pilot analysis of cross-task TTS patterns.Source: §4 · p. 4
• Replicating the strongest model family is generally more effective than mixing families; reasoning tasks favor smaller-model ensembles, while knowledge tasks prefer larger models.On MMLU in the 3B space, LLaMA×2 yields higher accuracy than LLaMA+Gemma or Gemma×2.Source: §4 · p. 5, Figure 3 · p. 5
• Both parallel width and sequential depth have task-dependent optima, after which performance plateaus or declines; MATH peaks at 8 parallel or 8 sequential nodes.The passage associates diminishing width gains with long-context limits and depth-related degradation with propagated errors.Source: §4 · p. 5, Figure 4 · p. 5
• Width and depth interact: on MATH with LLaMA-1B and wd ≤24, increasing width from 1 to 3 shifts the optimal depth from 8 to 4.Deeper refinement also shifts the optimal width forward.Source: §4 · p. 5, Figure 4 · p. 5

5 THE PROPOSED FRAMEWORK – AGENT-REINFORCE

• Agent-REINFORCE replaces REINFORCE’s sample–gradient–update loop with agent-led sampling, textual feedback, and probabilistic-graph updates under a fixed budget.The textual feedback serves as the gradient, while task-specific model preferences and width–depth interactions guide search.Source: §5 · p. 6

5.1 PROBABILISTIC GRAPH OPTIMIZATION PROBLEM

• The search distribution jointly represents edge presence, node-role selection, and model assignment probabilities for sampled collaboration graphs.Its parameters are θ for edges, π for roles, and ψ for model choices.Source: §5.1 · p. 6

5.2 AGENT-REINFORCE

• The agent initializes model assignments using task-specific family, size, and instance preferences, then uses trial feedback to select an initial configuration.Subsequent sampling draws edges and roles from the probabilistic graph while retaining node model assignments.Source: §5.2 · p. 7
• The update process uses task-specific width–depth optima and their fixed-budget trade-offs to steer later graph search.The agent consults current and past trial feedback when adjusting the graph toward a width–depth balance.Source: §5.2 · p. 7
• The Environment executes candidate graphs and returns performance feedback, while the Archive retains graph, trial, and feedback histories for later updates.Candidates exceeding the budget are brought within it by removing the smallest nodes and making the corresponding graph update.Source: §5.2 · pp. 7–8

Experimental Setup

• Experiments cover MATH, MMLU, and HumanEval with LLaMA and Gemma models, comparing traditional, gradient-based, and LLM-based search baselines.Methods run for up to 30 search iterations, with validation-based stopping and final evaluation on the test set.Source: §6 · p. 8

Main Results

• Agent-REINFORCE achieves the highest average test-set score and converges substantially faster than the compared methods.The authors attribute this to insight-guided search; they also report that TextGrad has the highest inference latency among the methods.Source: §6 · p. 8

Ablation Study

• Removing Insights 2 or 3 reduces accuracy, and removing role setting degrades MATH performance; MMLU remains stable without Insight 3 or role settings.The authors connect the MMLU pattern to its preference for larger models with fewer nodes.Source: Table 2 · p. 8, §6 · p. 9

Performance Under Various Budget Settings

• At MATH FLOPs budgets of 18 and 42, Agent-REINFORCE consistently delivers superior efficiency and accuracy.The authors note that some baselines do better at a smaller budget than at 80, consistent with task-specific compute optima.Source: §6 · p. 9, Table 3 · p. 9

Latency-aware Joint Optimization Objective

• On MATH with an 18 FLOPs budget, joint optimization yields 5.6 seconds average latency per test query versus 16.7 seconds under performance-only optimization, while performance falls from 47 to 44.This experiment uses multidimensional feedback to balance accuracy and efficiency.Source: §6 · p. 9, Table 3 · p. 9

Generalization to the Dollar Cost as Budget

• With a fixed API budget of $5E-4 per query, Agent-REINFORCE excels in accuracy and efficiency when searching among 4×8B to 6×8B models.The authors present this as generalization from FLOPs budgets to dollar-cost budgets.Source: §6 · p. 9, Table 3 · p. 9

Visualization

• The MATH graph under budget 80 favors small-model ensembles and hybrid scaling biased toward sequential refinement, with width 3 and depth 4.The authors relate the structure to the task’s relatively low difficulty and the benefit of iterative self-refinement for multi-step math reasoning.Source: §6 · p. 9, Figure 6 · p. 9

7 CONCLUSION

• The paper studies task-specific, compute-optimal test-time scaling over multi-LLM collaboration graphs under a fixed budget.The design space is described as exponentially large in model choices and nodes.Source: §7 · p. 9
• Pilot analysis identifies task-specific model-family preferences, width and depth optima, and interaction between width and depth.Small-model ensembles are favored when their incremental gains are high, while additional compute beyond width or depth optima can degrade performance.Source: §7 · p. 10
• Agent-REINFORCE uses budget-aware, feedback-driven graph search and outperforms traditional and LLM-based baselines in search efficiency and performance.The conclusion also reports that it can find optimal graphs under a joint performance–latency objective.Source: §7 · p. 10

A.1 TEST-TIME SCALING: MODES AND BUILDING BLOCKS

• Parallel scaling repeatedly samples outputs for aggregation, while sequential scaling iteratively self-refines an output.Fixed hybrids combine the two; dynamic scaling searches architectures and assigns heterogeneous models under a compute budget.Source: §A.1 · p. 16, Figure 7 · p. 16
• Dynamic TTS is represented as a collaboration graph with role-assigned nodes, directed information flow, and a terminal aggregator.The section names repeated sampling, fusion, and self-refinement as the graph’s building blocks.Source: §A.1 · p. 16, Figure 8 · p. 16

• The graph runs in topological order, with successors of the input node generating the initial outputs.Source: §A.2 · p. 16
• Nodes activate when their in-degree reaches zero, then fusers aggregate or assistants refine outputs and propagate results forward.Source: §A.2 · p. 16
• The unique sink node produces the final answer.Source: §A.2 · p. 16

• MATH achieves its best accuracy with a hybrid graph topology and a mixture of 3B models.Source: §A.3 · p. 16, Table 4 · p. 17
• MMLU favors pure parallel graph topologies and a single 8B model.Source: §A.3 · p. 16, Table 4 · p. 17
• Table 4 organizes accuracy by dataset, model combination, and topology, with results averaged over 10 random graphs.Source: Table 4 · p. 17, §A.3 · p. 17

• The passage distinguishes non-isomorphic DAG counts from labeled DAG counts according to whether node identities are ignored or retained.Source: §A.4 · p. 17

Indistinguishable nodes

• For indistinguishable nodes, fixing a topological order permits only forward edges, yielding 2^(n(n−1)/2) candidate adjacency matrices.Many candidates are isomorphic, so canonical labeling merges equivalent graphs; nodes are grouped by in-degree and out-degree to reduce the permutations considered.Source: §A.4 · p. 17

Distinguishable nodes

• For distinguishable nodes, the recurrence A(0)=1 and A(n)=sum over k of (−1)^(k+1) times the binomial coefficient and 2^(k(n−k))A(n−k) counts labeled DAGs.A(n) denotes the number of labeled DAGs on n nodes; the recurrence accounts for edge configurations while counting only acyclic structures.Source: §A.4 · p. 17

Indistinguishable nodes

• The canonical-form implementation counts non-isomorphic DAGs, reaching 1,259,209 structures at n=7.The reported counts are 1, 2, 8, 54, 762, 21,542, and 1,259,209 for n=1 through n=7, respectively.Source: §A.4 · pp. 17–18

A.5 TASKS, DATASETS, AND MODELS

• MATH, MMLU, and HumanEval evaluate mathematical reasoning, broad knowledge and reasoning, and code generation, using final-answer accuracy, precision, and passk, respectively.The experiments use 750 MATH and 285 MMLU training examples with 100 test examples each, and 128 HumanEval training instances with the remainder for testing.Source: §A.5 · p. 19

A.6 DETAILED INSIGHTS

• Model-family and size preferences vary by task, parallel and sequential scaling have task-dependent optima, and graph width and depth shift each other’s optima.MATH favors 3B×3 while MMLU favors 8B×1: on MATH, LLaMA 3B rises from 39% to 46% with two instances versus 49% for one 8B; on MMLU, it rises from 41% to 45% versus 64% for 8B×1. On MATH, performance peaks at 8 parallel or sequential nodes; under wd ≤24, optimal depth falls from 8 at width 1 to 4 at width 3, while optimal-depth accuracy is 38 at width 1, 47 at width 3, and 45 at width 4.Source: §A.6 · pp. 19–20

A.7 fCOST(G, T) WITH THE FLOPS COMPUTE METRIC

• The FLOPs cost model combines token-wise projection/MLP costs with attention costs based on prefill and decode lengths, then sums node costs across the graph.The scheme counts a multiply-add as 2 FLOPs and accounts for cached keys and values; monetary-cost surrogates can replace FLOPs-based node terms while retaining graph aggregation.Source: §A.7 · pp. 20–21

A.8 DETAILED BUDGET DEFINITION

• The normalized graph budget divides its task-specific cost by the cost of one inference using the pool’s smallest model, making different model sizes and topologies comparable.Cost may represent FLOPs, wall-clock runtime, or monetary cost; this work adopts FLOPs as its primary metric.Source: §A.8 · pp. 21–22

A.9 DETAILED OPTIMIZATION WITH JOINT OBJECTIVE

• For joint accuracy-and-latency optimization, Agent-REINFORCE gives the agent candidate inference times and stores feedback in an archive to assess latency-performance trade-offs.Instructions can also provide prior knowledge about graph structures and latency, such as latency’s sensitivity to node count and graph width.Source: §A.9 · p. 22

A.10 DETAILED DOLLAR COST-BASED BUDGET

• The dollar-cost budget uses per-million-token inference costs from Together AI and Compare Ai Models, with estimated reference values for three unquoted models.The reference values are 0.02 for LLaMA-3.2 1B, 0.02 for Gemma-3 1B, and 0.06 for Gemma-1.1 2B.Source: §A.10 · p. 22, Table 5 · p. 22

A.11 DETAILED REINFORCE ALGORITHM

• REINFORCE represents edge, role, and model choices as graph-distribution probabilities, then iterates candidate sampling, utility evaluation, and parameter updates under the budget.After optimization, maximum-a-posteriori decoding produces a deterministic graph, which is pruned if needed to meet the budget.Source: §A.11 · pp. 23–24

A.12 PROMPT DESIGN IN AGENT-REINFORCE

• Structured prompts initialize model-family and size preferences and instance counts, then use task context, distilled insights, constraints, and feedback to refine graph connectivity and budget allocation.Graph updates draw on the saturation and width-depth interdependence insights to improve the candidate structure.Source: §A.12 · p. 24
• Initialization prompts select model family, size, and instance counts for each task, reflecting the reported task-specific preferences under budget constraints.The requested output is a JSON dictionary of model combinations; the graph is later optimized into a DAG, with each edge passing one agent’s output into the next agent’s context.Source: §A.12 · p. 25

TASK CONTEXT

• Each graph-update iteration receives the previous graph, batch accuracy, and edge probabilities, then returns a next-round DAG and updated probability table.The prompt describes outward expansion from FinalDecision and probability increases only for edges shown useful across sampling rounds.Source: HISTORICAL SNAPSHOT · p. 26, TASK CONTEXT · p. 26, WHAT TO RETURN · p. 26

OPTIMIZATION RULES

• Updates preserve the model-count budget, assign assistant or fuser roles, and raise only sampled useful edges while avoiding cycles, excessive in-degree, and deep chains.Expansion starts at FinalDecision and proceeds through parents; task preferences also guide width and depth, whose optimal balance is interdependent.Source: OPTIMIZATION RULES · p. 26, DATA AND INSIGHT · p. 26

A.13 BASELINES

• The adapted baselines span LLM-guided TextGrad and MaaO, REINFORCE-based GPTSwarm, Bayesian Optimization, and unguided Random Search.TextGrad uses textual graph refinements; MaaO alternates REINFORCE with LLM updates; GPTSwarm uses gradient updates alone; BO uses a surrogate and acquisition function; Random Search samples budget-feasible graphs without history or prior guidance.Source: §A.13 · p. 27

A.14 CONVERGENCE AND EFFICIENCY ON MATH DATASET

• Agent-REINFORCE achieves the best accuracy and fastest convergence on MATH, while TextGrad slows from budget overuse and GPTSwarm/MaaO stall in local optima.Figure 9 shows validation-accuracy trajectories over 20 iterations.Source: Figure 9 · p. 28, §A.14 · p. 28

A.15 DETAILED RELATED WORK

• Prior test-time scaling combines sequential refinement, parallel candidates, and hybrid strategies, but fixed structures and underexplored multi-model sampling motivate graph-based, heterogeneous-model search.The related-work discussion notes sequential scaling’s risk of corrupting prior reasoning and parallel scaling’s limited depth for complex tasks.Source: §A.15 · pp. 28–29
• LLMs have been used as black-box optimizers or alongside gradient methods; this work applies their task-aware initialization and textual updates to gradient-available probabilistic graph optimization.The discussion frames this combination as retaining gradient methods’ data efficiency while using LLM guidance for initialization and semantic parameter updates.Source: §A.15 · pp. 29–30