InfAlign: Inference-aware language model alignment

• URL: http://arxiv.org/abs/2412.19792v2
• Date: Mon, 30 Dec 2024 09:37:33 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 has become a critical step in training modern generative language models.<n>We show that the alignment framework is sub-optimal in view of inference-time methods.<n>We propose a framework for inference-aware alignment (IAPO)
• Score: 58.66389179049758
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Language model alignment has become a critical step in training modern generative language models. The goal of alignment is to finetune a reference model such that the win rate of a sample from the aligned model over a sample from the reference model is high, subject to a KL divergence constraint. 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. However, the alignment objective does not capture such inference-time decoding procedures. We show that the existing alignment framework is sub-optimal in view of such inference-time methods. We then modify the alignment objective and propose a framework for inference-aware alignment (IAPO). We prove that for any inference-time decoding algorithm, the optimal solution that optimizes the inference-time win rate of the aligned policy against the reference policy is the solution to the typical RLHF problem with a transformation of the reward. This motivates us to provide the KL-regularized calibrate-and-transform RL (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. We particularize our study to two important inference-time strategies: best-of-N sampling and best-of-N jailbreaking, where N responses are sampled from the model and the one with the highest or lowest reward is selected. We propose specific transformations for these strategies and demonstrate that our framework offers significant improvements over existing state-of-the-art methods for language model alignment. Empirically, we outperform baselines that are designed without taking inference-time decoding into consideration by 8-12% and 4-9% on inference-time win rates over the Anthropic helpfulness and harmlessness dialog benchmark datasets.

Related papers

• 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)
• 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.