Train for Many, Update with One: Improving Multi-Sample Reasoning via Fenchel-Bregman Posterior Learning

• Unlike related methods that train a single model for multi-sample inference, this work learns a distribution over solver parameters for fixed-budget candidate coverage.Source: §1 · p. 1
• Finite-K posterior learning trains qϕ to minimize the probability all K candidates fail, weighting problems more when other candidates are unlikely to solve them.Source: §1 · p. 1
• Expected per-solver loss treats draws individually, whereas multi-sample success depends on at least one of K draws solving the problem.Source: §1 · p. 1
• Direct-K is unbiased but requires K independent posterior draws per update.Source: §1 · p. 2
• FB combines an exact Fenchel reformulation with closed-form Bregman weight tracking, enabling one posterior draw per update; its IVON instantiation is FBI.Source: §1 · p. 2
• FBI improves candidate coverage across recursive and autoregressive reasoning, with substantially lower training computation than Direct-K in a matched comparison.Source: §1 · p. 2

2.1 MULTI-SAMPLE REASONING WITH POSTERIOR DRAWS

• For a fixed candidate budget, the learned distribution qϕ over solver parameters generates the pool through K independent draws, each producing one candidate.Source: §2.1 · p. 3

2.2 FINITE-K POSTERIOR LEARNING

• Finite-K posterior learning minimizes the probability that all K candidates fail, aligning training with pool-level coverage rather than individual-draw loss.Source: §2.2 · p. 3
• The objective also adds KL regularization toward reference prior p0 to discourage the parameter distribution from concentrating at a point estimate.Source: §2.2 · p. 3

2.3 UNDERSTANDING AND OPTIMIZING THE FINITE-K OBJECTIVE

• For K ≥2, a problem receives more gradient weight when the other K−1 candidates are likely to fail; at K = 1, this reweighting disappears.Source: §2.3 · p. 4
• Direct-K gives an unbiased estimate of the all-fail data term but requires K independent posterior draws per update.Source: §2.3 · p. 4

3.1 FENCHEL REFORMULATION AND ONE-DRAW POSTERIOR UPDATE

• Fenchel conjugacy rewrites Fbarϕ(x)^K as an optimization over problem-specific weights λx, separating the weights from posterior parameters for fixed-weight updates.Source: §3.1 · p. 4
• Under standard conditions, one posterior draw gives an unbiased fixed-weight gradient estimate that recovers the finite-K data gradient at the optimal weight.Source: §3.1 · p. 5

3.2 BREGMAN UPDATE FOR THE DUAL WEIGHT

• A closed-form Bregman update tracks each dual weight through an exponential running average of observed failures, without an additional posterior draw.Source: §3.2 · p. 5

3.3 POSTERIOR UPDATES WITH IVON

• FBI instantiates posterior updates with IVON, which learns a diagonal Gaussian posterior over solver parameters.Source: §3 · p. 4, §3.3 · p. 6
• The same failure observations update the posterior and prepare the next weight; when the bounded score differs from binary failure, the method uses a score-based surrogate objective.Source: §3.3 · pp. 5–6

4 EXPERIMENTS

• The evaluation covers four recursive benchmarks and five mathematical-reasoning benchmarks, using TRM in the main recursive comparison and Qwen3-4B-Base for autoregressive reasoning.Matched runs share the backbone, posterior initialization, training data, update budget, and evaluation budget; the study also compares recursive backbones and Bayesian methods and analyzes training targets, weighting, recovery, and cost.Source: §4 · p. 6, §4.1 · p. 6, §4.2 · p. 7

4.1 RECURSIVE REASONING MODELS

• On four recursive benchmarks, FBI achieves or ties the highest value in every metric column and improves candidate coverage and selected accuracy on Sudoku-Extreme and Maze-Hard.On Maze-Hard, Pass@10 rises from W-PTRM’s 89.80% to 91.30%; FBI also reaches the Pass@10 of PTRM, W-PTRM, and Direct-10 using 80%, 60%, and 20% fewer candidates, respectively.Source: §4.1 · pp. 6–7, Table 1 · p. 7, Figure 2 · p. 7

4.2 AUTOREGRESSIVE LANGUAGE MODELS

