Source-linked AI summary

The Asymptotics of Language Model Alignment with Memory

Haricharan Balasundaram, V. Arvind Rameshwar

arXiv:2610.01828v1cs.CLcs.ITmath.IT

TL;DR

Earlier asymptotic comparisons between alignment methods assumed i.i.d. outputs, leaving unclear whether their closeness holds with memory. This paper analyzes Markov outputs with local rewards and finds asymptotic closeness, while identifying a zero-KL condition at m=1.

Abstract

from arXiv · show

Language model (LM) alignment broadly aims to perturb a given LM $Q$ into an aligned LM $q$ such that i) the outputs produced by $q$ and $Q$ are 'close' in probability, ii) $q$ has a higher expected reward than $Q$. Two common techniques for LM alignment are: KL-constrained RL, which requires knowledge of the LM distribution and is computationally expensive, and the best-of-$n$ algorithm, which requires only sampling from the LM. The work of Yang et al. established asymptotic closeness between the distributions produced by the two alignment methods for an $m$--length i.i.d. token sequence output by the LM, in the limit as $m$ increases to infinity. However, the i.i.d. assumption is not representative of practical LMs, whose output sequences often have memory. In this paper, we extend the asymptotic closeness result to the case when the $m$--length token sequence outputted by the LM is Markovian. Further, for finite-length output sequences -- particularly, when $m=1$ -- we provide a complete characterization of LM distributions and reward functions for which the KL-divergence between the distributions produced by the two alignment methods is zero -- a question first posed in Yang et al.

I. INTRODUCTION … C. Our contributions towards (Q1)

This paper extends the asymptotic connection between KL-constrained RL and best-of-n alignment to Markov outputs with locally decomposable rewards, and studies alignment KL divergence for finite outputs, particularly m = 1.

I. INTRODUCTION

  • Alignment seeks higher expected reward; KL-constrained RL optimizes under a divergence bound, while best-of-n selects the highest-reward sequence among n LM samples.The former requires knowledge of the reference distribution, while best-of-n operates by sampling from it.Source: §I · p. 1
  • Earlier work established asymptotic closeness for i.i.d. tokens with additive token rewards, motivating the question of whether the connection persists with Markov dependence.That result takes n = Θ(exp(m)); practical LM outputs can depend on previous contexts.Source: §I · p. 1
  • For m = 1, the paper identifies reference LM distributions and rewards yielding small alignment KL divergence, and studies vocabulary-size dependence for special classes.The formulation asks which distributions and rewards make the divergence zero or close to zero.Source: §I · p. 2, B. Problem Formulation · p. 3

A. Notation and Assumptions

  • The analysis assumes an irreducible, aperiodic Markov output chain initialized at stationarity and a reward that sums rewards over adjacent terms.These assumptions formalize Markovian output memory and local reward decomposition.Source: A. Notation and Assumptions · p. 2

C. Our contributions towards (Q1)

  • The KL-constrained RL solution is an exponentially tilted distribution that can also be represented as a non-stationary Markov process.The process is characterized by an initial distribution and time-varying stochastic transition matrices.Source: C. Our contributions towards (Q1) · p. 3
  • As sequence length grows, the characterized transition matrices converge to one another when their positions recede from the sequence end, indicating an approach to stationarity.The theorem states that transition matrices at positions whose distances from the end both grow have the same limit.Source: C. Our contributions towards (Q1) · p. 3
  • Under the reference LM, the average reward satisfies a large deviation principle with a good rate function.The average reward is defined over the length-m output sequence.Source: C. Our contributions towards (Q1) · p. 3
  • For Markov LM outputs with locally decomposable rewards, the KL-constrained RL and best-of-n distributions are asymptotically close as sequence length grows.The paper takes n = Θ(exp(m)) in the best-of-n setting.Source: §I · p. 2, C. Our contributions towards (Q1) · p. 3

D. Our contributions towards (Q2)

