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 minimal computational overhead at inference time. IAMA trains a base model such that it can be effectively aligned to multiple task optima 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.

AI chatbots must be adjusted so their responses match what users want but different users want different things. Some prefer brief answers; others want detailed ones. Some prioritize helpfulness; others safety. These preferences often conflict, making it impossible for a single model to satisfy everyone at once. We propose a training method that prepares a single model to produce diverse outputs, so that a simple selection step at usage time, such as generating a few candidate answers and picking the best one, can quickly satisfy whichever preference is needed. We also develop an efficient training algorithm with mathematical guarantees that it finds the best solution. This allows deploying one model that flexibly adapts to many user needs without maintaining separate models for each preference or requiring heavy computation at usage time, making personalized AI assistants more practical and resource-efficient.