• FBI achieves the highest average Pass@8 across five mathematical benchmarks at 60.44%, 0.84 percentage points above DMPO.It overtakes the three compared baselines at K = 8 and reaches 66.06% at K = 16, 1.90 percentage points above Pass@K Training.Source: Table 2 · p. 8, §4.2 · p. 8

4.3 PERFORMANCE ACROSS BACKBONES AND BAYESIAN METHODS

• FB-based methods outperform corresponding baselines across evaluated recursive backbones and variational optimizers, with FBE leading in candidate coverage and FBI in selected accuracy.With FPRM, FBI reaches 94.80% Pass@10 versus W-FPRM’s 88.70%; on TRM, FBI improves Pass@10 from IVON’s 90.8% to 91.3%, while FBE improves it from EVON’s 91.1% to 92.7%.Source: Figure 3 · p. 8, §4.3 · pp. 8–9

4.4 ANALYSIS OF FINITE-K POSTERIOR LEARNING

• Matching Ktrain to the evaluation budget K gives the highest reported Pass@K at every tested budget on Maze-Hard.At K = 2, Ktrain = 2 ties with Ktrain = 5 at 89.1%; raising Ktrain from 1 to 10 increases Pass@10 from 90.8% to 91.3% while Pass@1 changes from 87.8% to 87.7%.Source: Figure 4 · p. 9, §4.4 · p. 9
• At Ktrain = 10, FBI achieves comparable mean Maze-Hard Pass@10 to Direct-10 using 12.3% of its measured training time per update.FBI uses 10.4% of Direct-10’s FLOPs and 0.611 GB peak GPU memory versus 3.707 GB; it uses one posterior draw per update rather than ten.Source: Figure 4 · p. 9, §4.4 · p. 10
• Replacing problem-specific weights with the minibatch mean lowers coverage at every tested budget, including Pass@10 from 91.3% to 90.7%.The reported gain grows from 0.2 percentage points at K = 1 to 0.6 points at K = 5 and 10.Source: §4.4 · pp. 9–10
• After first-candidate failures on one fixed Maze-Hard seed, FBI recovers an exact solution in 32.5% of cases (40/123), the highest rate among compared methods.Its repeated-output rate is 56.9%, versus 64.0–87.8% for baselines; the weight ablation does not isolate the effect of weighting on conditional recovery.Source: Figure 5 · p. 10, §4.4 · p. 10

5 CONCLUSION

• FBI combines an exact reformulation with Bregman tracking for one-draw finite-K posterior learning.The conclusion reports improved candidate coverage across recursive and autoregressive reasoning.Source: §5 · p. 10
• On matched Maze-Hard, FBI reaches comparable Pass@10 to Direct-10 using 12.3% of its measured training time per update.On Qwen3-4B-Base, it attains the highest mean Pass@8 across five mathematical-reasoning benchmarks; FB gains also extend across backbones and to a structured Gaussian posterior.Source: §5 · p. 10
• Surrogate mismatch, stochastic gradient noise, or posterior drift may weaken failure tracking; broader tasks and posterior families could extend applicability.The conclusion suggests improving surrogate calibration where applicable, using predictive or variance-reduced tracking, and evaluating adaptive or correlated candidate pools and selector-aware training.Source: §5 · p. 10

AI USE STATEMENT

• OpenAI GPT and Anthropic Claude were used to improve manuscript clarity and readability.The authors also report using them to refine research ideas and assist with mathematical derivations and proofs.Source: AI USE STATEMENT · p. 10
• The authors reviewed all AI-assisted material and checked mathematical arguments step by step, including assumptions and derivations.Source: AI USE STATEMENT · p. 10
• The authors take full responsibility for the final content, including AI-assisted text and mathematical claims.Source: AI USE STATEMENT · p. 10

B EXPERIMENTAL DETAILS

• The appendix supplies derivations for the finite-K objective and one-draw Fenchel–Bregman update, including binary estimation, Fenchel weights, fixed-weight gradients, and Bregman tracking.The experimental-details section describes evaluation protocols, baseline controls, posterior-training settings, autoregressive evaluation, and training-cost measurements.Source: A DERIVATIONS FOR FINITE-K POSTERIOR… · p. 14, B EXPERIMENTAL DETAILS · p. 15

A.1 BINARY ALL-FAIL PROBABILITY AND DIRECT MONTE CARLO ESTIMATION