For output length m = 1, the paper characterizes when KL-constrained RL and best-of-n distributions coincide exactly, and shows that in a discrete uniform setting their KL separation decays as vocabulary size grows.

  • The preceding asymptotic result establishes closeness for Markov outputs under mild assumptions when n = exp(mδ), but applies only as m tends to infinity.Source: D. Our contributions towards (Q2) · p. 3
  • For n ≥ 2 and almost-surely distinct rewards, the distributions coincide exactly if and only if the reward CDF is suitably shifted exponential.This characterization concerns the continuous m = 1 case.Source: D. Our contributions towards (Q2) · p. 4
  • For n ≥ 4, a uniform distribution over a finite vocabulary with the continuous-case reward function has nonzero KL separation at finite size, but the separation decays to zero as vocabulary size grows.The reward is r(y) = b + ((n−1)/β)·log(y).Source: D. Our contributions towards (Q2) · p. 4
  • For n = 10, Figure 1 compares the discrete-case KL divergence with the theoretical asymptotic rate 1/(8d^2), and the text reports that the rate is verified.Source: Figure 1 · p. 4, D. Our contributions towards (Q2) · p. 4

III. KL-CONSTRAINED RL SOLUTION

The section derives the KL-RL solution as a tilted distribution and formalizes Theorem II.2 by defining ν and p(t), which are valid probability objects.

III. KL-CONSTRAINED RL SOLUTION

  • A Lagrangian yields the KL-RL solution as a tilted distribution, analogous to [1, Lemma 1] and using.This establishes Theorem II.1.Source: §III · p. 4
  • The formal Theorem II.2 defines h(t) recursively to specify ν and p(t), with 1_d denoting the all-ones vector in R^d.ν is a valid probability distribution, and each p(t) is a valid stochastic matrix.Source: §III · p. 5

A. Long-term behavior

The section shows that the long-term transition matrices converge to the solution of KL-RL with a flow constraint. That solution is stationary Markov and is characterized using the Perron eigenvalue and eigenvector of Gβ.

A. Long-term behavior

  • The limiting matrix p is the solution to KL-RL with an added flow constraint, and the resulting q is stationary Markov.Source: A. Long-term behavior · p. 5
  • With P aperiodic and irreducible, the flow-constrained solution is characterized using the Perron eigenvalue λβ and right eigenvector h of Gβ, with β chosen so dKL(p, P) = δ.The stationary distribution μ satisfies μT p = μT.Source: A. Long-term behavior · p. 5
  • For fixed δ > 0, the solution value sδ is related to the rate of change of log(λβ) with respect to β.β is chosen as in Theorem III.2.Source: A. Long-term behavior · p. 5
  • As m →∞, p(t) converges to p whenever lim_{m→∞}(m−t) →∞.The limit relates the unconstrained KL-RL matrices p(t) to the flow-constrained solution p.Source: A. Long-term behavior · p. 5

IV. LDP FOR AVERAGE REWARD · V. CONNECTING THE KL-CONSTRAINED RL AND BEST-OF-n SOLUTIONS

The sections establish a large-deviation principle for empirical transitions in the Markovian setting and use it to connect average-reward optimization with the asymptotic behavior of best-of-n rewards. They show that the KL-RL and average-reward formulations share an optimum, while the maximum reward among exp(mδ) samples converges to that optimum.

IV. LDP FOR AVERAGE REWARD

  • The transition-based LDP is non-trivial only under the stated constraint on the empirical transition probabilities θ [26].The supplied passage truncates the constraint after θj = P, so its full form is not available here.Source: §IV · p. 6
  • The empirical transition probabilities θ satisfy a large-deviation principle with a good rate function, providing the Markovian basis for the average-reward result.Applying the contraction principle to the continuous average-reward function yields the desired result; the authors note this argument is simpler than the i.i.d. argument in.Source: §IV · p. 6
  • Theorem II.4 minimizes KL divergence for a given expected reward, whereas KL-RL with Flow maximizes expected reward subject to a KL constraint.The two formulations have different objectives, but Lemma IV.2 establishes that their optimum values coincide: sδ = tδ.Source: §IV · p. 6

