InfAlign: Inference-aware language model alignment

• URL: http://arxiv.org/abs/2412.19792v4
• Date: Thu, 31 Jul 2025 03:02:43 GMT
• Title: InfAlign: Inference-aware language model alignment
• Authors: Ananth Balashankar, Ziteng Sun, Jonathan Berant, Jacob Eisenstein, Michael Collins, Adrian Hutter, Jong Lee, Chirag Nagpal, Flavien Prost, Aradhana Sinha, Ananda Theertha Suresh, Ahmad Beirami,
• Abstract summary: Language model alignment is a critical step in training modern generative language models.<n>We show that this train/test mismatch makes standard RLHF framework sub-optimal in view of inference-time methods.<n>We propose a framework for inference-aware alignment (InfAlign), which aims to optimize inference-time win rate of the aligned policy against the base model.
• Score: 58.66389179049758
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Language model alignment is a critical step in training modern generative language models. Alignment targets to improve win rate of a sample from the aligned model against the base model. Today, we are increasingly using inference-time algorithms (e.g., Best-of-N, controlled decoding, tree search) to decode from language models rather than standard sampling. We show that this train/test mismatch makes standard RLHF framework sub-optimal in view of such inference-time methods. To this end, we propose a framework for inference-aware alignment (InfAlign), which aims to optimize inference-time win rate of the aligned policy against the base model. We prove that for any inference-time decoding procedure, the optimal aligned policy is the solution to the standard RLHF problem with a transformation of the reward. This motivates us to provide the calibrate-and-transform RL (InfAlign-CTRL) algorithm to solve this problem, which involves a reward calibration step and a KL-regularized reward maximization step with a transformation of the calibrated reward. For best-of-N sampling and best-of-N jailbreaking, we propose specific transformations offering up to 3-8% improvement on inference-time win rates. Finally, we also show that our proposed reward calibration method is a strong baseline for optimizing standard win rate.

Related papers

• Psi-Sampler: Initial Particle Sampling for SMC-Based Inference-Time Reward Alignment in Score Models [10.542645300983878]
  $Psi$-Sampler is an SMC-based framework incorporating pCNL-based initial particle sampling.<n>Inference-time reward alignment with score-based generative models has gained significant traction.
  arXiv  Detail & Related papers  (2025-06-02T05:02:33Z)
• Alignment-Augmented Speculative Decoding with Alignment Sampling and Conditional Verification [33.05591553169347]
  We present a training-free alignment-augmented speculative decoding algorithm.<n>Our method achieves a mean acceptance length up to 2.39 and speed up generation by 2.23.
  arXiv  Detail & Related papers  (2025-05-19T14:55:41Z)
• Sample, Don't Search: Rethinking Test-Time Alignment for Language Models [55.2480439325792]
  We introduce QAlign, a new test-time alignment approach. As we scale test-time compute, QAlign converges to sampling from the optimal aligned distribution for each individual prompt. By adopting recent advances in Markov chain Monte Carlo for text generation, our method enables better-aligned outputs without modifying the underlying model or even requiring logit access.
  arXiv  Detail & Related papers  (2025-04-04T00:41:40Z)
• Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design [87.58981407469977]
  We propose a novel framework for inference-time reward optimization with diffusion models inspired by evolutionary algorithms. Our approach employs an iterative refinement process consisting of two steps in each iteration: noising and reward-guided denoising.
  arXiv  Detail & Related papers  (2025-02-20T17:48:45Z)
• Provably Efficient Online RLHF with One-Pass Reward Modeling [59.30310692855397]
  We propose a one-pass reward modeling method that does not require storing the historical data and can be computed in constant time.<n>We provide theoretical guarantees showing that our method improves both statistical and computational efficiency.<n>We conduct experiments using Llama-3-8B-Instruct and Qwen2.5-7B-Instruct models on the Ultrafeedback-binarized and Mixture2 datasets.
  arXiv  Detail & Related papers  (2025-02-11T02:36:01Z)
