Tail-Aware Information-Theoretic Generalization for RLHF and SGLD

Abstract: Classical information-theoretic generalization bounds typically control the generalization gap through KL-based mutual information and therefore rely on boundedness or sub-Gaussian tails via the moment generating function (MGF). In many modern pipelines, such as robust learning, RLHF, and stochastic optimization, losses and rewards can be heavy-tailed, and MGFs may not exist, rendering KL-based tools ineffective. We develop a tail-dependent information-theoretic framework for sub-Weibull data, where the tail parameter $θ$ controls the tail heaviness: $θ=2$ corresponds to sub-Gaussian, $θ=1$ to sub-exponential, and $0<θ<1$ to genuinely heavy tails. Our key technical ingredient is a decorrelation lemma that bounds change-of-measure expectations using a shifted-log $f_θ$-divergence, which admits explicit comparisons to Rényi divergence without MGF arguments. On the empirical-process side, we establish sharp maximal inequalities and a Dudley-type chaining bound for sub-Weibull processes with tail index $θ$, with complexity scaling as $\log^{1/θ}$ and entropy$^{1/θ}$. These tools yield expected and high-probability PAC-Bayes generalization bounds, as well as an information-theoretic chaining inequality based on multiscale Rényi mutual information. We illustrate the consequences in Rényi-regularized RLHF under heavy-tailed rewards and in stochastic gradient Langevin dynamics with heavy-tailed gradient noise.

• The paper introduces a tail-adaptive framework using sub-Weibull models to derive robust generalization bounds under heavy-tailed loss conditions.
• It adapts information-theoretic and decorrelation techniques to replace KL-based methods with Rényi divergences for stable RLHF and SGLD performance.
• Empirical validation shows that Rényi regularization prevents catastrophic reward collapse, ensuring improved reward stability in heavy-tailed settings.

Tail-Aware Information-Theoretic Generalization for RLHF and SGLD

Introduction and Motivation

The paper addresses the lack of robust generalization guarantees for algorithmic pipelines exposed to heavy-tailed loss distributions, with a particular focus on reinforcement learning from human feedback (RLHF) and stochastic gradient Langevin dynamics (SGLD). Classical generalization theory, both via uniform-convergence (VC theory, Rademacher, Gaussian complexity) and information-theoretic approaches (KL-based mutual information), fundamentally depends on light-tailed (sub-Gaussian or sub-exponential) assumptions—specifically the existence of moment generating functions (MGFs). However, empirical evidence in modern large-scale optimization and RLHF demonstrates the prevalence of genuinely heavy-tailed (sub-Weibull, $0 < \theta < 1$) distributions, where existing techniques and KL-based metrics become ineffective because the relevant MGFs diverge.

The paper develops a tail-adaptive information-theoretic framework for generalization analysis under sub-Weibull assumptions. The critical ingredient is a novel decorrelation lemma and divergence comparison, which yields explicit and sharp generalization and reward bounds even in the absence of exponential tail decay.

Sub-Weibull Processes and Tail-Dependent Complexity

The authors leverage the sub-Weibull family as a flexible model for tail behavior, parameterized by $\theta$: $\theta=2$ (sub-Gaussian), $\theta=1$ (sub-exponential), and $0<\theta<1$ (genuinely heavy-tailed). The central technical result is a maximal inequality (Lemma~\ref{pp-MaximalSW}) demonstrating that, for a family of $n$ i.i.d. sub-Weibull($\theta$) variables,

$$\mathbb{E}\left[\max_{i\in[n]}|X_i|\right] = O((\log n)^{1/\theta}),$$

with explicit, non-asymptotic leading constants and without loss from norm-to-moment conversions.

For processes indexed by metric spaces, the paper extends Dudley's entropy integral: $$\mathbb{E}\Bigl[\sup_{t\in T} X_t\Bigr] \leq 4C K_\theta \int_0^\infty [\log N(T, d, \epsilon)]^{1/\theta} d\epsilon,$$ quantifying the inflation of metric entropy complexity as $\theta \downarrow 0$.

Information-Theoretic Generalization for Heavy Tails

KL-based mutual information bounds (cf. [Russo20]) fail when the loss/reward is heavy-tailed. The authors construct a family of shifted-log $\theta$0-divergences, $\theta$1, matching the sub-Weibull tails. They establish two central comparison lemmas, showing that

$\theta$2