• Under independent draws, the product of K binary failure indicators is an unbiased estimator of the all-fail probability, whose complement is candidate coverage.Without independence, its expectation need not factor into the Kth power of marginal failure probability; reusing one draw does not recover the same estimator.Source: §A.1 · p. 14

A.2 DERIVATION OF THE FENCHEL REPRESENTATION

• For K ≥ 2, the Fenchel representation has a unique maximizing weight; for K = 1, the data term is linear and needs no Fenchel step.The derivation establishes strict convexity of ψK and identifies the maximizer through the inverse derivative.Source: §A.2 · p. 14

A.3 ONE-DRAW GRADIENT WITH A FIXED WEIGHT

• A fixed weight chosen before the posterior draw makes the one-draw gradient estimator unbiased for the corresponding fixed-weight data gradient.It estimates the finite-K data gradient only when the weight is optimal; the KL term remains unchanged in the fixed-weight posterior subproblem.Source: §A.3 · pp. 14–15

A.4 BREGMAN UPDATE FOR THE DUAL WEIGHT

• The Bregman-proximal update yields the exponential running average of observed failures, so the EMA is a closed-form update rather than an independently chosen heuristic.The update has a unique solution, with β = 1/(1 + η).Source: §A.4 · p. 15

B.1 RECURSIVE EVALUATION PROTOCOLS

• Unless stated otherwise, evaluation uses K = 10, with smaller budgets taken as ordered prefixes; validation selects several training and posterior settings.Posterior methods independently sample parameters, and stochastic solvers also resample internal randomness.Source: §B.1 · p. 16
• Recursive protocols evaluate Sudoku-Extreme at depth 64, Maze-Hard at depth 16, and ARC-AGI at depth 16 with task-specific selection or aggregation rules.ARC-AGI evaluates 25 fixed posterior draws per augmentation and computes Pass@K over ranked predictions after aggregation.Source: §B.1 · p. 16

B.2 BASELINE DETAILS

• PTRM samples latent trajectories, while W-PTRM and W-FPRM independently perturb checkpoint weights; Table 3 describes fixed-Gaussian controls with sampled-scalar variance σ^2.The controls use task-specific scales, selected on validation for Sudoku-Extreme and Maze-Hard; ARC-AGI-1 uses a fixed multiplier.Source: §B.2 · p. 16, Table 3 · p. 17
• FBI exceeds GRAM on both selected accuracy and Pass@10 across Sudoku-Extreme and Maze-Hard.The table reports percentages; selected accuracy is selector-returned exact-solution accuracy, while Pass@10 is exact-solution coverage among ten candidates.Source: Table 4 · p. 17

B.3 RECURSIVE POSTERIOR TRAINING

• Direct-10 uses FBI’s failure score and IVON posterior machinery but estimates the Ktrain = 10 data term with ten independent posterior draws per update.The posterior mean and diagonal precision are fine-tuned from a task checkpoint.Source: §B.3 · p. 17
• The score-based objective matches the exact all-fail target only when the posterior mean failure score equals the one-candidate failure probability; otherwise it is a surrogate.FBI constructs its bounded score as f = 1 − exp(−ℓ/T) from the solver loss before auxiliary selector losses are added.Source: §B.3 · p. 17
• The recursive implementation fixes and detaches the tracker weight before drawing parameters, averages repeated group-key failures before one tracker update, and leaves absent groups unchanged.The Fenchel penalty is omitted at fixed weight; the auxiliary Q loss contributes to gradient and curvature estimates but is not part of the failure score or multiplied by the weight.Source: §B.3 · pp. 17–18
• TRM-based FBI uses Ktrain = 10 and one posterior draw per update; the appendix also reports ARC-AGI and FPRM training configurations.TRM configurations use failure-score temperature 1.0; ARC-AGI uses a diagonal IVON posterior over 6,829,058 dense parameters, and the FPRM Maze-Hard run spans 2,000 updates.Source: Table 5 · p. 18, Table 6 · p. 18, §B.3 · p. 18

B.4 AUTOREGRESSIVE EXPERIMENTS

