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

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.

💡 Research Summary

The paper tackles the long‑standing problem of aligning large language models (LLMs) to multiple, often conflicting, human preferences. While reinforcement‑learning‑from‑human‑feedback (RLHF) and single‑reward approaches can only capture a narrow slice of desired behavior, inference‑time alignment methods such as best‑of‑N (BoN), soft‑BoN, and self‑consistency allow a model to be steered toward different criteria at generation time. However, these methods require multiple forward passes, making them computationally expensive, especially when the base model has not been trained with this usage pattern in mind.

To address this, the authors propose Inference‑aware Meta‑Alignment (IAMA), a two‑stage framework. In the first stage the base model π is meta‑trained with an “inference‑aware objective” that explicitly incorporates a set of inference‑time alignment operators {T_i} and their associated reward functions {r_i}. The objective is
R