SAIL-RevKL: Stable and Provable Self-Improving Online LLM Alignment (UAI 2026)

Our paper, “On the Convergence of Self-Improving Online LLM Alignment,” appeared at the 42nd Conference on Uncertainty in Artificial Intelligence (UAI 2026). It addresses a basic theoretical gap behind self-improving alignment.

Research Question

Can we establish a global convergence guarantee for online RLHF formulated as a bilevel optimization problem?

Proceedings · arXiv · Poster · MuJoCo Code · ** Code**

TL;DR

Our answer is conditional but positive. We first show that the original Self-Improving Alignment (SAIL) objective is guaranteed to have favorable curvature only near its initialization. We then introduce SAIL-RevKL, which adds

$$ \gamma\,\mathbb{E}_x!\left[ D_{\mathrm{KL}}!\left( \pi_{\mathrm{ref}}(\cdot\mid x)\,\Vert\,\pi_\theta(\cdot\mid x) \right) \right] $$

as a penalty. For a fixed reference and a log-linear policy, the Hessian of this reverse KL is the Fisher information matrix. With enough regularization, that positive curvature dominates the adverse curvature of SAIL throughout the full bounded parameter region. This yields global strong concavity and a Polyak-Lojasiewicz (PL) condition for the regularized surrogate, together with near-linear dependence on the inverse target accuracy.

Why Bilevel-Formulated Online RLHF Is Hard

Offline preference optimization trains on a fixed dataset. Online RLHF instead couples two learning problems: a reward model is fitted from preferences over responses generated by the current policy, and the policy is updated against that learned reward. This interaction is naturally formulated as a bilevel optimization problem because the solution of one level determines the objective and data distribution of the other.

SAIL uses the reward-policy equivalence to reduce this computationally expensive bilevel problem to an efficient single-level surrogate. If $\mathcal{D}_\theta$ denotes the policy-dependent response and preference distribution, that surrogate has the form

$$ J(\theta) = \mathbb{E}_{(x,y_w,y_l)\sim\mathcal{D}_\theta} !\left[ \log\sigma!\left( \beta\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\mathrm{SFT}}(y_w\mid x)} - \beta\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\mathrm{SFT}}(y_l\mid x)} \right) \right]. $$

This dependence matters. Standard performance-difference arguments assume that the reward or evaluation measure is independent of the policy appearing inside the advantage. That assumption no longer holds here, so convergence cannot be inherited directly from ordinary policy-gradient analysis. We instead analyze the geometry of $J$ itself.

The negative Hessian of $J$ decomposes into four terms. Some are stabilizing, while interaction terms can become unfavorable as $\theta$ moves away from the initial SFT parameters $\theta_0$. The proof must therefore control the objective’s curvature throughout the region the algorithm may visit, not only at its starting point.

Assumptions and Scope

The global theorem is explicit about its scope. It applies to the following setting:

Two construction choices are also essential. The reference policy $\pi_{\mathrm{ref}}$ is fixed during the analyzed update, and projected gradient ascent keeps every iterate inside $\Theta$, the radius-$B_\theta$ Euclidean ball centered at $\theta_0$.

Scope of the result. The word “global” below means global within this full bounded feasible set. The LoRA experiments test whether the method remains useful beyond these assumptions; the last-layer-only experiments are the closest empirical match to A1.

Core Theorem: Global Convergence of SAIL-RevKL

The original SAIL surrogate is guaranteed to have favorable geometry only near its initialization. We add a reverse-KL penalty and optimize

$$ J_\gamma(\theta) = J(\theta) - \gamma\,\mathbb{E}_x!\left[ D_{\mathrm{KL}}!\left( \pi_{\mathrm{ref}}(\cdot\mid x)\,\Vert\,\pi_\theta(\cdot\mid x) \right) \right]. $$

For a fixed reference and a log-linear policy, the reverse-KL Hessian is the Fisher information matrix, so

$$ -\nabla^2J_\gamma(\theta) = -\nabla^2J(\theta)+\gamma F(\theta). $$

