Distributional Alignment Games for Answer-Level Fine-Tuning

TL;DR

Proposes a distributional alignment game framework for answer-level fine-tuning, utilizing GRPO for efficient training, achieving significant performance gains.

cs.LG 🔴 Advanced 2026-04-30 65 views
Mehryar Mohri Jon Schneider Yifan Wu
deep learning fine-tuning game theory distribution alignment reasoning

Key Findings

Methodology

This work formulates answer-level fine-tuning as a two-player distributional game, leveraging Fenchel duality to convert intractable marginal probability optimization into a tractable projection problem. An auxiliary target distribution q is introduced, with the policy π (generator) minimizing divergence from q, while q adapts to promote diversity, coherence, or safety. The core algorithm, GRPO, samples groups of trajectories, computes log q(E(y)) as rewards, and updates π via advantage-weighted policy gradients. The framework unifies multiple objectives, provides theoretical guarantees linking Nash equilibrium to optimal solutions, and scales efficiently with existing policy gradient methods.

Key Results

  • On GSM8K and TriviaQA, the proposed Coherence-GRPO and Pairwise-GRPO algorithms improved mathematical reasoning accuracy by 3.18% to 9.18%, with relative gains of 4.8% to 12.46%. In question answering, the models achieved up to 42.06% higher exact match and 18.12% F1 scores, outperforming baseline fine-tuning.
  • The experiments validated the theoretical claims, demonstrating that the equilibrium of the formulated game corresponds to the optimal answer distribution, reducing variance and enhancing training stability.
  • The framework effectively incorporates multiple objectives—diversity, coherence, safety—within a unified mathematical structure, enabling flexible multi-goal optimization in large language models.

Significance

This research addresses the core computational challenge in answer-level fine-tuning by transforming intractable marginalization into a game-theoretic projection problem. The unified framework not only advances theoretical understanding but also offers practical algorithms compatible with existing policy optimization techniques. It paves the way for more robust, diverse, and safe large language models, with broad implications for reasoning, generation, and alignment tasks. The approach bridges the gap between heuristic heuristics and rigorous optimization, fostering new directions in multi-objective model training.

Technical Contribution

The paper introduces a novel duality-based formulation of answer-level fine-tuning as a distributional game, providing a rigorous theoretical foundation. It develops scalable algorithms based on GRPO, capable of optimizing multiple objectives simultaneously. The framework unifies diverse goals—diversity, coherence, safety—under a single convex-concave min-max formulation, with proven convergence guarantees. The theoretical analysis links Nash equilibria to optimal solutions, offering insights into the design of reward functions and target distributions. Extensive experiments validate the effectiveness of the approach across reasoning and question-answering benchmarks.

Novelty

This work is the first to systematically model answer-level fine-tuning as a distributional game grounded in Fenchel duality, providing a unified theoretical and algorithmic framework. Unlike prior heuristic or single-objective methods, it captures multiple alignment goals within a convex-concave min-max structure, enabling principled multi-objective optimization. The integration of GRPO with the game-theoretic formulation offers a scalable, theoretically sound approach to complex fine-tuning tasks in large language models, marking a significant advancement in the field.

Limitations

  • The method relies on accurate estimation of target distributions, which may be challenging in highly diverse or unseen answer spaces, potentially affecting robustness.
  • Computational costs remain high for very large models or extensive multi-objective settings, limiting immediate scalability in resource-constrained environments.
  • Parameter sensitivity and hyperparameter tuning are critical, with no fully automated mechanism, possibly impacting stability and reproducibility.

Future Work

Future directions include developing more robust target distribution estimation techniques, reducing computational overhead, and extending the framework to multi-modal and multi-task settings. Further theoretical analysis on convergence and equilibrium stability will enhance reliability. Exploring automatic hyperparameter tuning and integrating reinforcement learning strategies could improve adaptability. Ultimately, the goal is to enable real-time, multi-objective fine-tuning for large-scale models in diverse application domains.

AI Executive Summary

In the rapidly evolving field of natural language processing, the ability of large language models (LLMs) to produce correct and coherent answers is paramount. Traditional fine-tuning approaches focus on aligning model outputs with specific training data or preferences, but they struggle with answer-level optimization due to the computational intractability of marginalizing over latent reasoning paths. This paper introduces a groundbreaking framework that reconceptualizes answer-level fine-tuning as a strategic game between a generator policy and an auxiliary target distribution. Leveraging Fenchel duality, the authors transform the complex marginalization problem into a tractable projection task, where the target distribution guides the policy towards desired properties such as diversity, coherence, or safety.

The core innovation lies in formulating this as a distributional alignment game, where the Nash equilibrium corresponds to the optimal answer distribution. This theoretical insight unifies various recent heuristics—like inverse-frequency rewards for diversity and consensus-based methods for self-improvement—under a single convex-concave min-max framework. To operationalize this, the authors adapt Group Relative Policy Optimization (GRPO), enabling efficient, scalable training by sampling groups of trajectories and computing rewards based on the log probability under the target distribution.

Extensive experiments on GSM8K and TriviaQA demonstrate the effectiveness of the approach. The proposed algorithms, Coherence-GRPO and Pairwise-GRPO, achieve significant improvements in reasoning accuracy and question-answering metrics, with relative gains up to 42%. These results validate the theoretical guarantees and showcase the framework’s capacity to handle multiple objectives simultaneously.

This work advances the theoretical understanding of answer-level alignment, providing a unified, scalable, and principled approach to multi-objective fine-tuning. It opens new avenues for developing more robust, diverse, and safe language models, with broad implications for AI research and applications. Future work will focus on improving target distribution estimation, reducing computational costs, and extending the framework to multi-modal and multi-task scenarios, aiming for real-time, multi-objective model adaptation in complex environments.

Deep Dive

Abstract

We focus on the problem of \emph{Answer-Level Fine-Tuning} (ALFT), where the goal is to optimize a language model based on the correctness or properties of its final answers, rather than the specific reasoning traces used to produce them. Directly optimizing answer-level objectives is computationally intractable due to the need to marginalize over the vast space of latent reasoning paths. To overcome this, we propose a general game-theoretical framework that lifts the problem to a \emph{Distributional Alignment Game}. We formulate ALFT as a two-player game between a Policy (the generator) and a Target (an auxiliary distribution). We prove that the Nash Equilibrium of this game corresponds exactly to the solution of the original answer-level optimization problem. This variational perspective transforms the intractable marginalization problem into a tractable projection problem. We demonstrate that this framework unifies recent approaches to diversity and self-improvement (coherence) and provide efficient algorithms compatible with Group Relative Policy Optimization (GRPO), such as Coherence-GRPO, yielding significant complexity gains in mathematical reasoning tasks.

cs.LG cs.GT