V. CONNECTING THE KL-CONSTRAINED RL AND BEST-OF-n SOLUTIONS

  • For n = exp(mδ), the maximum reward Tm,exp(mδ) among n i.i.d. reward samples from Q converges in probability to sδ.The proof uses the LDP rate function from Theorem II.4.Source: §V · p. 6
  • Because Tm,exp(mδ) is bounded, its expected value also converges to sδ.This is the stated corollary following the convergence-in-probability result.Source: §V · p. 6
  • The proof of the KL-RL and best-of-n connection reduces to showing a remaining expression is zero, established by an algebraic lemma.The expression uses Λβ = log(λβ) and sδ = Λ′, as recalled from Theorem III.3.Source: §V · p. 6

VI. CLOSENESS OF BEST-OF-n AND THE KL-RL DISTRIBUTION FOR m = 1 … VIII. AI USE STATEMENT

For Markov LM outputs with locally decomposed rewards, KL-constrained RL and best-of-n yield asymptotically close aligned distributions as output length grows, extending the i.i.d. result in. The paper also identifies finite-regime closeness classes and develops an m = 1 comparison.

VI. CLOSENESS OF BEST-OF-n AND THE KL-RL DISTRIBUTION FOR m = 1

  • The best-of-n and tilted distributions have the same conditional law given reward, and their reward-law KL divergence equals their full-distribution KL divergence.The stated conditional-law identity is qβ(Y′ ∈ B | L′ = l) = Q(Y ∈ B | L = l); Corollary VI.2 gives dKL(pL′, pL) = dKL(qβ, Q).Source: §VI · p. 7
  • For continuous reward CDFs, the best-of-n distribution is absolutely continuous with respect to Q, with Radon–Nikodym derivative nPL(r(y))^(n−1).The section uses this expression alongside the tilted distribution to obtain Theorem II.6, which characterizes when the two distributions have zero KL divergence.Source: §VI · p. 7

A. Discrete Case

  • In the discrete case, qβ is independent of β and satisfies the KL constraint; finite-d expressions are expanded to derive Theorem II.7.The derivation expands the denominator to an in−1-order term and cancels corresponding numerator and denominator terms to study dependence on d and n.Source: A. Discrete Case · p. 7

VII. CONCLUSIONS

  • For Markov outputs with locally decomposed rewards, KL-constrained RL and best-of-n produce asymptotically close aligned distributions as output length grows, extending.The paper also identifies classes of LM distributions and reward functions for which the distributions remain close in a finite, non-asymptotic regime.Source: §VII · p. 7
  • The paper identifies future work on asymptotic closeness under mixing assumptions and non-asymptotic closeness for arbitrary discrete distributions.Source: §VII · p. 7

VIII. AI USE STATEMENT

  • OpenAI GPT 5.5 generated Python Matplotlib plot code and helped clean up proofs; GPT 5.6 proofread and edited grammar and style, with the authors retaining full responsibility.Source: §VIII · p. 7

APPENDIX A PROOF OF THEOREM III.1 · APPENDIX B PROOF OF THEOREM III.2

The appendices verify the Markov-chain construction used in Theorem III.1 and derive the Perron-eigenvector form, constraints, and stationary-distribution properties used in Theorem III.2. They also connect the resulting stationary distribution to the expected reward calculation.

APPENDIX A PROOF OF THEOREM III.1

  • Theorem III.1’s proof checks that the constructed distribution and transition matrices satisfy probability and stochasticity requirements.It separately examines the denominator and numerator, verifies non-negativity for ν, and checks that each p(t) is stochastic.Source: Appendix A · p. 9