Theorem (Global convergence on the feasible set)

Let $\Theta$ be the radius-$B_\theta$ Euclidean ball centered at $\theta_0$. Suppose A1-A4 hold, $\pi_{\mathrm{ref}}$ remains fixed during optimization, and projected updates keep every iterate in $\Theta$. Choose $\gamma$ sufficiently large to satisfy the explicit curvature threshold in Theorem 2 of the paper, so that the Fisher term dominates the adverse curvature of the original SAIL objective throughout $\Theta$.

Then $J_\gamma$ is $\mu$-strongly concave on all of $\Theta$ and satisfies the PL inequality

$$ \Vert\nabla J_\gamma(\theta)\Vert_2^2 \ge 2\mu\left(J_\gamma(\theta_\gamma^{\ast})-J_\gamma(\theta)\right), \qquad \theta\in\Theta, $$

where $\theta_\gamma^{\ast}=\arg\max_{\theta\in\Theta}J_\gamma(\theta)$ is the unique maximizer. Moreover, for projected mini-batch stochastic gradient ascent with a valid constant stepsize, there exist constants $q\in(0,1)$ and $C>0$ such that

$$ \mathbb{E}!\left[ J_\gamma(\theta_\gamma^{\ast})-J_\gamma(\theta_T) \right] \le q^T\left(J_\gamma(\theta_\gamma^{\ast})-J_\gamma(\theta_0)\right) + \frac{C\sigma_g^2}{B_s}. $$

Thus the optimization error decreases geometrically until it reaches the mini-batch noise floor. Taking $B_s=\widetilde{\mathcal{O}}(\varepsilon^{-1})$ and $T=\widetilde{\mathcal{O}}(\log\varepsilon^{-1})$ gives an $\varepsilon$ expected function-value gap with total sample complexity

$$ \widetilde{\mathcal{O}}!\left( \varepsilon^{-1}\log\varepsilon^{-1} \right). $$

The theorem is proved in a deliberately controlled log-linear setting. We next test whether the same regularized objective is useful when training real language models.

LLM Alignment Results

The main experiments use LoRA fine-tuning and GPT-4 as an LLM judge. PKU-SafeRLHF evaluates the balance between helpfulness and harmlessness, while UltraFeedback focuses on general instruction following. Pairwise win rate is the primary comparison below; tie rate and mean GPT score difference are reported separately because they capture different aspects of the judge outcomes.

PKU-SafeRLHF, Qwen 0.5B

UltraFeedback

The headline comparison is the pairwise win rate against vanilla SAIL: 43% vs. 29% on Qwen with PKU-SafeRLHF, 30% vs. 15% on Phi-3 with UltraFeedback, and 34% vs. 23% on LLaMA-3 with UltraFeedback.

A theory-aligned test

LoRA changes nonlinear features and is not a direct instantiation of the log-linear policy in A1. We therefore also freeze each backbone and train only its final linear layer on UltraFeedback U10. This isolates the regime that most closely matches the theory.

SAIL-RevKL improves the best win rate over vanilla SAIL on all three backbones. The preferred coefficient is scale-dependent: $10^{-2}$ for Qwen 0.5B and $10^{-1}$ for Phi-3 3.8B and LLaMA-3 8B. This is consistent with the theory’s curvature-bias trade-off and shows why $\gamma$ should be tuned rather than treated as a universal constant.

Citation

@inproceedings{wu2026convergence,
  title     = {On the Convergence of Self-Improving Online {LLM} Alignment},
  author    = {Wu, Xudong and Liu, Pangpang and Aggarwal, Vaneet and Chen, Jiayu},
  booktitle = {Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence},
  pages     = {7433--7467},
  year      = {2026},
  volume    = {337},
  series    = {Proceedings of Machine Learning Research}
}

Acknowledgements

This work was supported in part by the Seed Fund for PI Research, Basic Research, from The University of Hong Kong under Project Code 2502251784.