Inference-Aware Meta-Alignment of LLMs via Non-Linear GRPO

TL;DR

This work tackles the challenge of aligning LLMs to multiple, potentially conflicting preferences under limited inference-time compute. It introduces inference-aware meta-alignment (IAMA), a two-stage framework that meta-trains a base policy to be effectively adaptable to multiple inference-time alignment criteria, and formulates the objective as a non-linear optimization in the space of probability measures. To solve the resulting non-linear problem, it develops non-linear GRPO, a first-order-variation based extension of GRPO/TRPO-like methods, with convergence guarantees for BoN-type objectives and a provable bound in the inexact setting. The approach is empirically validated on tasks such as length control and RLHF-style objectives, showing that the meta-aligned model can produce diverse responses and achieve better Pareto frontiers when BoN or SoftBoN inference-time transformations are used. Overall, IAMA provides a principled, efficient route to multi-criteria alignment suitable for practical deployment with constrained inference-time resources, along with theoretical guarantees and broad applicability in RLHF contexts.

Abstract

Aligning large language models (LLMs) to diverse human preferences is fundamentally challenging since criteria can often conflict with each other. Inference-time alignment methods have recently gained popularity as they allow LLMs to be aligned to multiple criteria via different alignment algorithms at inference time. However, inference-time alignment is computationally expensive since it often requires multiple forward passes of the base model. In this work, we propose inference-aware meta-alignment (IAMA), a novel approach that enables LLMs to be aligned to multiple criteria with limited computational budget at inference time. IAMA trains a base model such that it can be effectively aligned to multiple tasks via different inference-time alignment algorithms. To solve the non-linear optimization problems involved in IAMA, we propose non-linear GRPO, which provably converges to the optimal solution in the space of probability measures.

Figures (8)

• Figure 1: Illustration of inference-aware meta-alignment (IAMA). The base model $\pi_{\mathrm{ref}}$ is meta-trained so that it can be effectively adapted to multiple task optima $\{\pi^*_{i}\}_{i=1}^m$ via inference-time alignment algorithms $\{\mathcal{T}_i\}_{i=1}^m$. Compared to deploying $\pi_{\mathrm{ref}}$ directly, IAMA enables obtaining better aligned models for each task with limited computational budget at inference time.
• Figure 2: Optimal IAMA policies $\pi^*_{\mathrm{IAMA}}$ with BoN sampling ($N=2, 4, 8$). The optimal policy $\pi^*_{\mathrm{IAMA}}$ with large $N$ produces diverse outputs to cover both modes at $y=0$ and $y=1$.
• Figure 3: Response length distributions for different alignment methods on the length reward task. The meta-aligned model via non-linear GRPO can effectively adapt to both short and long response preferences via (Soft) BoN sampling, while the aligned model via standard GRPO generates only medium-length responses.
• Figure 4: Averaged rewards on the evaluation set. The meta-aligned model via non-linear GRPO can effectively trade off between helpfulness and harmlessness via BoN ($N=4$) sampling.
• Figure 5: Expected error of linearized loss estimation with different $M$ and $N$ values. From left to right, $N=1,4,8,16$. The shaded area indicates the bias-variance decomposition of the expected error.

Theorems & Definitions (18)

• Proposition 3.1
• Remark 3.2
• Definition 4.1
• Proposition 4.2
• Proposition 4.3
• Theorem 5.2
• Theorem 5.3
• Lemma 5.4
• Lemma 5.5
• Lemma 5.6