• The autoregressive evaluation generates sixteen candidates per problem, scores binary correctness with Math-Verify, and estimates Pass@8 from those outcomes.Candidates use temperature 0.6 and top-p 0.95; the reported aggregate averages benchmark scores without weighting.Source: §B.4 · p. 19
• Mathematical-reasoning coverage is reported as unweighted mean Pass@K across five benchmarks under a common protocol with sixteen candidates per problem.The benchmarks are MATH500, OlymMATH, AMC23, AIME24, and AIME25; Table 8 reports K values 1, 4, 8, and 16.Source: §B.4 · p. 19, Table 8 · p. 19
• FBI trains a diagonal IVON posterior over Qwen3-4B-Base parameters for 300 steps, using one parameter draw shared by eight responses per prompt and RLOO policy-gradient estimates.Evaluation independently samples 16 parameter vectors and generates one response per problem from each.Source: §B.4 · p. 19

B.5 BAYESIAN METHOD COMPARISON

• In the Bayesian comparison, IVON and EVON use single-candidate training, while FBI and FBE target Ktrain = 10 with one posterior draw per update.IVON and FBI use diagonal Gaussian posteriors; EVON and FBE use Gaussian posteriors with diagonal covariance in a Kronecker-factored preconditioner’s eigenbasis.Source: §B.5 · p. 19

B.6 TRAINING COMPUTE ACCOUNTING

• The parameter-space resource comparison matches backbone, checkpoint, data, seed, batch size, precision, IVON schedule, and GPU, changing the estimator and posterior draws.FBI uses one posterior draw per update and Direct-10 uses ten; both use Ktrain = 10.Source: §B.6 · p. 20
• The timed loop includes batch retrieval, optimizer updates, and EMA; FLOPs profile matrix multiplication and fused attention, while peak memory is measured across 30 updates.The FLOP counter omits some elementwise operations, and ordinary and curvature-anchor updates are combined in their 19:1 frequency.Source: §B.6 · p. 20
• FBI uses 12.3% of Direct-10’s time and 10.4% of its counted FLOPs per update in the matched one-GPU benchmark.Both methods use Ktrain = 10; the benchmark varies the gradient estimator and number of posterior draws.Source: Table 9 · p. 20, §B.6 · p. 20
• On Maze-Hard, FBI’s mean Pass@10 is 91.3%, versus 88.8% for latent-space Direct-10, a 2.5 percentage-point difference.In a separate resource benchmark, FBI uses 50.5% of the update time, 10.5% of counted FLOPs, and 23.5% of peak allocated GPU memory of the latent-space variant.Source: Table 10 · p. 20, §B.6 · p. 21
• Unlike latent-space Direct-10, FBI and parameter-space Direct-10 learn posterior means and parameter-wise sampling scales, allowing weight-perturbation magnitude to adapt to training data.The latent-space variant injects independent Gaussian noise into the reasoning latent and processes ten candidates in one batch on one GPU.Source: §B.6 · pp. 20–21

C HYPERPARAMETER SENSITIVITY

• Both selected accuracy and Pass@10 are highest at β = 0.9 and relative KL strength 2 among the tested settings.The figure uses a frozen prefix-Q selector and Pass@10 measures coverage among ten posterior candidates on the test set with one fixed seed.Source: Figure 6 · p. 21, C HYPERPARAMETER SENSITIVITY · p. 21
• At β = 0.9, selected accuracy is 88.1% and Pass@10 is 91.7%; both metrics rise from β = 0, then fall at β = 0.99 and 0.999.Across the EMA sweep, selected accuracy ranges from 87.2% to 88.1% and Pass@10 from 90.2% to 91.7%.Source: C HYPERPARAMETER SENSITIVITY · p. 21
• In the relative-KL sweep, selected accuracy and Pass@10 rise from 87.4% and 90.0% at strength 0.25 to 88.1% and 91.7% at strength 2, then fall to 87.6% and 90.9% at strength 4.The sweep varies IVON’s effective sample size while inversely adjusting weight decay to keep the implied Gaussian prior fixed.Source: C HYPERPARAMETER SENSITIVITY · p. 21
• The β = 0 setting is a boundary control outside the finite-η Bregman-proximal setting analyzed above.The EMA sweep holds relative KL strength at 2, and the relative-KL sweep uses β = 0.9; the same random seed is used throughout.Source: C HYPERPARAMETER SENSITIVITY · p. 21