so R\'enyi divergences (with appropriately chosen $\theta$3) can substitute for KL in generalization analysis.

A new decorrelation lemma provides a change-of-measure inequality for expectations under dependent couplings (e.g., algorithm outputs depending on the sample) in terms of R\'enyi or $\theta$4 information and an Orlicz-norm tail term. This yields PAC-Bayes, expected generalization, and multi-scale chaining bounds that remain meaningful for sub-Weibull losses.

RLHF: Trust Regions, Goodhart and R\'enyi Regularization

A key application is RLHF, wherein KL-constrained trust regions have become a default for regularizing policy updates away from a reference (pretrained) model. The authors demonstrate, both theoretically and empirically, that even strict KL control can lead to catastrophic "Goodhart" effects when the reward model is heavy-tailed [(2604.10727), kwa2024catastrophic]: as the KL constraint is relaxed, the proxy reward can be aggressively optimized while the true (gold/evaluator) reward collapses.

This phenomenon is visualized in:

Figure 1: As the power exponent $\theta$5 increases, the proxy reward keeps increasing with KL divergence, while the gold reward peaks and then collapses, signaling catastrophic Goodhart in the heavy-tailed regime.

To mitigate this, the authors analyze and advocate for R\'enyi-regularized RLHF, where R\'enyi divergence constraints provide a power-type penalty on large density-ratio departures, thereby directly controlling extreme policy deviations and yielding sharply bounded increases in expected reward on sub-Weibull tails: $\theta$6 where $\theta$7 is the Orlicz norm of the reward under the reference.

Figure 2: R\'enyi-regularized RLHF with heavy-tailed proxy rewards yields gold reward curves that increase and then plateau, in contrast to KL which exhibits collapse.

Furthermore, in the context of best-of-$\theta$8 policies, the data-processing inequality for R\'enyi divergence ensures no catastrophic reward inflation, and the reward gain scales benignly as $\theta$9.

Experimental Validation and Diagnostics

The empirical section implements an RLHF pipeline with power-transformed proxy rewards to induce controlled heavy tails. Distributional diagnostics confirm the emergence of heavy tails under higher $\theta=2$0 (order-preserving transformation):

Figure 3: Proxy reward distribution before and after an order-preserving heavy-tailed transformation for $\theta=2$1, showing amplified extremes in the latter.

The experiments report striking divergence between proxy and gold rewards under KL penalties as tails become heavier: gold reward rises and then collapses with increasing divergence. Under R\'enyi regularization $\theta=2$2, the relationship stabilizes, with the gold reward being monotone or plateauing (see Figures~\ref{divergence_vs_reward11} and~\ref{divergence_vs_reward1210}). This holds across transformations and for additive Weibull-noise constructions.

Figure 4: Proxy--proxy-gold reward coupling under R\'enyi regularization remains monotone for all tested heavy-tail scenarios, while KL collapses for large $\theta=2$3.

Implications and Future Directions

The work rigorously establishes that KL-based generalization and RLHF regularization are inadequate for heavy-tailed regimes—empirical and theoretical analyses reveal the possibility of arbitrarily poor reward transfer under even minor KL departures. In contrast, R\'enyi and, more generally, $\theta=2$4-divergence metrics, provide tail-adaptive, robust bounds connecting algorithmic information and statistical generalization.

This has practical consequences for the construction of RLHF trust regions and the design of regularization in large-scale LLM alignment: tail-aware divergence penalties are critical to avoid overoptimization on rare, spurious reward artifacts and to stabilize transfer to more accurate evaluators.

Theoretically, the framework serves as a blueprint for extending empirical process and information-theoretic learning theory to non-sub-Gaussian, real-world settings—an open question is how to optimize bounds and algorithms when tail indices or process increments are only locally controlled or nonparametric. Methodological extensions to complex dependency structures, non-i.i.d. data (RL trajectories), more general optimization dynamics, and hybrid divergence metrics are natural next steps.

Conclusion

The paper provides a comprehensive, technically sophisticated framework for heavy-tailed generalization in statistical learning, with pivotal implications for RLHF and stochastic optimization. Tail-adaptive information measures, notably R\'enyi divergence, are established as robust and quantifiable tools for both theoretical analysis and algorithm design under heavy-tailed reward and loss distributions. This challenges KL-centric orthodoxy in regularized policy optimization, and points toward more resilient, tail-sensitive algorithmic pipelines for both AI alignment and large-scale learning (2604.10727).