APPENDIX B PROOF OF THEOREM III.2

  • Theorem III.2’s proof formulates the optimization problem with a Lagrangian and obtains the candidate transition probabilities from its stationarity conditions.The resulting expression uses nonnegative constants C_a and D_b.Source: Appendix B · p. 9
  • Primitivity of Gβ gives a unique Perron eigenvalue and right eigenvector, which determine the constants in the candidate transition probabilities.The proof sets D_b=h_b and C_a=1/(λβh_a).Source: Appendix B · p. 9
  • The proof verifies row-stochasticity and normalization by choosing μ as a normalized left eigenvector of p and β so that d_KL(p, P)=δ.It then establishes uniqueness of μ as the normalized positive left eigenvector of p.Source: Appendix B · pp. 9–10
  • For the expected-reward calculation, the proof identifies the stationary distribution as μ_j=g_jh_j, using left and right Perron eigenvectors normalized so g^T h=1.The vector g is the left eigenvector of Gβ for λβ.Source: Appendix B · p. 10
  • The proof differentiates Gβh=λβh with respect to β and uses the definitions of p and μ to evaluate sδ, the expected reward under the constructed chain.This links the Perron-eigenvalue calculation to the reward quantity defined in the proof.Source: Appendix B · p. 10

APPENDIX D PROOF OF LEMMA III.4 … APPENDIX G PROOF OF LEMMA IV.2

The appendices establish the Perron–Frobenius setup for Lemma III.4, derive equations in Theorem III.5’s proof, and use contraction and matching optimization bounds to prove the stated results.

APPENDIX D PROOF OF LEMMA III.4

  • Lemma III.4’s proof invokes Perron–Frobenius for Gβ, with unique Perron eigenvalue λβ and unique left and right Perron eigenvectors g and h.The eigenvectors are normalized so that gT h = 1.Source: Appendix D · p. 10

APPENDIX E PROOF OF THEOREM III.5

  • The cited proof invokes Lemma III.4 and uses the two equations above to obtain a result for a, b ∈ Y.Source: Appendix E · p. 11

APPENDIX F PROOF OF THEOREM II.4

  • Theorem II.4’s proof applies the contraction principle to the continuous function S_m = ϕ(θ), obtaining an LDP with a good rate function.The proof identifies the rate using Theorem IV.1.Source: Appendix F · p. 11
  • Theorem II.4’s proof obtains the theorem’s rate via Theorem IV.1 and establishes that ˜J(s) is good by the contraction principle.Source: Appendix F · p. 11

APPENDIX G PROOF OF LEMMA IV.2

  • For the forward bound, an optimizer (µ⋆, p⋆) of KL-RL with Flow is feasible for defining ˜J(sδ), yielding sδ ≤ tδ.Source: Appendix G · pp. 11–12
  • For the reverse bound, feasible pairs for ˜J(s) make s feasible in KL-RL with Flow, yielding tδ ≤ sδ; together the bounds give tδ = sδ.Source: Appendix G · p. 12

APPENDIX H PROOF OF LEMMA V.1 · APPENDIX I PROOF OF LEMMA V.3

The appendix proves concentration of the best-of-exp(mδ) reward statistic around sδ using large-deviation bounds and convexity of the rate function. It also completes Lemma V.3 by using stationarity and the KL constraint.

APPENDIX H PROOF OF LEMMA V.1

  • The probability that Tm,exp(mδ) exceeds sδ + ϵ tends to zero, since ˜J(s) > δ above sδ and the proof applies the large-deviation principle.The proof invokes a union bound and the LDP for Sm.Source: Appendix H · p. 12
  • ˜J(s) is convex because J(θ) is convex, Sm is linear in θ, and minimizing the resulting convex L(s, θ) yields ˜J(s).The proof attributes J(θ)'s convexity to convexity of the Markovian KL divergence in its first argument.Source: Appendix H · p. 12
  • For the lower-tail bound, the proof establishes inf_{s>sδ−ϵ} ˜J(s) ≤ δ − γ for some γ > 0 by considering cases around sδ and the mean s0.It uses ˜J(sδ) ≤ δ and, when needed, ˜J(s0) = 0 at s0 = E[Sm].Source: Appendix H · pp. 12–13
  • P(Tm,exp(mδ) ≤ sδ − ϵ) is bounded by exp(−exp(mγ + o(m))) for some γ > 0, so the probability tends to zero.The bound follows from (78) and the fact that inf_{s>sδ−ϵ} ˜J(s) ≤ δ − γ for some γ > 0.Source: Appendix H · pp. 12–13