• Gradient Correction in Federated Learning with Adaptive Optimization [19.93709245766609]
  We propose tt FAdamGC, the first algorithm to integrate client-drift compensation into adaptive optimization.<n>We show that tt FAdamGC consistently outperforms existing methods in total communication and cost across varying levels of data.
  arXiv  Detail & Related papers  (2025-02-04T21:21:30Z)
• Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control [26.195547996552406]
  We cast reward fine-tuning as optimal control (SOC) for dynamical generative models that produce samples through an iterative process. We find that our approach significantly improves over existing methods for reward fine-tuning, achieving better consistency, realism, and generalization to unseen human preference reward models.
  arXiv  Detail & Related papers  (2024-09-13T14:22:14Z)
• Decoding-Time Language Model Alignment with Multiple Objectives [116.42095026960598]
  Existing methods primarily focus on optimizing LMs for a single reward function, limiting their adaptability to varied objectives. Here, we propose $textbfmulti-objective decoding (MOD)$, a decoding-time algorithm that outputs the next token from a linear combination of predictions. We show why existing approaches can be sub-optimal even in natural settings and obtain optimality guarantees for our method.
  arXiv  Detail & Related papers  (2024-06-27T02:46:30Z)
• PARL: A Unified Framework for Policy Alignment in Reinforcement Learning from Human Feedback [106.63518036538163]
  We present a novel unified bilevel optimization-based framework, textsfPARL, formulated to address the recently highlighted critical issue of policy alignment in reinforcement learning. Our framework addressed these concerns by explicitly parameterizing the distribution of the upper alignment objective (reward design) by the lower optimal variable. Our empirical results substantiate that the proposed textsfPARL can address the alignment concerns in RL by showing significant improvements.
  arXiv  Detail & Related papers  (2023-08-03T18:03:44Z)
• Modular Conformal Calibration [80.33410096908872]
  We introduce a versatile class of algorithms for recalibration in regression. This framework allows one to transform any regression model into a calibrated probabilistic model. We conduct an empirical study of MCC on 17 regression datasets.
  arXiv  Detail & Related papers  (2022-06-23T03:25:23Z)
• Neural Improvement Heuristics for Graph Combinatorial Optimization Problems [49.85111302670361]
  We introduce a novel Neural Improvement (NI) model capable of handling graph-based problems where information is encoded in the nodes, edges, or both. The presented model serves as a fundamental component for hill-climbing-based algorithms that guide the selection of neighborhood operations for each.
  arXiv  Detail & Related papers  (2022-06-01T10:35:29Z)
• Fast Convex Optimization for Two-Layer ReLU Networks: Equivalent Model Classes and Cone Decompositions [47.276004075767176]
  We develop software for convex optimization of two-layer neural networks with ReLU activation functions. We show that convex gated ReLU models obtain data-dependent algorithms for the ReLU training problem.
  arXiv  Detail & Related papers  (2022-02-02T23:50:53Z)
• Obtaining Adjustable Regularization for Free via Iterate Averaging [43.75491612671571]
  Regularization for optimization is a crucial technique to avoid overfitting in machine learning. We establish an averaging scheme that converts the iterates of SGD on an arbitrary strongly convex and smooth objective function to its regularized counterpart. Our approaches can be used for accelerated and preconditioned optimization methods as well.
  arXiv  Detail & Related papers  (2020-08-15T15:28:05Z)
• Pre-training Is (Almost) All You Need: An Application to Commonsense Reasoning [61.32992639292889]
  Fine-tuning of pre-trained transformer models has become the standard approach for solving common NLP tasks. We introduce a new scoring method that casts a plausibility ranking task in a full-text format. We show that our method provides a much more stable training phase across random restarts.
  arXiv  Detail & Related papers  (2020-04-29T10:54:40Z)
• A Speaker Verification Backend for Improved Calibration Performance across Varying Conditions [21.452221762153577]
  We present a discriminative backend for speaker verification that achieved good out-of-the-box calibration performance. All parameters of the backend are jointly trained to optimize the binary cross-entropy for the speaker verification task. We show that this simplified method provides similar performance to the previously proposed method while being simpler to implement, and having less requirements on the training data.
  arXiv  Detail & Related papers  (2020-02-05T15:37:46Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.