Generalized Distributional Alignment Games for Unbiased Answer-Level Fine-Tuning

Abstract: The Distributional Alignment Game framework provides a powerful variational perspective on Answer-Level Fine-Tuning (ALFT). However, standard algorithms for these games rely on estimating logarithmic rewards from small batches, introducing a systematic bias due to Jensen's inequality that can destabilize training. In this paper, we systematically resolve this structural estimation bias. First, we generalize the alignment game to arbitrary Bregman divergences, showing that for a family of geometries inducing polynomial rewards, we can construct provably exact and unbiased estimators using U-statistics. Second, for the canonical KL divergence game where an exact solution is impossible, we derive a globally robust minimax polynomial estimator that is provably optimal, achieving the fundamental statistical error limit of $Θ(1/K^2)$, which we establish via the Ditzian-Totik theorem. Finally, we synthesize these two approaches to propose a novel Variance-Optimal Augmented Polynomial Optimization Program (AQP) Estimator, proving that by systematically reducing variance, our method achieves not only optimal bias but also provably accelerated game convergence, leading to more efficient and stable training with zero online computational overhead.

• The paper presents a novel framework that replaces empirical logarithms with unbiased polynomial estimators to counter small-batch estimation bias in answer-level fine-tuning.
• It develops a minimax polynomial estimator for KL-divergence based games, achieving an optimal Θ(1/K²) bias bound and robust performance even for rare answers.
• The study further proposes a variance-optimal AQP estimator that accelerates convergence and stabilizes reward computation, offering a practical toolkit for improved RL training.

Generalized Distributional Alignment Games for Unbiased Answer-Level Fine-Tuning

Motivation and Problem Statement

The paper addresses a critical statistical challenge in Answer-Level Fine-Tuning (ALFT) of LLMs: the estimation bias inherent to reward computation from small-sample batches. Conventional algorithms rely on per-answer rewards derived from empirical log-probabilities ($\log \hat{q}(z)$) computed within a limited set of outputs (e.g., $K=16$). Jensen's inequality causes a systematic bias $\mathcal{O}(1/K)$ in reward estimation, penalizing rare or exploratory outputs and accelerating mode collapse. This bias is not specific to any single RL algorithm, but a fundamental bottleneck in optimizing distributional objectives with policy gradient (PG) methods.

Theoretical Framework: Distributional Alignment Games

ALFT is formalized as a min-max Distributional Alignment Game, where the generative policy $\pi(y|x)$ contends against an auxiliary target answer distribution $q(z)$. Via Fenchel duality, the game separates the intractable marginalization over latent reasoning paths from tractable optimization in answer space. The Nash equilibrium recovers the optimal answer-level policy, and existing algorithms implement this game using alternating best-response dynamics (GRPO) applied to finite sample batches.

Practically, empirical reward functions use shaped signals such as $R(z) = \log \hat{q}(z)$, where $\hat{q}(z)$ is the observed answer frequency in the batch. The non-linearity of $h(\hat{q})$ introduces a negative bias for rare answers, resulting in unstable RL training and premature collapse. The paper mathematically characterizes this as an unavoidable $O(1/K)$ structural bias for strictly concave mappings (logarithmic, KL-divergence).

Generalization via Bregman Divergences: Unbiased Estimation

To resolve the bias, the framework generalizes the Distributional Alignment Game from KL divergence to arbitrary Bregman divergences. The reward mapping $h(q)$ is shown to be determined by the gradient of a distance-generating function. For reward mappings that are polynomials of degree $K=16$0, the paper demonstrates that unbiased estimators can be constructed via U-statistics. For example, the collision penalty and Gini-style diversity objectives result in estimators where empirical batch frequencies produce unbiased rewards.

The paper rigorously characterizes the structure required for unbiased estimation: only games with rewards as bounded-degree polynomials admit exact unbiased estimators in small batches. This generalization enables efficient and stable ALFT optimization when objectives do not critically require the logarithmic geometry of KL divergence.

Optimal Estimation in KL-Regularized Games: Minimax Polynomial Estimator

For canonical KL-alignment games, unbiased estimation is impossible due to the unbounded curvature of the logarithm at the boundary. The paper establishes, via the Ditzian-Totik theorem, a fundamental lower bound of $K=16$1 on the worst-case gradient bias achievable by any estimator, regardless of its functional form.

The authors develop a globally robust minimax polynomial estimator for the KL setting. Offline, the estimator is optimized as a semi-infinite linear program in the Bernstein basis, yielding a vector of coefficients $K=16$2 for batch answer counts $K=16$3. This estimator achieves the $K=16$4 bias bound uniformly, including for rare answers, and is computationally trivial to deploy: an $K=16$5 table lookup replaces the empirical logarithm during RL training.

Further, the paper proves that local Taylor corrections, even with boundary adjustment (Laplace smoothing or hybrid schemes), cannot uniformly achieve better than $K=16$6 rate, and only the minimax polynomial estimator can optimally neutralize boundary bias.

Variance-Optimal Augmented Polynomial Optimization (AQP) Estimator

Variance reduction is addressed by formulating an Augmented Quadratic Program (AQP) to minimize the estimator's second moment under a specified bias budget $K=16$7. The AQP traces the bias-variance Pareto frontier and yields variance-optimal estimators for small batch sizes. Sample splitting techniques are shown to provide no actual variance reduction due to the completeness of the binomial sufficient statistic (Rao-Blackwell theorem). The variance-optimal estimator, computed offline as an SOCP, combines bias minimization with optimized convergence rates for alternating game dynamics.

The paper presents formal bounds asserting accelerated convergence in Nash equilibrium approximation, with detailed guarantees on bias, variance, and boundary stability.

Practical Toolkit and Broader Implications

Three unbiased estimation strategies are provided:

1. Polynomial Geometry with U-statistics: For diversity objectives, exact unbiased estimation is achieved efficiently.
2. Minimax Polynomial Estimator for KL Games: Theoretically optimal bias for canonical KL objectives, robust to rare answer bias.
3. Variance-Optimal AQP Estimator: Accelerates convergence, minimizes variance with no computational overhead.

These approaches are recommended based on alignment targets, training speed, and stability requirements. The toolkit extends beyond ALFT to entropy-regularized RL (including Soft Actor-Critic) and any setting where reward computation involves batch-estimated non-linear functionals. By replacing empirical logarithms with precomputed optimal estimators, structural anti-exploration penalties are eliminated, enabling exploration dynamics as intended in modern RL paradigms.

Conclusion

The paper systematically resolves the estimation bias in answer-level fine-tuning of LLMs by generalizing Distributional Alignment Games, introducing unbiased polynomial estimators and optimal low-bias reward mappings for KL-divergence objectives, and establishing variance-optimality via AQP. The practical implication is that state-of-the-art RL algorithms for ALFT and entropy-regularized domains can be upgraded to unbiased, stable, and computationally efficient pipelines by simply substituting the empirical logarithm with an offline-computed array lookup. This enables rigorous convergence to optimal answer-level policies, enhances training stability, and provides an actionable toolkit with broad implications for RL, LLM alignment, and distributional learning objectives.