Paper deep dive
Inference-Aware Meta-Alignment of LLMs via Non-Linear GRPO
Shokichi Takakura, Akifumi Wachi, Rei Higuchi, Kohei Miyaguchi, Taiji Suzuki
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 93%
Last extracted: 3/11/2026, 1:02:32 AM
Summary
The paper introduces Inference-Aware Meta-Alignment (IAMA), a framework for aligning Large Language Models (LLMs) to multiple, potentially conflicting reward functions by optimizing a base model for inference-time adaptation. To address the resulting non-linear optimization problem, the authors propose 'non-linear GRPO', a novel algorithm that uses first-order functional variations to linearize the objective, enabling efficient training and proving linear convergence to optimal solutions.
Entities (5)
Relation Signals (3)
IAMA → utilizes → Non-linear GRPO
confidence 95% · To solve the non-linear optimization problems involved in IAMA, we propose non-linear GRPO
Non-linear GRPO → optimizes → LLM
confidence 90% · IAMA trains a base model such that it can be effectively aligned to multiple tasks via different inference-time alignment algorithms.
IAMA → generalizes → MAML
confidence 85% · This can be regarded as a generalization of celebrated MAML framework in the space of probability measures.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2602.01603
- Canonical: https://arxiv.org/abs/2602.01603
Trouble viewing inline? Open PDF directly →
Full Text
91,250 characters extracted from source content.
Expand or collapse full text
Inference-Aware Meta-Alignment of LLMs via Non-Linear GRPO Shokichi Takakura Akifumi Wachi Rei Higuchi Kohei Miyaguchi Taiji Suzuki 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. Machine Learning, ICML 1 Introduction Aligning large language models (LLMs) via reinforcement learning from human feedback (RLHF) (christiano2017deep; ouyang2022training) or verifiable rewards (RLVR) (lambert2024tulu; shao2024deepseekmath; jaech2024openai) is a crucial step to ensure that they generate responses aligned with human preferences or specific criteria. However, in real-world applications, the requirements of LLM responses are often diverse depending on the situation and cannot be captured by a single reward function. Several works have explored the idea of aligning LLMs to multiple reward functions simultaneously (agnihotri2025multi; dai2024safe), but it is fundamentally challenging or even impossible since criteria can often conflict with each other (bai2022training). Recently, the paradigm of inference-time alignment has gained popularity, where the base model is aligned at inference time using techniques such as best-of-N (BoN) sampling (nakano2021webgpt), variants of BoN (verdun2025soft; jinnai2025regularized), and self-consistency (wang2023self). Since these inference-time alignment algorithms can modify the outputs of the base model at inference time based on different reward functions or criteria, the model can be aligned to diverse preferences. However, inference-time alignment is computationally expensive since it often requires multiple forward passes of the base model at inference time. With limited computational budget, the performance of inference-time alignment can be poor if the base model is not properly trained. For example, if we aggregate multiple rewards into a single metric and align the model to it, this results in a model that yields mediocre responses across all rewards, leading to a loss of diversity. Consequently, this makes it difficult to sample responses tailored to a specific reward using inference-time alignment methods such as BoN. Thus, we pose the following question: Can we train a model that can be efficiently aligned to diverse preferences via inference-time alignment? To address this issue, we propose inference-aware meta-alignment (IAMA), a novel and general framework for aligning LLMs to multiple reward functions via two-stage alignment procedure: 1) meta-train the base model with the inference-aware objective and 2) inference-time alignment with different reward functions. We show that the two-stage alignment process can be regarded as a model-agnostic meta-learning (MAML) (finn2017model) in the space of probability measures and meta-learning problem can be formulated as a non-linear optimization problem with respect to the policy distribution. This is in contrast to the standard expected reward maximization problems, which are linear with respect to the policy distribution. Due to the non-linearity of the objective, we cannot directly apply existing policy optimization algorithms such as TRPO (schulman2015trust), PPO (schulman2017proximal), and GRPO (shao2024deepseekmath) to solve the problem. To address this issue, we propose a novel algorithm called non-linear GRPO, which linearizes the objective at each iteration using the first-order variation. While the optimization landscape of LLMs is generally highly non-convex, our objective in the space of probability measures is shown to be convex for certain inference-time alignment algorithms such as BoN. Using this property, we prove the linear convergence of the proposed algorithm to the optimal solution. Empirically, we demonstrate that the proposed method can effectively align LLMs to multiple reward functions simultaneously. Our contribution can be summarized as follows: • We propose inference-aware meta-alignment (IAMA), a novel and general framework which optimizes a single base model so that it can adapt effectively to diverse preferences via inference-time alignment. This can be regarded as a generalization of celebrated MAML framework in the space of probability measures. • We show that IAMA can be formulated as a non-linear optimization problem in the space of measures and to solve the problem, we develop non-linear GRPO, which can be regarded as a mirror descent with respect to KL divergence. • Theoretically, we prove the convexity of the objective for BoN-type inference-time alignment methods and show that non-linear GRPO converges linearly to the optimal solution. We also empirically demonstrate the effectiveness of the proposed method on multiple tasks and models. While implementation is straightforward, it significantly pushes Pareto frontiers in a multi-objective RLHF task. 1.1 Other Related Work Inference-aware Alignment Several recent works have explored the idea of incorporating inference-time alignment algorithms into the training objective (chow2025inference; tang2025optimizing; balashankarinfalign). Another line of work has investigated the pass@k@k or max@k@k objectives, which is equivalent to applying BoN sampling at inference time (walder2025pass; bagirov2025best). However, these works typically focus on a single inference-time alignment algorithm and a single reward function. Therefore, it is still unclear how to effectively align LLMs to multiple reward functions simultaneously via two-stage alignment. In addition, they consider only BoN sampling and do not provide a general framework that can handle various inference-time alignment algorithms. The only exception is InfAlign framework (balashankarinfalign) but it focuses on a single inference-time alignment algorithm. In addition, while their problem formulation can be applied to general inference-time alignment algorithms, the proposed InfAlign-CTRL method is applicable only for BoN-type alignment algorithms and limits the reward functions to be win rate vs. a reference model. Therefore, it is still unclear how to solve the general inference-aware alignment problem even for a single inference-time alignment algorithm. See Appendix B for more discussion. Meta-learning Meta-learning (thrun1998learning) aims to train models that can quickly adapt to new tasks via few-shot learning or fast adaptation. One of the most popular approaches to meta-learning is MAML (finn2017model), which learns model parameters such that it can be effectively adapted to multiple tasks via few gradient steps. However, the inference-time alignment procedure does not update the model parameters and thus, cannot be directly handled by the original MAML framework. MetaAlign (zhang2025-metaalign) achieves dynamic alignment by conditioning the model on system prompts that specify preferences. While effective for explicit instruction following, it may struggle to capture complex or implicit preferences that are not easily articulated in prompts. Optimization in the Space of Probability Measures Optimization problems in the space of probability measures appear in various machine learning problems such as reinforcement learning, variational inference, and generative modeling (chu2019probability). A line of work has explored particle-based methods such as mean-field Langevin dynamics and particle dual averaging (mei2018mean; nitanda2021particle). However, these methods require retraining for each resampling. Another line of work has studied mirror descent (frankowski2022mirror; yao2024minimizing) or dual averaging (kawata2025direct) methods in the space of probability measures. Although these methods are theoretically straightforward, they focus on tractable setups, where the first-order variation can be computed in closed form or includes computationally expensive steps such as normalizing flows (yao2024minimizing). Therefore, it is still unclear how to solve non-linear optimization problems involved in LLM alignment. 2 Problem Settings We consider a language model which is represented as a probability distribution (or policy) π over responses y given a context x. We denote the space of all possible contexts and responses as X and Y, respectively. We assume that the contexts are sampled from a distribution ρ over X. We define the space of all possible policies as P and the reference policy as πref∈ _ref . Throughout this paper, we often omit the dependence on the context x for simplicity if the notation does not cause confusion. For instance, we write π(y)π(y) instead of π(y∣x)π(y x). Training-time Alignment Let r:×→[0,rmax]r:X×Y→[0,r_ ] be a reward function which measures the quality of the response y given the context x. The objective of alignment is to find a policy π(y∣x)π(y x) that maximizes the expected reward while staying close to a reference policy πref _ref. Typically, the objective is formulated as a KL-regularized expected reward maximization problem: ℛ[π]=y∼π(⋅∣x),x∼ρ[r(x,y)]−βKL[π∣πref], [π]=E_y π(· x),x ρ [r(x,y) ]- [π _ref], where KLKL is the Kullback-Leibler divergence defined as KL[π∣πref]=y∼π(⋅∣x),x∼ρ[logπ(y∣x)πref(y∣x)]KL[π _ref]=E_y π(· x),x ρ [ π(y x) _ref(y x) ] and β≥0β≥ 0 is a regularization parameter. Inference-time Alignment Rarely do we use the aligned model as it is at inference time. Instead, we often employ inference-time alignment algorithms that modify the outputs of the base model π at inference time. One of the most popular methods is best-of-N (BoN) sampling (nakano2021webgpt), where we sample multiple responses from the base model π and select the response with the highest reward. Soft BoN (verdun2025soft) is a soft version of BoN, which (approximately) samples responses from an exponentially tilted distribution of rewards. These are equivalent to transforming the base model π to a new policy and sampling a response from the transformed policy. We denote such an inference-time alignment algorithm as a functional :→T:P that maps a base model π∈π to an aligned model [π]T[π]. In the case of BoN with N samples, the transformed policy is given as N⋅C[π](y)N−1π(y)N· C[π](y)^N-1π(y) (beirami2025theoretical; gui2024bonbon; amini2025variational) where C[π](y)=∫r(y′)≤r(y)π(y′)y’⋅C[π](y)= _r(y )≤ r(y)π(y ) y' is the cumulative distribution function of rewards under π, and in the case of soft BoN with temperature parameter τ>0τ>0, it is given as [π](y)∝exp(r(y)/τ)π(y)T[π](y) (r(y)/τ)π(y) (verdun2025soft). Other than (Soft) BoN, there exist various inference-time alignment algorithms such as self-consistency (wang2023self) and a MCMC-based method (karan2025reasoning) and our general framework can handle these algorithms as well. 3 Proposed Framework: Inference-aware Meta-Alignment Pπref _refπ1∗π^*_1π2∗π^*_2π3∗π^*_31T_12T_23T_31T_12T_23T_3Meta alignmentInference-time alignment Figure 1: Illustration of inference-aware meta-alignment (IAMA). The base model πref _ref is meta-trained so that it can be effectively adapted to multiple task optima πi∗i=1m\π^*_i\_i=1^m via inference-time alignment algorithms ii=1m\T_i\_i=1^m. Compared to deploying πref _ref directly, IAMA enables obtaining better aligned models for each task with limited computational budget at inference time. As we have seen in the previous section, modern alignment procedure often involves two stages: 1) training-time alignment of the base model π and 2) inference-time alignment via an algorithm T. Considering the multiple inference-time alignment algorithms ii=1m\T_i\_i=1^m and corresponding reward functions rii=1m\r_i\_i=1^m, a natural question arises: What is the best model π that can be effectively aligned to diverse rewards via inference-time alignment algorithms? To address this question, we propose the following objective for training-time alignment: ℛ[π]=g(R1[π],…,Rm[π])⏟R[π]−βKL[π∣πref], [π]= g(R_1[π],…,R_m[π])_R[π]- [π _ref], (1) where Ri[π]=y∼i[π],x∼ρ[ri(x,y)]R_i[π]=E_y _i[π],x ρ [r_i(x,y) ], and g:ℝm→ℝg:R^m is an aggregation function. The aggregation function g can be weighted sum ∑i=1mwiRi[π] _i=1^mw_iR_i[π] for some weights wi≥0w_i≥ 0, worst-case performance miniRi[π] _iR_i[π], or smooth minimum −1γlog(∑i=1mwiexp(−γRi[π]))- 1γ ( _i=1^mw_i (-γ R_i[π])) for γ>0γ>0, which interpolates between the average and worst-case performance. This objective can be regarded as a generalization of InfAlign framework (balashankarinfalign) to multiple inference-time alignment algorithms and general reward functions beyond win rate against a reference model. Similar motivation has been explored in the context of model-agnostic meta-learning (MAML) (finn2017model), where the goal is to find parameters that can be effectively adapted to multiple tasks via few gradient steps. On the other hand, in the inference-time alignment, the model parameters are not updated during adaptation, but the base model π is transformed to a new policy i[π]T_i[π] in the space of probability measures. We illustrate the concept of proposed framework in Fig. 1. While the formulation looks similar to the standard KL-regularized expected reward maximization, the key difference is that the reward term R[π]R[π] is non-linear with respect to the base model π even for a weighted sum due to the presence of the inference-time alignment algorithms ii=1m\T_i\_i=1^m. This is in contrast to the expected reward Eπ[r]E_π[r], which is linear with respect to π. Due to the non-linearity, we need to develop new algorithms to solve the optimization problem. 3.1 Comparison with Naive Approach A naive approach to align LLMs to multiple reward functions is to train a base model π using weighted sum of expected rewards without considering inference-time alignment. That is, we optimize the following objective: ℛnaive[π]=∑i=1mwiy∼π,x∼ρ[ri(x,y)]−βKL[π∣πref]. _naive[π]= _i=1^mw_iE_y π,x ρ [r_i(x,y) ]- [π _ref]. To see the difference between the naive approach and IAMA, let us consider a simple case where =[0,1]Y=[0,1]. Suppose that there are two reward functions: r1(y) r_1(y) =1−y2,r2(y)=1−(1−y)2. =1-y^2,r_2(y)=1-(1-y)^2. These rewards are completely conflicting since r1r_1 is maximized at y=0y=0 while r2r_2 is maximized at y=1y=1. If we optimize the naive objective with equal weights w1=w2=0.5w_1=w_2=0.5 and ignore the KL regularization for simplicity, the optimal policy is given as πnaive∗=δ0.5π^*_naive= _0.5, which always outputs y=0.5y=0.5. Here, δa _a is the Dirac delta distribution centered at a. On the other hand, if we consider inference-time alignment via BoN sampling with N≥2N≥ 2, we have the following result: Proposition 3.1. Consider the above reward functions r1r_1 and r2r_2. Let 1T_1 and 2T_2 be BoN sampling with N≥2N≥ 2. Then, the optimal policy that maximizes the IAMA objective (1) with equal weights w1=w2=0.5w_1=w_2=0.5 and β=0β=0 is given as πIAMA∗(y)=αyα−1(1−y)α−1(yα+(1−y)α)2 π^*_IAMA(y)= α y^α-1(1-y)^α-1(y^α+(1-y)^α)^2 where α=1N−1α= 1N-1. See Appendix G.1 for the proof. We show the optimal policy πIAMA∗π^*_IAMA for N=2,4,8N=2,4,8 in Fig. 2. Figure 2: Optimal IAMA policies πIAMA∗π^*_IAMA with BoN sampling (N=2,4,8N=2,4,8). The optimal policy πIAMA∗π^*_IAMA with large N produces diverse outputs to cover both modes at y=0y=0 and y=1y=1. In the case of BoN, the base policy generates diverse outputs to cover both modes at y=0y=0 and y=1y=1, which allows the inference-time alignment to select high-reward outputs for both reward functions. This is in contrast to the naive optimal policy πnaive∗π^*_naive, which produces only mediocre outputs. In addition, the optimal policy varies depending on the computational budget at inference time (i.e., the number of samples N). This simple example clearly addresses the question of why IAMA should be considered. Remark 3.2. Another possible approach is to align multiple models πii=1m\ _i\_i=1^m to each reward function rii=1m\r_i\_i=1^m separately and use one of them at inference time depending on the user’s preference (barrett2008learning; rame2023rewarded). However, this approach requires storing multiple models, which can be impractical for large models used in real-world applications. In addition, the preference of users may not be known by the provider even at inference time. In such cases, this approach cannot be applied while IAMA can still be used if users select best responses based on their own preferences. 4 Proposed Method: Non-linear GRPO Since the objective in Eq. (1) is non-linear with respect to the policy π, we cannot directly apply existing policy optimization algorithms such as TRPO (schulman2015trust), PPO (schulman2017proximal), and GRPO (shao2024deepseekmath). To address this issue, we propose a non-linear GRPO algorithm, which is a non-trivial extension of GRPO-type algorithms (shao2024deepseekmath; zheng2025group; liu2025understanding), which sample multiple responses for a given context. From a different perspective, balashankarinfalign proposed InfAlign-CTRL algorithm to solve a special case of Eq. (1) using reward transformation but it is applicable only for 1) BoN-type algorithms, 2) win-rate maximization, and 3) single reward setting. This limitation is inevitable since such transformation can be computed only under these constraints. Our approach eliminates all three of these constraints and is therefore of independent interest. First, we revisit the standard TRPO-type algorithms (including TRPO, PPO, GRPO, and their variants). For notational convenience, we consider the minimization of ℒ[π]:=−ℛ[π]L[π]:=-R[π] instead of maximization of ℛ[π]R[π]. Let πt _t be the current policy at iteration t. Then, TRPO-type algorithms can be interpreted as an approximation of the following mirror descent method in the space of probability distributions: πt+1 _t+1 =argminπℒ~[π]+1ηKL[π∣πt], = *argmin_π L[π]+ 1ηKL[π _t], ℒ~[π] L[π] :=−y∼πt[π(y)πt(y)r(y)]+βKL[π∣πref], :=-E_y _t [ π(y) _t(y)r(y) ]+ [π _ref], where η>0η>0 is a step-size parameter. In TRPO, the proximal term KL[π∣πt]KL[π _t] is replaced with a constraint KL[π∣πt]≤εKL[π _t]≤ for some ε>0 >0 and in PPO and GRPO, it is replaced with the following clipped surrogate objective. πt+1 _t+1 =argminπℒ~[π], = *argmin_π L[π], ℒ~[π] L[π] :=−y∼πt[min(π(y)πt(y)r(y),clipϵ(π(y)πt(y))r(y))] :=-E_y _t [ ( π(y) _t(y)r(y),clip_ε ( π(y) _t(y) )r(y) ) ] +βKL[π∣πref], + [π _ref], where clipϵ(z)=max(min(z,1+ϵ),1−ϵ)clip_ε(z)= ( (z,1+ε),1-ε) for some ϵ>0ε>0. We extend the above TRPO-type algorithms to handle the non-linear reward term R[π]R[π]. First, we introduce the notion of first-order variation (functional derivative). Definition 4.1. A functional Rπ:×→ℝ Rπ:P×Y is called the first-order variation (or functional derivative) of a functional R:→ℝR:P if for all π,π′∈π,π , dR[π+ϵ(π′−π)]dϵ|ϵ=0 dR[π+ε(π -π)]dε |_ε=0 =∫Rπ[π](y)⋅(π′−π)(y)y⋅. = Rπ[π](y)·(π -π)(y) y. Note that the first variation admits a freedom under constant shifts. In the following, we assume that the first variation Rπ[π](y) Rπ[π](y) exists for all π∈π and y∈y . To simplify the notation, we denote Rπ[π] Rπ[π] as R⋅[π] R[π] and ∫R⋅[π](y)⋅g(y)y⋅ R[π](y)· g(y) y as ⟨R⋅[π],g⟩ R[π],g for a function g:→ℝg:Y . Using the first-order variation, we approximate the objective R[π]R[π] around the current policy πt _t as R[π] R[π] ≈R[πt]+∫∂R∂π[πt](y)⋅(π(y)−πt(y))y⋅. ≈ R[ _t]+ Rπ[ _t](y)·(π(y)- _t(y)) y. Using this approximation, we define the surrogate loss as ℒ~[π] L[π] =−(R[πt]+∫∂R∂π[πt](y)⋅(π(y)−πt(y))y⋅) =- (R[ _t]+ Rπ[ _t](y)·(π(y)- _t(y)) y ) +βKL[π∣πref] + [π _ref] =−y∼πt[πt∂R∂π[πt](y)]+βKL[π∣πref]+const. =-E_y _t [ π _t Rπ[ _t](y) ]+ [π _ref]+const. Then, we update the policy by solving the following proximal objective: πt+1=argminπℒ~[π]+1ηKL[π∣πt]. _t+1= *argmin_π L[π]+ 1ηKL[π _t]. Aside from the fact that the term KL[π∣πref]KL[π _ref] remains the same, the above update can be interpreted as a mirror descent in the space of probability measures with KL divergence as the Bregman divergence. While the above update is conceptually straightforward, it is not directly implementable in practice since the first-order variation ∂R∂π[πt](y) Rπ[ _t](y) generally depends on the current policy πt _t in a complex manner. To overcome this issue, we propose to approximate it using empirical samples. In GRPO-type algorithms, we sample multiple responses from the current policy πt _t for a given context x. We denote the sampled responses as yjj=1M\y_j\_j=1^M and the empirical distribution of the samples as π^t=1M∑j=1Mδyj π_t= 1M _j=1^M _y_j. Note that we can use other approximations such as kernel density estimation to obtain a smoother approximation of the current policy (silverman2018density). Then, we approximate the first-order variation as ∂R∂π[π^t](y) Rπ[ π_t](y). Using this approximation, we define the surrogate loss as ℒ^[π] L[π] =−y∼πt[πt∂R∂π[π^t](y)]+βKL[π∣πref], =-E_y _t [ π _t Rπ[ π_t](y) ]+ [π _ref], and we update the policy by approximately solving the following proximal objective: πt+1≃argminπℒ^[π]+1ηKL[π∣πt]. _t+1 *argmin_π L[π]+ 1ηKL[π _t]. The above update can be interpreted as a GRPO update with a modified reward function r~t(y)=∂R∂π[π^t](y) r_t(y)= Rπ[ π_t](y). That is, we can modify the original GRPO by only introducing a drop-in reward function while keeping the rest of the algorithm unchanged. This makes the non-linear GRPO algorithm easy to implement in practice. Note that we can also use other GRPO-type algorithms such as GSPO (zheng2025group) as the base algorithm instead of GRPO. In addition, since non-linear GRPO is quite general, it can be applied to various problems beyond IAMA and includes some existing algorithms as special cases. See Appendix C for more discussion. We summarize the non-linear GRPO algorithm in Algorithm 1. In practice, we can further approximate the above update by replacing the expectation with empirical samples and proximal term with the clipped surrogate objective as in PPO and GRPO. Please refer to Algorithm 2 for the full description of practical implementation of non-linear GRPO. Algorithm 1 Non-linear GRPO 1: Input: Initial policy π0 _0, reference policy πref _ref, objective R, regularization parameter β, step size η, number of iterations T 2: for t=0,1,…,T−1t=0,1,…,T-1 do 3: Sample M responses yjj=1M\y_j\_j=1^M from πt _t 4: Compute approximated functional derivative: r~t(y)=∂R∂π[π^t](y) r_t(y)= Rπ[ π_t](y). 5: Define surrogate loss: ℒ^[π]=−Ey∼πt[π(y)πt(y)⋅r~t(y)]+βKL[π∣πref] L[π]=-E_y _t[ π(y) _t(y)· r_t(y)]+ [π _ref] 6: Update policy: πt+1≃argminπℒ^[π]+1ηKL[π∣πt] _t+1 _π L[π]+ 1ηKL[π _t] 7: end for 8: Return: Final policy πT _T 4.1 First-Order Variation for (Soft) BoN Since our formulation is quite general, it can accommodate various inference-time alignment algorithms. In this section, we consider two important cases including 1) BoN and 2) soft BoN (exponential tilting) and derive the concrete forms of the functional derivatives required to implement the non-linear GRPO algorithm. From the chain rule, the first variation of R can be computed as ∂R∂π[π](y) Rπ[π](y) =∑i=1m∂g∂Ri⋅∂Ri∂π[π](y). = _i=1^m gR_i· R_iπ[π](y). Thus, to compute the first variation of R, it suffices to compute the first-order variations of RiR_i for each inference-time alignment algorithm iT_i and reward function rir_i. The following proposition provides formulas for the first-order variations of RiR_i for BoN and soft BoN, respectively. Proposition 4.2. Let R[π]=y∼BoNN[π][r(y)]R[π]=E_y _N[π] [r(y) ]. Assume that the distribution of r(y)r(y) under π has a density. Then, the first-order variation of R can be computed as ∂R∂π[π](y) Rπ[π](y) =−∫r(y)rmaxN⋅(C[π](r))N−1r⋅. =- _r(y)^r_ N·(C[π](r))^N-1 r. Proposition 4.3. Let R[π]=y∼SoftBoNτ[π][r(y)]R[π]=E_y _τ[π] [r(y) ]. Then, the first-order variation of R can be computed as ∂R∂π[π](y) Rπ[π](y) =r(y)exp(r(y)/τ)∫exp(r(z)/τ)⋅π(z)z⋅ = r(y) (r(y)/τ) (r(z)/τ)·π(z) z −exp(r(y)/τ)⋅∫r(z)exp(r(z)/τ)⋅π(z)z⋅ (∫exp(r(z)/τ)⋅π(z)z⋅)2. - (r(y)/τ)· r(z) (r(z)/τ)·π(z) z ( (r(z)/τ)·π(z) z)^2. See Appendix F for the proofs. The above formulas involve integrals but they are reduced to sum over empirical samples if we replace π with the empirical distribution π π. See Appendix D.1 for details. Therefore, we can easily compute the approximated functional derivatives required in the non-linear GRPO algorithm. 5 Convergence Analysis In this section, we provide the convergence guarantees of the non-linear GRPO algorithm. The optimization of LLMs in the space of parameters is generally non-convex and challenging to analyze. On the other hand, if we regard the optimization in the space of probability distributions (policies), the problem becomes simpler. Specifically, in the case of weighted sum and (smooth) minimum of BoN-type objectives, R[π]R[π] is concave as shown in Section 5.3. We also discuss other cases where R[π]R[π] is concave in Appendix C. First, we provide convergence guarantees of non-linear GRPO for general smooth and concave objectives. Note that the concavity is required only for the convergence analysis and the proposed algorithm can be applied to non-concave objectives (e.g. Soft BoN, Worst-of-N (balashankarinfalign) case) as well. Here, we assume the following about the objective function: Assumption 5.1. The reward functional R is concave and L-smooth relative to the KL-divergence. That is, for all π,π′∈π,π , we have R[π′] R[π ] ≤R[π]+⟨R⋅[π],π′−π⟩, ≤ R[π]+ R[π],π -π , R[π′] R[π ] ≥R[π]+⟨R⋅[π],π′−π⟩−L⋅KL[π′∣π]. ≥ R[π]+ R[π],π -π -L·KL[π π]. for some L>0L>0. Since R[π]R[π] is concave, the objective ℒ[π]=−R[π]+βKL[π∣πref]L[π]=-R[π]+ [π _ref] is strongly convex with respect to the KL-divergence if β>0β>0. Therefore, the optimization problem has a unique optimal policy π∗∈ _* . 5.1 Convergence Guarantees with Exact Proximal Updates In this section, we assume that the proximal objective is solved exactly at each iteration. That is, at step t, πt+1 _t+1 can be computed as πt+1=argminπℒ~[π]+1ηKL[π∣πt]. _t+1= *argmin_π L[π]+ 1ηKL[π _t]. Then, we have the following convergence guarantee: Theorem 5.2. Let πtt=0T\ _t\_t=0^T be the sequence of policies generated by the (exact) non-linear GRPO algorithm with step size η=1/Lη=1/L. Then, we have ℒ[πT]−ℒ[π∗] [ _T]-L[ _*] ≤β⋅KL[π∗∣π0]((L+β)/L)T−1. ≤ β·KL[ _* _0]((L+β)/L)^T-1. See Appendix H.1 for the proof. The proof strategy basically follows that of frankowski2022mirror except for some modifications to handle KL-regularization. This theorem shows that the non-linear GRPO algorithm converges to the optimal policy exponentially fast if β>0β>0. 5.2 Convergence Guarantees with Inexact Proximal Updates Next, we consider the practical case, where the first-order variation is approximated using empirical samples and the proximal objective is solved approximately. That is, at step t, we compute πt+1 _t+1 as an approximate solution to πt+1≈argminπℒ^[π]+1ηKL[π∣πt]. _t+1≈ *argmin_π L[π]+ 1ηKL[π _t]. Note that ∂R∂π[π^t] Rπ[ π_t] is a random variable depending on the sampled responses yjj=1M\y_j\_j=1^M. The first-order optimality condition of proximal objective gives ∂ℒ^∂π[π]+1ηlogπt=const. Lπ[π]+ 1η π _t=const. In practice, due to the limited representation ability of the model and inexact optimization, there exists a residual term rt:=ℒ^π[πt+1]+1ηlogπt+1πt.r_t:= Lπ[ _t+1]+ 1η _t+1 _t. Then, we have the following convergence guarantee: Theorem 5.3. Assume that the approximated first-order variation satisfies [‖∂R∂π[π^t](y)−∂R∂π[πt](y)‖sp2]≤εE [ Rπ[ π_t](y)- Rπ[ _t](y)_sp^2 ]≤ for some ε≥0 ≥ 0, and the optimization residual satisfies [‖rt‖sp2]≤δE [ r_t_sp^2 ]≤δ for some δ≥0δ≥ 0 given πt _t. Here, ‖f‖sp=(supyf(y)−infyf(y))/2 f_sp=( _yf(y)- _yf(y))/2 is the span seminorm. Let πtt=0T\ _t\_t=0^T be the sequence of policies generated by the (inexact) non-linear GRPO algorithm with step size η=1/Lη=1/L and a random index t^∈1,…,T t∈\1,…,T\ following the distribution: Prob(t^=t)∝ (L+β/2L)tProb ( t=t ) ( L+β/2L)^t. Then, we have [ℒ[πt^]−ℒ[π∗]] [L[ _ t]-L[ _*] ] ≤(β/2)⋅KL[π∗∣π0]((L+β/2)/L)T−1+2(ε+δ)β. ≤ (β/2)·KL[ _* _0]((L+β/2)/L)^T-1+ 2( +δ)β. Here, the expectation is taken over the randomness in the approximation of the functional derivative, optimization procedure, and the choice of t t. See Appendix H.2 for the proof. Even with inexact proximal updates, the non-linear GRPO algorithm converges to a neighborhood of the optimal policy at an exponential rate but the bias term appears due to the approximation error in the functional derivative. In general, the required number of samples M to control the approximation error ε can be exponentially large in the dimension of the response space Y. However, in most cases, measure transformation T depends on the distribution of the reward r(y)r(y) rather than the response y itself. Thus, since r(y)r(y) is one-dimensional, we can expect that the approximation error ε can be controlled with a moderate number of samples. We provide a dimension-independent bound on the approximation error for BoN-type objectives in the next section. 5.3 Properties of BoN-type Objectives In this section, we prove the concavity and smoothness of R[π]R[π] for BoN-type alignment algorithms. Then, we derive a dimension-independent bound on the approximation error in the functional derivative. Specifically, we consider the following general form: [π](y)=f(C[π](r(y)))⋅π(y), [π](y)=f(C[π](r(y)))·π(y), where f:ℝ→ℝf:R is monotonically increasing and LfL_f-Lipschitz continuous in [0,1][0,1]. This form includes BoN with f(z)=N⋅zN−1f(z)=N· z^N-1 and Best of Poisson (BoP) (khalaf2025inference) (drawing the number of samples N from a Poisson distribution with parameter λ>0λ>0) with f(z)=λ⋅eλ(z−1)f(z)=λ· e^λ(z-1) as special cases. Concavity and Smoothness For BoN-type algorithms, we have the following results: Lemma 5.4. R[π]=∫r(y)T[π](y)y⋅R[π]= r(y)T[π](y) y is concave in π. Lemma 5.5. R[π]=∫r(y)[π](y)y⋅R[π]= r(y)T[π](y) y is Lf⋅rmaxL_f· r_ -smooth relative to the KL-divergence. See Appendix H.3 and H.4 for the proof, respectively. These lemmas show that BoN-type objectives satisfy the assumptions required for the convergence analysis. Note that the concavity is preserved even if we consider the weighted sum or (smooth) min aggregation of multiple BoN-type objectives. Approximation Error Finally, we analyze the approximation error in the functional derivative for BoN-type objectives. The following lemma bounds the approximation error in the functional derivative when using empirical samples. Lemma 5.6. Let yjj=1M\y_j\_j=1^M be i.i.d. samples from πt _t and π^t=1M∑j=1Mδyj π_t= 1M _j=1^M _y_j be the empirical distribution. Then, the approximation error in the functional derivative can be bounded as [‖∂R∂π[π^t](y)−∂R∂π[πt](y)‖sp2]≤Lf2⋅rmax2M. [ Rπ[ π_t](y)- Rπ[ _t](y)_sp^2 ]≤ L_f^2· r_ ^2M. See Appendix H.5 for the proof. This lemma shows that the approximation error decreases at the rate of O(1/M)O(1/M) and is independent of the dimension of the response space Y. We also demonstrate in the numerical experiments that the non-linear GRPO algorithm works well with a moderate number of samples M. 6 Numerical Experiments (a) 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. (b) 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=4N=4) sampling. In this section, we empirically evaluate the effectiveness of the proposed framework and algorithm through extensive experiments on multiple models and tasks. We use TRL (vonwerra2020trl) library to implement (non-linear) GRPO algorithm. As discussed in the previous sections, the non-linear GRPO can be implemented by simply replacing the reward function with the approximated functional derivative. Throughout the experiments, we set the number of sampled responses M=8M=8, which is a default value in TRL library. See Appendix I for the detailed experimental settings and Appendix J for additional experimental results, including experiments with different hyperparameters and base models, as well as an analysis of the effect of sample size M. We also provide the implementation code in the supplementary material for reproducibility. 6.1 Length Reward We begin with a simple task using length reward to illustrate the difference between aligned policy and meta-aligned policy and see that the proposed method can obtain IAMA solution effectively. Specifically, we consider the following two rewards: r1(y) r_1(y) =− (|y−L1|Lmax)2,r2(y)=− (|y−L2|Lmax)2, =- ( y-L_1L_ )^2,r_2(y)=- ( y-L_2L_ )^2, where L1=50,L2=150L_1=50,L_2=150 are the target lengths and Lmax=256L_ =256 is the maximum length. These reward functions encourage the model to generate responses with lengths close to L1L_1 and L2L_2, respectively. This setup corresponds to the situation where some users prefer short responses while others prefer long and detailed ones. This is a similar setting to the one in Proposition 3.1. We use UltraFeedback dataset (cui2023ultrafeedback) as the context data and train Mistral 7B Instruct (jiang2023mistral7b) model to optimize the objective ℛ[π]=0.5R1[π]+0.5R2[π]−βKL[π∣πref]R[π]=0.5R_1[π]+0.5R_2[π]- [π _ref] with β=10−4β=10^-4 using the non-linear GRPO algorithm. Here, Ri[π]=y∼i[π][ri(y)]R_i[π]=E_y T_i[π] [r_i(y) ] with iT_i being BoN or soft BoN sampling. As a baseline, we also train the base model using standard GRPO to maximize the average expected reward π[(r1(y)+r2(y))/2]E_π [(r_1(y)+r_2(y))/2 ] - βKL[π∣πref] [π _ref]. Fig. 6 (left) shows the response length distributions generated by the IAMA models for BoN objective (N=2,4N=2,4) and Fig. 6 (right) shows those for soft BoN sampling (τ=0.09,0.1τ=0.09,0.1). Note that this is the base model distribution, i.e., the distribution without (soft) BoN sampling. We can see that the meta-aligned model produces diverse responses that cover both short and long lengths effectively while the baseline model generates only medium-length responses. In particular, larger N in BoN and smaller τ in soft BoN lead to stronger polarization of response lengths. We also experimented with BoN with N=8,16N=8,16, which are equal or greater than the number of sampled responses during training (M=8M=8) and confirmed that non-linear GRPO works well even in such cases. The results are provided in Appendix J.2 due to space limitations. 6.2 RLHF with helpfulness and harmlessness Next, we evaluate the proposed method on the RLHF task using h-rlhf dataset (bai2022training). We optimize a reproduced version of Alpaca 7B (taori2023alpaca; dai2024safe) model to maximize w1Rhelpful[π]+(1−w1)Rharmless[π]w_1R_helpful[π]+(1-w_1)R_harmless[π], where Rhelpful[π]R_helpful[π] and Rharmless[π]R_harmless[π] are BoN objectives with N=4N=4 for helpfulness and harmlessness rewards, respectively. We vary the weight w1w_1 to see how well the model can trade off between helpfulness and harmlessness. As reward models, we train Qwen3 4B (yang2025qwen3) based on the Bradley-Terry objective (bradley1952rank). Since evaluation with human raters can be costly, we use the larger reward models based on Qwen 32B as golden rewards to evaluate the aligned models following previous works (balashankarinfalign; eisenstein2023helping). To fairly compare the performance of different methods and weights, we fix the target KL divergence instead of fixing the regularization coefficient β. Then, we adopt log-space proportional controller (ziegler2019fine) to adjust the KL-regularization coefficient β during training. In this experiment, we set the target KL divergence as 0.10.1. Fig. 6 shows the evaluation results on helpfulness and harmlessness rewards with and without BoN sampling at inference time. Without BoN sampling, IAMA models cannot outperform standard aligned models since the standard alignment approach can directly optimize the weighted sum of expected rewards. On the other hand, with BoN sampling, IAMA models push the Pareto front forward compared to the standard alignment approach. This is because the meta-aligned model produces diverse responses covering different trade-offs between helpfulness and harmlessness, which allows the inference-time alignment to select high-reward responses effectively. We also provide the results for the case of N=8N=8 in Appendix J.3. 7 Conclusion In this work, we proposed inference-aware meta-alignment (IAMA), a novel framework for aligning LLMs to diverse user preferences via two-stage alignment procedures. We formulated IAMA as a non-linear optimization problem in the space of measures and developed the non-linear GRPO algorithm to solve the problem. We proved the convexity and smoothness of the BoN-type objective and provided convergence guarantees of the proposed algorithm. While implementation is simple, experimental results demonstrated the effectiveness of IAMA via non-linear GRPO in aligning LLMs to multiple reward functions simultaneously. Acknowledgements TS and RH were partially supported by JSPS KAKENHI (24K02905) and JST CREST (JPMJCR2115). This research is supported by the National Research Foundation, Singapore, Infocomm Media Development Authority under its Trust Tech Funding Initiative, and the Ministry of Digital Development and Information under the AI Visiting Professorship Programme (award number AIVP-2024-004). Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore, Infocomm Media Development Authority, and the Ministry of Digital Development and Information. Impact Statement This work aims to advance the methodology for aligning large language models to diverse and potentially conflicting human preferences with limited computational resources at inference time. By improving the efficiency and flexibility of inference-time alignment, the proposed approach may facilitate broader deployment of aligned language models in practical applications. At the same time, as with other alignment techniques, misuse or mis-specification of alignment criteria could lead to unintended behaviors, underscoring the importance of careful design and evaluation of alignment objectives. Overall, we do not foresee unique ethical risks beyond those commonly associated with research on large language model alignment. References Appendix A Notation For the reader’s convenience, we summarize the notation used in this paper in Table 1. Table 1: Summary of Notation Notation Description P Set of probability distributions over the response space π,πrefπ, _ref Policy and reference policy in P r(y)r(y) Reward function mapping response y to a scalar reward BoNN[π]BoN_N[π] Best-of-N (BoN) aligned model based on base policy π with N samples SoftBoNτ[π]SoftBoN_τ[π] Soft BoN aligned model based on base policy π with temperature parameter τ ∂R∂π[π](y) Rπ[π](y) Functional derivative of R at policy π evaluated at response y R⋅[π] R[π] ∂R∂π[π] Rπ[π] as a linear functional over P ⟨R⋅[π],π′⟩ R[π],π Action of linear functional R⋅[π] R[π] on π′π defined as ∫∂R∂π[π](y)⋅π′(y)y⋅ Rπ[π](y)·π (y) y KL[π∣πref]KL[π _ref] Kullback–Leibler divergence from πref _ref to π Dϕ[π∣μ]D_φ[π μ] Bregman divergence induced by convex functional ϕφ between μ and π πr _r Distribution of reward r(y)r(y) when y∼πy π Appendix B Detailed Discussion on Related Work In this section, we provide a more detailed discussion on related work about inference-time alignment and contrast our proposed method with prior studies. Recently, balashankarinfalign have proposed InfAlign-CTRL, which apply non-linear transformation to the reward function. However, they focus on 1) single reward function, 2) BoN-type alignment algorithms, and 3) win rate maximization objective. These assumptions are essential for deriving the method because the nonlinear transformation depends on the distribution of rewards, which can be determined from the fact that, in the case of a single reward and win-rate maximization, the distribution is uniform. Note that when there are multiple rewards, even if we consider win-rate maximization, the reward distribution is no longer uniform due to the presence of correlations. Several works (walder2025pass; bagirov2025best) have developped unbiased gradient estimator for Pass@k (BoN) objectives to align LLMs via policy gradient methods. However, these methods cannot handle inference-time alignment algorithms beyond BoN and require M≥NM≥ N samples in each gradient estimation step. Therefore, they are not applicable to large N settings or Best of Poisson (BoP), which require (potentially) unbounded number of samples. On the other hand, our proposed non-linear GRPO algorithm can handle a wide range of inference-time alignment algorithms and multiple reward functions. Moreover, even when restricted to the BoN setting, these prior studies did not investigate the theoretical properties of the optimization problem. We analyze the loss landscape of the optimization problem and establish its convexity. Based on this analysis, we provide convergence guarantees. This constitutes an additional, independent contribution. Appendix C Other Non-linear Optimization Problems Related to LLM Alignment Since our proposed non-linear GRPO algorithm is quite versatile, it can be applied to other non-linear optimization problems related to LLM alignment. In addition, non-linear GRPO can be regarded as a generalization of existing algorithms for specific non-linear optimization problems. Note that the existing works have not explicitly discussed the connection to non-linear optimization in the space of probability measures, which enables us to formulate a general algorithm and its theoretical analysis. Risk-aware Optimization Conditional Value at Risk (CVaR) is a widely used risk measure in finance and reinforcement learning, which is defined as CVaRα[π]:=y∼π[r(y)∣r(y)≤qα[π]], _α[π]:=E_y π [r(y) r(y)≤ q_α[π] ], where qα[π]q_α[π] is the α-quantile of the reward distribution when y∼πy π. This objective is non-linear in π due to the presence of the quantile function qα[π]q_α[π]. CPPO proposed by ying2022towards is a policy optimization algorithm for CVaR objectives, which can be seen as a special case of our non-linear GRPO algorithm. jiang2025risk also proposed a risk-sensitive RL algorithm for LLMs, using Entropic Risk Measure (ERM) defined as RERM[π]:=1βlogy∼π[eβr(y)]. R_ERM[π]:= 1β _y π [e^β r(y) ]. If β>0β>0, this objective encourages risk-seeking behavior, and is concave in π since log is concave and y∼π[eβr(y)]E_y π [e^β r(y) ] is linear in π. Diversity-promoting Optimization Diversity-promoting optimization aims to learn a policy that generates diverse outputs. For example, chen2025enhancing proposed to use determinantal point processes (DPPs) as a diversity-promoting regulalrizer: RDPP[π]:=y1,…,yN∼π[logdet(K(y1,…,yN))], R_DPP[π]:=E_y_1,…,y_N π [ (K(y_1,…,y_N)) ], where [K(y1,…,yN)]i,j=k(yi,yj)[K(y_1,…,y_N)]_i,j=k(y_i,y_j) is a kernel matrix measuring the similarity between outputs y1,…,yNy_1,…,y_N. Other possible regularizers include Rao’s quadratic entropy: RRao[π]:=−y,y′∼π[k(y,y′)]. R_Rao[π]:=-E_y,y π [k(y,y ) ]. If kernel k is positive semi-definite, the above objective is concave in π since this can be regarded as a quadratic form induced by k. Appendix D Practical Implementation of Non-linear GRPO Here, we describe a practical implementation of the non-linear GRPO algorithm. This is almost identical to the standard GRPO algorithm proposed in shao2024deepseekmath except for the use of the approximated functional derivative ∂R∂π[π^t](y) Rπ[ π_t](y) as the reward function. There exist various techniques to improve the stability and efficiency of GRPO that are not listed here. They can be directly applied to our non-linear GRPO as well. Please refer to the TRL (vonwerra2020trl) document for details. Algorithm 2 Practical Non-linear GRPO 1: Input: Initial policy π0 _0, reference policy πref _ref, reward functions rii=1m\r_i\_i=1^m, weights wii=1m\w_i\_i=1^m, regularization parameter β, step size η, number of iterations T 2: for t=0,1,…,T−1t=0,1,…,T-1 do 3: Sample M responses yjj=1M\y_j\_j=1^M from πt _t 4: Compute approximated functional derivative: r~t(y)=∂R∂π[π^t](y) r_t(y)= Rπ[ π_t](y). 5: Compute advantages: Aj=r~t(yj)−btstA_j= r_t(y_j)-b_ts_t where bt=1M∑j=1Mr~t(yj)b_t= 1M _j=1^M r_t(y_j) and st=1M∑j=1M(r~t(yj)−bt)2s_t= 1M _j=1^M( r_t(y_j)-b_t)^2 6: Define surrogate loss: ℒ^[π]=1M∑i=1Mmin[π(yi)πt(yi)⋅Ai,clipϵ(π(yi)πt(yi))⋅Ai]+βKL^[π∣πref] L[π]= 1M _i=1^M [ π(y_i) _t(y_i)· A_i,clip_ε ( π(y_i) _t(y_i) )· A_i ]+β KL[π _ref], where KL^[π∣πref]=1M∑i=1Mπref(yi)π(yi)−logπref(yi)π(yi)−1 KL[π _ref]= 1M _i=1^M _ref(y_i)π(y_i)- _ref(y_i)π(y_i)-1 7: Approximately solve the minimization problem minπℒ^[π] _π L[π] with optimizer (e.g., Adam) to obtain πt+1 _t+1 8: end for 9: Return: Final policy πT _T D.1 Computation of Functional Derivative In this section, we describe how to compute the functional derivatives for BoN and Soft BoN objectives with empirical distributions. For both cases, computational complexity depends on the number of samples M instead of the size N of BoN. BoN If we approximate the functional derivative ∂R∂π[π^t](y) Rπ[ π_t](y) with empirical distribution π^t π_t constructed from M samples yii=1M\y_i\_i=1^M drawn from πt _t, the computational bottleneck is sorting the rewards to compute the empirical quantiles. Thus, we can compute the functional derivative at these samples in O(MlogM)O(M M) time as follows: Algorithm 3 Linearized Reward Computation for BoN Objective 1: Input: Reward values rjj=1M\r_j\_j=1^M, BoN parameter N 2: Sort indices 1,…,M\1,…,M\ in ascending order of rewards so that r(1)≤r(2)≤⋯≤r(M)r_(1)≤ r_(2)≤·s≤ r_(M), where (i)(i) denotes the index of the i-th smallest reward. 3: Initialize cumulative sum c←0c← 0 4: Set reference reward rnext←r(M)r_next← r_(M) 5: Initialize r~j←0 r_j← 0 for all j=1,…,Mj=1,…,M 6: for i=M,M−1,…,1i=M,M-1,…,1 do 7: Compute reward gap Δr←rnext−r(i) r← r_next-r_(i) 8: Compute empirical CDF q←i/Mq← i/M 9: Update cumulative contribution c←c+NqN−1Δrc← c+Nq^N-1 r 10: Assign linearized reward r~(i)←−c r_(i)←-c 11: Update reference reward rnext←r(i)r_next← r_(i) 12: end for 13: Return: r~jj=1M\ r_j\_j=1^M Soft BoN If we approximate the functional derivative ∂R∂π[π^t](y) Rπ[ π_t](y) with empirical distribution π^t π_t constructed from M samples yii=1M\y_i\_i=1^M drawn from πt _t, we can compute the functional derivative in O(M)O(M) time as follows: Algorithm 4 Linearized Reward Computation for Soft BoN Objective 1: Input: Reward values rjj=1M\r_j\_j=1^M, temperature parameter τ 2: Compute normalization constant Z←1M∑j=1Mexp(rj/τ)Z← 1M _j=1^M (r_j/τ) 3: Compute softmax-weighted expected reward r¯←1M∑j=1Mrjexp(rj/τ)/Z r← 1M _j=1^Mr_j (r_j/τ)/Z 4: for j=1,…,Mj=1,…,M do 5: Compute linearized reward r~j←rjexp(rj/τ)/Z−exp(rj/τ)r¯/Z r_j← r_j (r_j/τ)/Z- (r_j/τ)\, r/Z 6: end for 7: Return: Linearized rewards r~jj=1M\ r_j\_j=1^M Appendix E Auxiliary Lemmas Lemma E.1. Let C[π](r)C[π](r) be the cumulative distribution function (CDF) defined as C[π](r)=∫[r(y)≤r]π(y)y⋅C[π](r)= 1[r(y)≤ r]π(y) y. Then, the functional derivative of C[π](r)C[π](r) is given by ∂C∂π[π](r)=[r(y)≤r] Cπ[π](r)= 1[r(y)≤ r]. Proof. For any π,π′∈π,π and ϵ∈ℝε , we have C[π+ϵ(π′−π)](r)−C[π](r) C[π+ε(π -π)](r)-C[π](r) =∫[r(y)≤r]⋅(π+ϵ(π′−π))(y)y⋅−∫[r(y)≤r]⋅π(y)y⋅ = 1[r(y)≤ r]·(π+ε(π -π))(y) y- 1[r(y)≤ r]·π(y) y =ϵ∫[r(y)≤r]⋅(π′−π)(y)y⋅ =ε 1[r(y)≤ r]·(π -π)(y) y and thus dC[π+ϵ(π′−π)](r)dϵ|ϵ=0 dC[π+ε(π -π)](r)dε |_ε=0 =∫[r(y)≤r]⋅(π′−π)(y)y⋅. = 1[r(y)≤ r]·(π -π)(y) y. Hence, we have ∂C∂π[π](r)=[r(y)≤r] Cπ[π](r)= 1[r(y)≤ r]. ∎ Lemma E.2. For a reward function r(y)r(y) and a base policy π, assume that the distribution of r(y)r(y) when y∼πy π is continuous. Then, the BoN policy BoNN[π]BoN_N[π] satisfies BoNN[π](y)=N⋅C[π](r(y))N−1⋅π(y) _N[π](y)=N· C[π](r(y))^N-1·π(y) Proof. See Lemma 5 of balashankarinfalign. ∎ Lemma E.3 (Three-point inequality). Let ,ψ,ϕG,ψ,φ be convex functionals. Then, for any μ,νμ,ν and ν¯:=argmin[ν]+βDψ(ν∣πref)+LDϕ(ν∣μ)∈ ν:= [ν]+β D_ψ(ν _ref)+LD_φ(ν μ) , we have [ν¯]+βDψ(ν¯∣πref)+LDϕ(ν¯∣μ) [ ν]+β D_ψ( ν _ref)+LD_φ( ν μ) ≤[ν]+βDψ(ν∣πref)+LDϕ(ν∣μ) [ν]+β D_ψ(ν _ref)+LD_φ(ν μ) −LDϕ(ν∣ν¯)−βDψ(ν∣ν¯). -LD_φ(ν ν)-β D_ψ(ν ν). Proof. This is an extension of Lemma 3 in frankowski2022mirror and we follow the same proof strategy. Let f=[ν]+βDψ(ν∣πref)+LDϕ(ν∣μ)f=G[ν]+β D_ψ(ν _ref)+LD_φ(ν μ). Note that f is convex since G, DψD_ψ, and DϕD_φ are convex. Then, we have Df(ν∣ν¯) D_f(ν ν) =D(ν∣ν¯)+βDψ(ν∣ν¯)+LDϕ(ν∣ν¯) =D_G(ν ν)+β D_ψ(ν ν)+LD_φ(ν ν) ≥βDψ(ν∣ν¯)+LDϕ(ν∣ν¯) ≥β D_ψ(ν ν)+LD_φ(ν ν) from the non-negativity of Bregman divergence. By the optimality of ν¯ ν, we have Df(ν∣ν¯)=f(ν)−f(ν¯)−⟨∂f∂π[ν¯],ν−ν¯⟩≤f(ν)−f(ν¯)D_f(ν ν)=f(ν)-f( ν)- fπ[ ν],ν- ν ≤ f(ν)-f( ν). Thus, we have f(ν)−f(ν¯)≥βDψ(ν∣ν¯)+LDϕ(ν∣ν¯). f(ν)-f( ν)≥β D_ψ(ν ν)+LD_φ(ν ν). Rearranging the above inequality, we obtain the desired result. ∎ Appendix F Proof for Propositions 4.2 and 4.3 Case 1: BoN BoN defines the aligned model as BoNN[π](y)=f(C[π](r(y)))⋅π(y), _N[π](y)=f(C[π](r(y)))·π(y), where f(r)=N⋅rN−1f(r)=N· r^N-1. Thus, the objective function is given by R(π) R(π) =y∼BoNN[π][r(y)] =E_y BoN_N[π] [r(y) ] =∫r(y)⋅f(C[π](r(y)))⋅π(y)y⋅ = r(y)· f(C[π](r(y)))·π(y) y =∫0rmaxr⋅f(C[πr](r))⋅πr(r)r⋅ = _0^r_ r· f(C[ _r](r))· _r(r) r = [r⋅F(C[πr](r))]0rmax−∫0rmaxF(C[πr](r))r⋅ = [r· F(C[ _r](r))]_0^r_ - _0^r_ F(C[ _r](r)) r where πr _r is the distribution of r(y)r(y) when y∼πy π and F is an antiderivative of f. Here, the last inequality follows from integration by parts. Then, the functional derivative can be computed as R[π]π(z) R[π]π(z) =−∫f(C[πr](r))C[πr]π(z)r⋅ =- f(C[ _r](r)) C[ _r]π(z) r =−∫f(C[πr](r))⋅r≥r(z)r⋅ =- f(C[ _r](r))· 1_r≥ r(z) r =−∫r(z)rmaxf(C[πr](r))r⋅. =- _r(z)^r_ f(C[ _r](r)) r. This completes the proof. Case 2: Soft BoN Let λ=1/τλ=1/τ. Soft BoN defines the aligned model as SoftBoNτ[π](y)=exp(λr(y))⋅π(y)∫exp(λr(z))⋅π(z)z⋅. _τ[π](y)= (λ r(y))·π(y) (λ r(z))·π(z) z. Thus, the objective function is given by R(π) R(π) =y∼SoftBoNλ[r(y)] =E_y _λ [r(y) ] =∫r(y)⋅exp(λr(y))⋅π(y)∫exp(λr(z))⋅π(z)z⋅y⋅ = r(y)· (λ r(y))·π(y) (λ r(z))·π(z) z y =dλlog(∫exp(λr(z))⋅π(z)z⋅). = ddλ ( (λ r(z))·π(z) z ). Then, the functional derivative can be computed as ∂R∂π[π](y) Rπ[π](y) =dλ (∂πlog (∫exp(λr(z))⋅π(z)z⋅)) = ddλ ( π ( (λ r(z))·π(z) z)) =dλ (exp(λr(y))∫exp(λr(z))⋅π(z)z⋅) = ddλ ( (λ r(y)) (λ r(z))·π(z) z) =r(y)exp(λr(y))∫exp(λr(z))⋅π(z)z⋅−exp(λr(y))⋅∫r(z)exp(λr(z))⋅π(z)z⋅ (∫exp(λr(z))⋅π(z)z⋅)2. = r(y) (λ r(y)) (λ r(z))·π(z) z- (λ r(y))· r(z) (λ r(z))·π(z) z ( (λ r(z))·π(z) z)^2. This completes the proof. Appendix G Proof for Section 3 G.1 Proof of Proposition 3.1 Let Y1,…,YN∼i.i.d.πY_1,…,Y_N i.i.d. π be N i.i.d. samples from π. Since r1r_1 is strictly decreasing and r2r_2 is strictly increasing, the samples maximizing r1r_1 and r2r_2 are the minimum and maximum order statistics, Y(1) Y_(1) :=mink≤NYk, := _k≤ NY_k, Y(N) Y_(N) :=maxk≤NYk. := _k≤ NY_k. Hence, the objective is R[π]:=[−Y(1)2]+[−(1−Y(N))2]. R[π]:=E[-Y_(1)^2]+E[-(1-Y_(N))^2]. For any y∈[0,1]y∈[0,1], Prob(Y(1)>y) (Y_(1)>y ) =(1−C[π](y))N, =(1-C[π](y))^N, Prob(Y(N)≤y) (Y_(N)≤ y ) =C[π](y)N. =C[π](y)^N. Using [X2]=∫012tℙ(X>t)tE[X^2]= _0^12t\,P(X>t)\,dt for X∈[0,1]X∈[0,1] and a change of variables, R[π]=−2∫01 [y(1−C[π](y))N+(1−y)C[π](y)N]y. R[π]=-2 _0^1 [y(1-C[π](y))^N+(1-y)C[π](y)^N]\,dy. Thus maximizing R[π]R[π] is equivalent to minimizing I(π):=∫01ϕy(C[π](y))y, I(π):= _0^1 _y(C[π](y))\,dy, where ϕy(p):=y(1−p)N+(1−y)pN _y(p):=y(1-p)^N+(1-y)p^N. For fixed y∈(0,1)y∈(0,1), we have ϕy′(p)=N(N−1)[y(1−p)N−2+(1−y)pN−2]>0 _y (p)=N(N-1) [y(1-p)^N-2+(1-y)p^N-2 ]>0 for p∈(0,1)p∈(0,1). Therefore, ϕy _y is strictly convex on (0,1)(0,1) and has a unique minimizer satisfying 0 0 =ϕy′(p)=−Ny(1−p)N−1+N(1−y)pN−1 = _y (p)=-Ny(1-p)^N-1+N(1-y)p^N-1 ⟺(p1−p)N−1=y1−y. ( p1-p )^N-1= y1-y. That is, the minimizer is given by C∗(y)=yαyα+(1−y)α,α:=1N−1. C^*(y)= y^αy^α+(1-y)^α, α:= 1N-1. Moreover, C∗(y)C^*(y) is strictly increasing with C∗(0)=0C^*(0)=0 and C∗(1)=1C^*(1)=1, which implies that π∗:=dC∗dy _*:= dC^*dy is in P and a minimizer of I(π)I(π). This completes the proof. Appendix H Proofs for Section 5 H.1 Proof of Theorem 5.2 For notational simplicity, we define F[π] F[π] :=−R[π], :=-R[π], ℒ[π] [π] =F[π]+βKL[π∣πref]. =F[π]+ [π _ref]. Note that F is convex and L-smooth since R is concave and L-smooth. From the L-smoothness of F[π]F[π], we have ℒ[πt+1]−ℒ[πt]≤⟨F⋅[πt],πt+1−πt⟩+LKL[πt+1∣πt]+βKL[πt+1∣πref]−βKL[πt∣πref]. [ _t+1]-L[ _t]≤ F[ _t], _t+1- _t +LKL[ _t+1 _t]+ [ _t+1 _ref]- [ _t _ref]. Applying Lemma E.3 to the convex functional ⟨F⋅[πt],π−πt⟩ F[ _t],π- _t , we obtain ⟨F⋅[πt],πt+1−πt⟩+LKL[πt+1∣πt]+βKL[πt+1∣πref] F[ _t], _t+1- _t +LKL[ _t+1 _t]+ [ _t+1 _ref] ≤⟨F⋅[πt],π−πt⟩+LKL[π∣πt]+βKL[π∣πref] ≤ F[ _t],π- _t +LKL[π _t]+ [π _ref] −(L+β)KL[π∣πt+1] -(L+β)KL[π _t+1] for any π∈π . Combining the above two inequalities, the following inequality holds: ℒ[πt+1]−ℒ[πt] [ _t+1]-L[ _t] ≤⟨F⋅[πt],π−πt⟩ ≤ F[ _t],π- _t +LKL[π∣πt] +LKL[π _t] +βKL[π∣πref] + [π _ref] −(L+β)KL[π∣πt+1] -(L+β)KL[π _t+1] −βKL[πt∣πref]. - [ _t _ref]. By setting π=πtπ= _t, we have ℒ[πt+1]−ℒ[πt] [ _t+1]-L[ _t] ≤−(L+β)KL[πt∣πt+1]≤0, ≤-(L+β)KL[ _t _t+1]≤ 0, which implies the monotonicity of ℒ[πt]L[ _t]. On the other hand, by setting π=π∗π= _*, we have ℒ[πt+1]−ℒ[πt] [ _t+1]-L[ _t] ≤⟨F⋅[πt],π∗−πt⟩ ≤ F[ _t], _*- _t +LKL[π∗∣πt] +LKL[ _* _t] +βKL[π∗∣πref] + [ _* _ref] −(L+β)KL[π∗∣πt+1] -(L+β)KL[ _* _t+1] −βKL[πt∣πref]. - [ _t _ref]. Since F is convex, we have ⟨F⋅[πt],π∗−πt⟩≤F(π∗)−F(πt). F[ _t], _*- _t ≤ F( _*)-F( _t). Thus, we obtain ℒ[πt+1]−ℒ[πt] [ _t+1]-L[ _t] ≤F(π∗)−F(πt) ≤ F( _*)-F( _t) +LKL[π∗∣πt] +LKL[ _* _t] +βKL[π∗∣πref] + [ _* _ref] −(L+β)KL[π∗∣πt+1] -(L+β)KL[ _* _t+1] −βKL[πt∣πref] - [ _t _ref] =ℒ[π∗]−ℒ[πt] =L[ _*]-L[ _t] +LKL[π∗∣πt] +LKL[ _* _t] −(L+β)KL[π∗∣πt+1]. -(L+β)KL[ _* _t+1]. Rearranging the above inequality, we have ℒ[πt+1] [ _t+1] ≤ℒ[π∗]+LKL[π∗∣πt]−(L+β)KL[π∗∣πt+1]. [ _*]+LKL[ _* _t]-(L+β)KL[ _* _t+1]. Summing the above inequality from t=0t=0 to T−1T-1, we obtain ∑t=1T (L+βL)tℒ[πt] _t=1^T ( L+βL)^tL[ _t] ≤∑t=1T (L+βL)tℒ[π∗]+L∑t=1T (L+βL)tKL[π∗∣πt−1]−(L+β)∑t=1T (L+βL)tKL[π∗∣πt] ≤ _t=1^T ( L+βL)^tL[ _*]+L _t=1^T ( L+βL)^tKL[ _* _t-1]-(L+β) _t=1^T ( L+βL)^tKL[ _* _t] =∑t=1T (L+βL)tℒ[π∗]+(L+β) [∑t=1T (L+βL)t−1KL[π∗∣πt−1]−∑t=1T (L+βL)tKL[π∗∣πt]] = _t=1^T ( L+βL)^tL[ _*]+(L+β) [ _t=1^T ( L+βL)^t-1KL[ _* _t-1]- _t=1^T ( L+βL)^tKL[ _* _t]] =∑t=1T (L+βL)tℒ[π∗]+(L+β) [∑t=0T−1 (L+βL)tKL[π∗∣πt]−∑t=1T (L+βL)tKL[π∗∣πt]] = _t=1^T ( L+βL)^tL[ _*]+(L+β) [ _t=0^T-1 ( L+βL)^tKL[ _* _t]- _t=1^T ( L+βL)^tKL[ _* _t]] =∑t=1T (L+βL)tℒ[π∗]+(L+β) [KL[π∗∣π0]− (L+βL)TKL[π∗∣πT]] = _t=1^T ( L+βL)^tL[ _*]+(L+β) [KL[ _* _0]- ( L+βL)^TKL[ _* _T]] ≤∑t=1T (L+βL)tℒ[π∗]+(L+β)KL[π∗∣π0]. ≤ _t=1^T ( L+βL)^tL[ _*]+(L+β)KL[ _* _0]. Dividing both sides by ∑t=1T (L+βL)t _t=1^T ( L+βL)^t, we have ℒ[πT]−ℒ[π∗] [ _T]-L[ _*] ≤∑t=1T (L+βL)tℒ[πt]∑t=1T (L+βL)t−ℒ[π∗] ≤ _t=1^T ( L+βL)^tL[ _t] _t=1^T ( L+βL)^t-L[ _*] ≤(L+β)KL[π∗∣π0]∑t=1T (L+βL)t ≤ (L+β)KL[ _* _0] _t=1^T ( L+βL)^t =(L+β)KL[π∗∣π0](L+β)L⋅((L+β)/L)T−1(L+β)/L−1 = (L+β)KL[ _* _0] (L+β)L· ((L+β)/L)^T-1(L+β)/L-1 =β⋅KL[π∗∣π0]((L+β)/L)T−1. = β·KL[ _* _0]((L+β)/L)^T-1. The first inequality follows from the monotonicity of ℒ[πt]L[ _t]. This completes the proof. H.2 Proof of Theorem 5.3 Let the empirical gradient be g^t:=F⋅[π^t] g_t:= F[ π_t] and define the estimation error et:=g^t−F⋅[πt]e_t:= g_t- F[ _t]. From the L-smoothness of F, we have ℒ[πt+1]−ℒ[πt] [ _t+1]-L[ _t] ≤F⋅[πt](πt+1−πt)+LKL[πt+1∣πt]+βKL[πt+1∣πref]−βKL[πt∣πref] ≤ F[ _t]( _t+1- _t)+L\,KL[ _t+1 _t]+β\,KL[ _t+1 _ref]-β\,KL[ _t _ref] =g^t(πt+1−πt)+LKL[πt+1∣πt]+βKL[πt+1∣πref]−βKL[πt∣πref]−et(πt+1−πt). = g_t( _t+1- _t)+L\,KL[ _t+1 _t]+β\,KL[ _t+1 _ref]-β\,KL[ _t _ref]-e_t( _t+1- _t). Due to the inexact optimization, there exists a residual term rtr_t defined as rt:=g^t+1ηlogπt+1πt+βlogπt+1πref. r_t:= g_t+ 1η _t+1 _t+β _t+1 _ref. Note that rtr_t is constant if the optimization is exact. This implies that πt+1 _t+1 is the optimal solution to the following optimization problem: minπ∈⟨g^t−rt,π−πt⟩+1ηKL[π∣πt]+βKL[π∣πref]. _π g_t-r_t,π- _t + 1ηKL[π _t]+ [π _ref]. Applying Lemma E.3 to the convex functional ⟨g^t−rt,π−πt⟩ g_t-r_t,π- _t , we obtain ⟨g^t−rt,πt+1−πt⟩+LKL[πt+1∣πt]+βKL[πt+1∣πref] g_t-r_t, _t+1- _t +LKL[ _t+1 _t]+ [ _t+1 _ref] =⟨g^t,πt+1−πt⟩+LKL[πt+1∣πt]+βKL[πt+1∣πref]−⟨rt,πt+1−πt⟩ = g_t, _t+1- _t +LKL[ _t+1 _t]+ [ _t+1 _ref]- r_t, _t+1- _t ≤⟨g^t−rt,π∗−πt⟩+LKL[π∗∣πt]+βKL[π∗∣πref]−(L+β)KL[π∗∣πt+1] ≤ g_t-r_t, _*- _t +LKL[ _* _t]+ [ _* _ref]-(L+β)KL[ _* _t+1] =⟨g^t,π∗−πt⟩+LKL[π∗∣πt]+βKL[π∗∣πref]−(L+β)KL[π∗∣πt+1]−⟨rt,π∗−πt⟩. = g_t, _*- _t +LKL[ _* _t]+ [ _* _ref]-(L+β)KL[ _* _t+1]- r_t, _*- _t . Combining the above two inequalities, we have ℒ[πt+1]−ℒ[πt] [ _t+1]-L[ _t] ≤g^t(π∗−πt)+LKL[π∗∣πt]+βKL[π∗∣πref] ≤ g_t( _*- _t)+LKL[ _* _t]+ [ _* _ref] −(L+β)KL[π∗∣πt+1]−βKL[πt∣πref]−⟨et,πt+1−πt⟩−⟨rt,π∗−πt+1⟩. -(L+β)KL[ _* _t+1]- [ _t _ref]- e_t, _t+1- _t - r_t, _*- _t+1 . (2) By convexity of F, F̣[πt](π∗−πt)≤F[π∗]−F[πt], \@text@daccentF[ _t]( _*- _t)≤ F[ _*]-F[ _t], and thus g^t(π∗−πt)=F⋅[πt](π∗−πt)+⟨et,π∗−πt⟩≤F[π∗]−F[πt]+⟨et,π∗−πt⟩. g_t( _*- _t)= F[ _t]( _*- _t)+ e_t, _*- _t ≤ F[ _*]-F[ _t]+ e_t, _*- _t . Substituting this into (2), we get ℒ[πt+1]−ℒ[πt] [ _t+1]-L[ _t] ≤F[π∗]−F[πt]+LKL[π∗∣πt]+βKL[π∗∣πref] ≤ F[ _*]-F[ _t]+LKL[ _* _t]+ [ _* _ref] −(L+β)KL[π∗∣πt+1]−βKL[πt∣πref]+⟨et,π∗−πt+1⟩−⟨rt,π∗−πt+1⟩. -(L+β)KL[ _* _t+1]- [ _t _ref]+ e_t, _*- _t+1 - r_t, _*- _t+1 . Rearranging the above inequality, we have ℒ[πt+1] [ _t+1] ≤ℒ[π∗]+LKL[π∗∣πt]−(L+β)KL[π∗∣πt+1]+⟨et,π∗−πt+1⟩−⟨rt,π∗−πt+1⟩. [ _*]+LKL[ _* _t]-(L+β)KL[ _* _t+1]+ e_t, _*- _t+1 - r_t, _*- _t+1 . (3) Let c=(supyet(y)+infyet(y))/2c=( _ye_t(y)+ _ye_t(y))/2. By Pinsker’s inequality and Hölder’s inequality, ⟨et,π∗−πt+1⟩ e_t, _*- _t+1 =⟨et−c,π∗−πt+1⟩ = e_t-c, _*- _t+1 ≤‖et−c‖∞‖π∗−πt+1‖1 ≤\|e_t-c\|_∞\| _*- _t+1\|_1 =‖et‖sp⋅‖π∗−πt+1‖1 = e_t_sp·\| _*- _t+1\|_1 ≤‖et‖sp⋅2KL[π∗∣πt+1]. ≤ e_t_sp· 2\,KL[ _* _t+1]. The second equality follows from the fact ‖et−c‖∞=max(supyet(y)−c,c−infyet(y))=(supyet(y)−infyet(y))/2=‖et‖sp e_t-c_∞= ( _ye_t(y)-c,c- _ye_t(y))=( _ye_t(y)- _ye_t(y))/2= e_t_sp. Applying Young’s inequality (ab≤a22α+αb22ab≤ a^22α+ α b^22), we obtain ⟨et,π∗−πt+1⟩≤‖et‖sp22α+αKL[π∗∣πt+1] e_t, _*- _t+1 ≤ \|e_t\|_sp^22α+ [ _* _t+1] for any α>0α>0. Similarly, we have −⟨rt,π∗−πt+1⟩≤‖rt‖sp22α+αKL[π∗∣πt+1] - r_t, _*- _t+1 ≤ \|r_t\|_sp^22α+ [ _* _t+1] for any α>0α>0. Substituting these into (3), we obtain ℒ[πt+1]≤ℒ[π∗]+LKL[π∗∣πt]−(L+β−2α)KL[π∗∣πt+1]+‖et‖sp2+‖rt‖sp22α. [ _t+1] [ _*]+L\,KL[ _* _t]-(L+β-2α)\,KL[ _* _t+1]+ \|e_t\|_sp^2+\|r_t\|_sp^22α. (4) Setting α=β/4α=β/4 (so that β−2α=β/2β-2α=β/2), we have ℒ[πt+1]≤ℒ[π∗]+LKL[π∗∣πt]−(L+β/2)KL[π∗∣πt+1]+2(‖et‖sp2+‖rt‖sp2)β. [ _t+1] [ _*]+LKL[ _* _t]-(L+β/2)\,KL[ _* _t+1]+ 2(\|e_t\|_sp^2+\|r_t\|_sp^2)β. Let μ:=L+β/2μ:=L+β/2. Summing over t=1,…,Tt=1,…,T with weights (μ/L)t(μ/L)^t yields ∑t=1T(μL)tℒ[πt] _t=1^T ( μL )^tL[ _t] ≤∑t=1T(μL)tℒ[π∗]+μ[KL[π∗∣π0]−(μL)TKL[π∗∣πT]]+∑t=1T(μL)t2(‖et−1‖sp2+‖rt−1‖sp2)β. ≤ _t=1^T ( μL )^tL[ _*]+μ [KL[ _* _0]- ( μL )^TKL[ _* _T] ]+ _t=1^T ( μL )^t 2(\|e_t-1\|_sp^2+\|r_t-1\|_sp^2)β. Let ST:=∑t=1T(μ/L)t=μL⋅(μ/L)T−1(μ/L)−1S_T:= _t=1^T(μ/L)^t= μL· (μ/L)^T-1(μ/L)-1. Then ∑t=1T(μL)tℒ[πt]ST−ℒ[π∗] _t=1^T ( μL )^tL[ _t]S_T-L[ _*] =∑t=1TProb(t^=t)ℒ[πt]−ℒ[π∗] = _t=1^TProb ( t=t )L[ _t]-L[ _*] ≤μKL[π∗∣π0]ST+2β⋅∑t=1T(μ/L)t(‖et−1‖sp2+‖rt−1‖sp2)ST. ≤ μ\,KL[ _* _0]S_T+ 2β· _t=1^T(μ/L)^t(\|e_t-1\|_sp^2+\|r_t-1\|_sp^2)S_T. Substituting μ=L+β/2μ=L+β/2 gives μKL[π∗∣π0]ST μ\,KL[ _* _0]S_T =(L+β/2)KL[π∗∣π0](L+β/2)L⋅((L+β/2)/L)T−1((L+β/2)/L)−1 = (L+β/2)\,KL[ _* _0] (L+β/2)L· ((L+β/2)/L)^T-1((L+β/2)/L)-1 =(β/2)⋅KL[π∗∣π0]((L+β/2)/L)T−1. = (β/2)·KL[ _* _0]((L+β/2)/L)^T-1. From the assumptions [‖et‖sp2]≤εE[\|e_t\|_sp^2]≤ and [‖rt‖sp2]≤δE[\|r_t\|_sp^2]≤δ, we have [∑t=1T(μ/L)t(‖et−1‖sp2+‖rt−1‖sp2)ST] [ _t=1^T(μ/L)^t(\|e_t-1\|_sp^2+\|r_t-1\|_sp^2)S_T ] ≤∑t=1T(μ/L)t(ε+δ)ST=ε+δ. ≤ _t=1^T(μ/L)^t( +δ)S_T= +δ. Thus, we have [ℒ[πt^]−ℒ[π∗]] [L[ _ t]-L[ _*] ] ≤(β/2)⋅KL[π∗∣π0]((L+β/2)/L)T−1+2(ε+δ)β. ≤ (β/2)·KL[ _* _0]((L+β/2)/L)^T-1+ 2( +δ)β. H.3 Proof of Lemma 5.4 Proof. Let F be the antiderivative of f. Since f is non-decreasing, F is convex. By integration by parts, we have R[π] R[π] =∫r(y)f(C[π](r(y)))π(y)y⋅ = r(y)f(C[π](r(y)))π(y) y =∫0∞rf(C[π](r))πr(r)r⋅ = _0^∞rf(C[π](r)) _r(r) r =∫0∞r⋅drF(C[π](r))r⋅ = _0^∞r· ddrF(C[π](r)) r =[rF[π](r)]0rmax−∫0rmaxF(C[π](r))r⋅ = [rF[π](r) ]_0^r_ - _0^r_ F(C[π](r)) r =rmaxF(1)−∫0rmaxF(C[π](r))r⋅. =r_ F(1)- _0^r_ F(C[π](r)) r. Thus, it is suffice to show that P[π]:=∫0rmaxF(C[π](r))r⋅P[π]:= _0^r_ F(C[π](r)) r is convex in π. For any π,π′∈π,π and λ∈[0,1]λ∈[0,1], we have P[(λπ+(1−λ)π′)] P[(λπ+(1-λ)π )] =∫0rmaxF(C[(λπ+(1−λ)π′)](r))r⋅ = _0^r_ F(C[(λπ+(1-λ)π )](r)) r =∫0rmaxF(λC[π](r)+(1−λ)C[π′](r))r⋅ = _0^r_ F (λ C[π](r)+(1-λ)C[π ](r) ) r ≤∫0rmax(λF(C[π](r))+(1−λ)F(C[π′](r)))r⋅ ≤ _0^r_ (λ F(C[π](r))+(1-λ)F(C[π ](r)) ) r =λP[π]+(1−λ)P[π′]. =λ P[π]+(1-λ)P[π ]. Here, the inequality follows from the convexity of F. This completes the proof. ∎ H.4 Proof of Lemma 5.5 Proof. From the integration by parts, we have R(π) R(π) =∫0rmaxr(y)f(C[π](r))π⋅(y) = _0^r_ r(y)f(C[π](r)) π(y) =[rF(C[π](r))]0rmax−∫0rmaxF(C[π](r))r⋅, = [rF(C[π](r)) ]_0^r_ - _0^r_ F(C[π](r)) r, =rmaxF(1)−∫0rmaxF(C[π](r))r⋅, =r_ F(1)- _0^r_ F(C[π](r)) r, where F is the antiderivative of f. Since the first term is constant, it suffices to show that ∫0rmaxF(C[π](r))r⋅ _0^r_ F(C[π](r)) r is Lf⋅rmaxL_f· r_ -smooth relative to the KL-divergence. First, we prove F(C[π](r))F(C[π](r)) is LfL_f-smooth relative to the KL-divergence for all r. For any π,π′π,π , we have F(C[π′](r))−F(C[π](r))−∫f(C[π](r))⋅r(y)≤r⋅(π′(y)−π(y))y⋅ F(C[π ](r))-F(C[π](r))- f(C[π](r))· 1_r(y)≤ r·(π (y)-π(y)) y =F(C[π′](r))−F(C[π](r))−f(C[π](r))⋅(C[π′](r)−C[π](r)) =F(C[π ](r))-F(C[π](r))-f(C[π](r))·(C[π ](r)-C[π](r)) ≤Lf2(C[π′](r)−C[π](r))2 ≤ L_f2(C[π ](r)-C[π](r))^2 =Lf2(∫(π′(y)−π(y))⋅r(y)≤ry⋅)2 = L_f2 ( (π (y)-π(y))· 1_r(y)≤ r y )^2 ≤Lf2(∫|π′(y)−π(y)|⋅r(y)≤ry⋅)2 ≤ L_f2 ( π (y)-π(y)· 1_r(y)≤ r y )^2 ≤Lf2(∫|π′(y)−π(y)|y⋅)2 ≤ L_f2 ( π (y)-π(y) y )^2 ≤2LfdTV(π,π′)2 ≤ 2L_fd_TV(π,π )^2 ≤LfKL[π′∣π]. ≤ L_fKL[π π]. Here, the first inequality follows from the LfL_f-smoothness of F, and the last inequality follows from Pinsker’s inequality. Thus, F(C[π](r))F(C[π](r)) is LfL_f-smooth relative to the KL-divergence for all r. Next, by integrating both sides over r∈[0,rmax]r∈[0,r_ ], we have ∫0rmaxF(C[π′](r))r⋅−∫0rmaxF(C[π](r))r⋅−∫0rmax∫f(C[π](r))⋅r(y)≤r⋅(π′(y)−π(y))y⋅r⋅ _0^r_ F(C[π ](r)) r- _0^r_ F(C[π](r)) r- _0^r_ f(C[π](r))· 1_r(y)≤ r·(π (y)-π(y)) y r ≤Lf⋅rmax⋅KL[π′∣π]. ≤ L_f· r_ ·KL[π π]. Thus, ∫0rmaxF(C[π](r))r⋅ _0^r_ F(C[π](r)) r is Lf⋅rmaxL_f· r_ -smooth relative to the KL-divergence. ∎ H.5 Proof of Lemma 5.6 Since the span seminorm satisfies ‖f‖sp≤‖f‖∞ f_sp≤ f_∞, it suffices to bound the infinity norm. The first-order variation can be computed as ∂R∂π[π](y) Rπ[π](y) =−∫r(y)rmaxf(C[π](r))r⋅. =- _r(y)^r_ f(C[π](r)) r. Thus, the approximation error can be bounded as follows: |∂R∂π[π^t](y)−∂R∂π[πt](y)| Rπ[ π_t](y)- Rπ[ _t](y) ≤∫0rmax|f(C[π^t](r))−f(C[πt](r))|r⋅ ≤ _0^r_ f(C[ π_t](r))-f(C[ _t](r)) r ≤Lfrmax‖C[π^t]−C[πt]‖∞ ≤ L_fr_ C[ π_t]-C[ _t]_∞ Let Z=‖C[π^t]−C[πt]‖∞Z= C[ π_t]-C[ _t]_∞ be a random variable. Using Dvoretzky-Kiefer-Wolfowitz (DKW) inequality, we have Prob(Z≥ε)≤2e−2Mε2, (Z≥ )≤ 2e^-2M ^2, which implies [Z2] [Z^2 ] =2∫01εProb(Z≥ε)ε⋅≤2∫01ε⋅2e−2Mε2ε⋅ =2 _0^1 (Z≥ ) ≤ 2 _0^1 · 2e^-2M ^2 ≤1M. ≤ 1M. Thus, we have [‖∂R∂π[π^t](y)−∂R∂π[πt](y)‖∞2]≤Lf2rmax2M, [ Rπ[ π_t](y)- Rπ[ _t](y)_∞^2 ]≤ L_f^2r_ ^2M, which completes the proof. Appendix I Experimental Details Here, we provide additional details regarding our experiments. Hyperparameters used in our experiments are summarized in Table 2. Compute Resources Our experiments were conducted on Intel(R) Xeon(R) Silver 4316 CPU @ 2.30GHz and 8 NVIDIA A100-SXM4-80GB GPUs. Length Reward Experiments We set the number of responses M=8M=8, learning rate 1e−61e-6, batch size 4848, number of iterations T=3000T=3000, and β=1e−4β=1e-4. We use constant learning rate schedule and turn off reward scaling in GRPO to stabilize training. Other hyperparameters are default values from the TRL library (vonwerra2020trl). H-RLHF Experiments We set the number of responses M=8M=8, learning rate 1e−61e-6, batch size 4848, number of iterations T=1000T=1000, and the initial β=1e−2β=1e-2. We adopt log-space proportional controller (ziegler2019fine) to adjust the KL-regularization coefficient β during training as follows: et e_t =clip (KL[πt∣πref]−KLtargetKLtarget,[−0.2,0.2]), =clip ( KL[ _t _ref]-KL_targetKL_target,[-0.2,0.2]), βt+1 _t+1 =βt⋅(1+0.1⋅et) = _t·(1+0.1· e_t) with target KL divergence KLtarget=0.1KL_target=0.1. To stabilize training, we turn off reward scaling in GRPO. Other hyperparameters are default values from the TRL library (vonwerra2020trl). For evaluation, we use the first 1000 samples from the validation set of the H-RLHF dataset to reduce the computational cost. Table 2: Hyperparameters used in our experiments. Hyperparameter Length Reward Experiments H-RLHF Experiments Number of responses M 8 8 Learning rate 10−610^-6 10−610^-6 Per device batch size 8 8 Number of iterations T 3000 1000 Initial KL coefficient β 10−410^-4 10−210^-2 Target KL KLtargetKL_target N/A 10−110^-1 Learning rate scheduler Constant Linear bf16 Training True True Optimizer AdamW AdamW Clipping ϵε 0.20.2 0.20.2 Appendix J Additional Experimental Results J.1 Effect of the Number of Responses M To investigate the effect of the number of responses M in the estimation of linearized loss, we conduct additional synthetic experiments. Since it is difficult to obtain true expected reward with inference-time alignment in general, we consider tractable setup where the reward distribution is a beta distribution on [0,1][0,1] and the inference-time alignment is performed by BoN with N responses. Let πα(y) _α(y) be a beta distribution with shape parameters (α,1)(α,1) on [0,1][0,1]. Then, π1(y) _1(y) is equivalent to the uniform distribution on [0,1][0,1]. As is well known, BoNN[πα](y)BoN_N[ _α](y) is a beta distribution παN(y) _α N(y) and its expected reward is given by Rbon[π1]=αNαN+1R_bon[ _1]= α Nα N+1. Therefore, we can analytically compute the true difference of expected reward as Rbon[πα]−Rbon[π1]=NαNα+1−N+1R_bon[ _α]-R_bon[ _1]= NαNα+1- NN+1. We expect that the estimated linearized loss ⟨R[⋅,πα−π1⟩ R[ π], _α- _1 approximates this true difference well when M is sufficiently large and α≃1α 1. Specifically, we construct π π by drawing M=2,…,32M=2,…,32 samples from π1 _1 and estimate the linearized loss for α=1+10−4α=1+10^-4. Then, the mean squared error between the estimated linearized loss and the true difference of expected reward is computed over 100,000 trials. To reduce the variance, linearized reward is centered by subtracting the mean as in GRPO algorithm. Figure 3 shows the comparison results for different M and N values. We can observe that large N requires larger M to achieve low expected error but it still achieves reasonable accuracy even with M<NM<N compared to N=1N=1 case, which corresponds to the standard GRPO. In particular, the bias term rapidly decreases as M increases and nearly vanishes for M≥NM≥ N. Figure 3: Expected error of linearized loss estimation with different M and N values. From left to right, N=1,4,8,16N=1,4,8,16. The shaded area indicates the bias-variance decomposition of the expected error. J.2 Length Reward Experiments In addition to the results with N=2,4N=2,4 in Figure 6, we provide the results with N=8,16N=8,16 in Figure 4. We also show the distribution of the reference model, which is the initial model before training. We observe that the distribution of the reference model is approximately uniform and the baseline model produces only medicore length responses concentrated around 100 tokens. On the other hand, IAMA models generate diverse responses and as N increases, the distribution is more polarized to produce both short and long responses. This aligns with the results in Section 3.1. Figure 4: Distribution of completion lengths of reference model πref _ref, baseline model aligned by standard GRPO, and our IAMA models trained with BoN (N=2,4,8,16N=2,4,8,16) objectives. J.3 H-RLHF Experiments We provide additional experimental results for H-RLHF experiments. Experiments with N=8N=8 Figure 5 (right) shows the RLHF results with N=8N=8. Similar to the N=8N=8 case in Figure 6, our method pushes the Pareto frontier significantly compared to the baseline methods. Figure 5: RLHF results with N=4N=4 (left) and N=8N=8 (right). Experiments on Mistral 7B We conduct additional experiments using Mistral 7B Instruct (jiang2023mistral7b), a state-of-the-art open-source LLM. We follow the same experimental setup as with Alpaca 7B except for β=0.01β=0.01 and KLtarget=0.1KL_target=0.1. Figure 6 shows the Pareto frontiers of IAMA and baseline with and without BoN sampling (N=4N=4). We see a similar tendency as Alpaca 7B results. Figure 6: H-RLHF results on Mistral 7B with N=4N=4.