APPENDIX I PROOF OF LEMMA V.3

  • Since µ is stationary for p, the proof cancels the second and third terms in (87) before combining the remaining terms.The supplied proof passages do not state the resulting expression.Source: Appendix I · p. 13
  • The proof concludes Lemma V.3 using sδ = Λ′ and choosing β so that dKL(p, P) = δ, as in (8).It refers to Theorem III.2 and the stationary distribution µ for p.Source: Appendix I · p. 13

APPENDIX J PROOF OF LEMMA VI.1 · APPENDIX K PROOF OF LEMMA VI.3 · APPENDIX L PROOF OF THEOREM II.6

The appendices prove the stated lemmas using conditional-expectation and symmetry arguments, then prove Theorem II.6’s support characterization and check its KL constraint.

APPENDIX J PROOF OF LEMMA VI.1

  • Under qβ, the joint event for Y′ and L′ equals the Q-expectation of its indicator weighted by g(L).Conditioning on L factors this as the weighted probability Q(Y ∈ B | L) on L ∈ C.Source: Appendix J · p. 13
  • Conditioning on L shows the tilted joint probability factors into g(L), the event L ∈ C, and Q(Y ∈ B | L).This follows by applying the tower law to the weighted indicator.Source: Appendix J · p. 13

APPENDIX K PROOF OF LEMMA VI.3

  • By symmetry, the probability the selected sample lies in B is n times the probability Y1 lies in B and every other sample has lower reward.The proof then conditions on Y1 and uses independence and (32).Source: Appendix K · p. 14

APPENDIX L PROOF OF THEOREM II.6

  • The forward direction of Theorem II.6 concludes that L has support (−∞, b] when its CDF is continuous and equals 1 for b or greater.The proof obtains the CDF condition from the almost-sure agreement of Radon–Nikodym derivatives.Source: Appendix L · p. 14
  • Conversely, assuming L has the CDF specified in Theorem II.6, the proof verifies the corresponding density relation and concludes the converse.The supplied passage does not show the density formula.Source: Appendix L · pp. 14–15
  • The KL check uses that L′ is reverse exponential and evaluates the divergence between exponential distributions using [29, Table 1].Choosing β so dKL(qβ, Q) = δ yields the desired result.Source: Appendix L · p. 15

APPENDIX M PROOF OF LEMMA VI.4

Appendix M proves Lemma VI.4 by establishing a majorization relation between qβ and qbest, then bounding the resulting KL divergence.

APPENDIX M PROOF OF LEMMA VI.4

  • Prefix-sum comparisons, with equality only at k=d, establish that qbest_{n,1} majorizes qβ.The proof uses that a sequence indexed by i in [d] is increasing.Source: Appendix M · p. 15
  • Karamata’s inequality and convexity of KL divergence in its first argument yield the stated divergence comparison.The proof cites [30, Theorem 1].Source: Appendix M · p. 15
  • The KL upper bound uses s(x)=nx^(n−1) and t(x)=1 on (0,1], applies data processing through f, and evaluates the integral by parts.The data processing inequality is cited as [31, Corollary 2.18]; combining the resulting equations gives the desired result.Source: Appendix M · p. 15

APPENDIX N PROOF OF THEOREM II.7

Appendix N proves Theorem II.7 through a calculation for n ≥4, using a log expansion, Faulhaber’s formula, and algebraic manipulation.

APPENDIX N PROOF OF THEOREM II.7

  • For n ≥4, the proof evaluates the target quantity using qβ(yi)(1 + δi) for all i and the expansion log(1 + δi) = δi − δi^2/2 + O(δi^3).The derivation also invokes (129).Source: Appendix N · p. 16
  • Faulhaber’s formula is used to simplify the calculation, and the proof applies it again in a later step.Source: Appendix N · p. 16
  • The proof completes the calculation with algebraic manipulations and a final use of Faulhaber’s formula to obtain the desired result.Source: Appendix N · p. 17
Loading 2610.01828v1…