PaperlayerView exact revision on arXiv

Source-linked AI summary

Unifying ICL, SFT, KL-Regularized RL Through a Bayesian Lens

Junxin Fan

arXiv:2609.05111v1cs.AI

TL;DR

The paper asks how ICL, SFT, and KL-regularized RL methods can be understood within one framework. It constructs Bayes or Gibbs posteriors from reference models and utility signals, then views learning as forward-KL projection, showing shared objectives and update forms while preserving differences in learning signals. The framework also motivates supervised cold starts and connects training-time amortization with test-time reasoning procedures.

  • Problem

    ICL, SFT, and KL-regularized RL are often treated as fundamentally different despite needing a common account of their relationships and differing learning signals.

  • Method

    The paper builds Bayes or Gibbs posteriors from reference models and utility signals, then analyzes in-context and parameterized learning as forward-KL projections onto model families.

  • Results

    KL-regularized RL, reward-weighted SFT, reward-weighted ICL, and advantage-weighted SFT share forward-KL projection objectives and corresponding first-order weighted-score updates.

  • Takeaways & Limitations

    The framework treats RLHF/RLVR as posterior design followed by projection and clarifies why supervised cold starts commonly precede high-variance on-policy updates.

  • Takeaways & Limitations

    The equivalences concern objectives and first-order updates, while supervision source and granularity remain different across SFT, RL, and ICL.

Abstract

from arXiv · show

Large language models are now trained and evaluated under a diverse set of paradigms: supervised fine-tuning (SFT), few-shot in-context learning (ICL), KL-regularized RLHF/RLVR, on-policy distillation (OPD), and test-time reasoning with search and chain-of-thought. These methods are often discussed as fundamentally different, and recent empirical results--such as the mixed impact of few-shot prompting on RL-tuned reasoning models--can appear puzzling. This note develops a Bayesian perspective that puts these procedures on the same footing. At the core is a two-step template: (i) construct a (generalized) Bayes or Gibbs posterior q* over outputs or actions given a context, using a prior/reference model and a utility signal (log-likelihood, reward, or advantage); and (ii) approximate q* by a forward-KL projection onto a parametric family, either in-weights (SFT/RL) or in-context (ICL). Part I formalizes few-shot ICL and SFT as amortized and-weights projections onto the Bayes posterior predictive. Parts II-IV show that KL-regularized RLHF/RLVR, reward-weighted SFT, reward-weighted ICL (RW-ICL), and advantage-weighted SFT (AWSFT) are all instances of forward-KL projection onto posteriors induced by rewards or advantages. We disentangle where these equivalences hold (objectives and first-order updates) and where they do not (source and granularity of the learning signal). Part V sketches implications for modern reasoning pipelines: RLHF/RLVR recipes as "posterior design + projection", why cold-start or supervised warm-up is practically unavoidable for importance-weighted KL projections, and DeepSeek-R1 and o1-style reasoning models as combining test-time Bayesian search with training-time KL amortization.

Part I: Few-shot ICL ≈SFT (Bayesian View)

The paper frames few-shot ICL and SFT as different approximations to the same Bayesian target: the Bayes posterior predictive. ICL amortizes this prediction in a forward pass, while SFT performs a forward-KL projection onto a parametric model; analogous projections extend to reward-based RL objectives.

  • SFT projection: SFT learns the teacher or Bayes predictive distribution as its forward-KL projection onto the parametric model family.The projection is equivalent to maximum-likelihood training and can use direct samples from the target predictive distribution or an empirical approximation.
  • Bayesian target: The Bayes posterior predictive is the normative target for few-shot in-context prediction.
  • Amortized ICL: A sufficiently expressive pretrained Transformer can amortize this target, approximating it directly from the context during its forward pass.This formulation does not require test-time gradient descent, although prior work also interprets ICL through first-order or meta-learning updates.
  • ICL–SFT connection: Thus, ICL and SFT are different approximations to the same underlying Bayesian predictor rather than identical procedures.
  • Reward-based extension: KL-regularized RL produces a Gibbs posterior from a reference policy and reward, while reward-weighted SFT is its forward-KL projection onto the policy family.The paper states this equivalence at the target-posterior and projection level, not as a claim that every practical RLHF implementation reaches the closed-form optimum.

Part III: Reward-weighted Few-shot ICL (RW-ICL)

Reward-weighted ICL constructs a context-dependent Gibbs posterior from a reference model and reward, then trains an in-context predictor as its forward-KL projection. This matches KL-regularized RLHF/RLVR at the distribution level, up to projection error.

  • RW-ICL defines a Gibbs posterior using the reference model as prior and exp(R/β) as a reward-based pseudo-likelihood.The reward aggregates information from human preferences, verifiers, or reward models conditioned on the context and query.
  • The RW-ICL population objective trains a meta-predictor Mθ to approximate this generalized posterior across tasks and contexts.The expectation follows the same hierarchical meta-distribution used for ordinary ICL, together with the process generating rewards.
  • RW-ICL is a forward-KL projection of the generalized Gibbs posterior onto the in-context model family.Stochastic gradient ascent implements the projection using reference-model samples reweighted by exp(R/β).
  • At the first-order update level, RW-ICL, reward-weighted SFT, and KL-regularized RL share reward-weighted score-gradient structure.The correspondence concerns update forms, with the model and signal changing across ICL, SFT, and RL.
  • At the distribution level, reward-weighted ICL approximately matches the KL-regularized RLHF/RLVR optimum, up to approximation error in the KL projection.

Part IV: Process-level KL-RLVR and Advantage-weighted SFT

Process-level KL-regularized RL produces statewise Gibbs policies over actions, and advantage-weighted SFT projects those policies onto a parametric family. The same weighted-update structure extends to trajectory-level RW-ICL, while the paradigms retain different learning-signal granularities.

  • Process-level KL-RLVR: The KL-regularized policy-improvement optimum is a Gibbs posterior over actions, with Q(s, a)/β determining action preference relative to a reference policy.
  • Process-level KL-RLVR: Advantage baselines shift Q by state-dependent constants that are absorbed into normalization, leaving the normalized policy unchanged.
  • Advantage-weighted SFT: AWSFT is the forward-KL projection of the statewise KL-RLVR Gibbs optimum onto the parametric policy family.The empirical objective uses estimated advantages as exponentiated weights over state-action pairs.
  • Trajectory-level RW-ICL: Trajectory-level RW-ICL uses per-step advantage weights and produces updates matching AWSFT, subject to state-distribution and context-window choices.
  • Learning signals and credit assignment: Across SFT, RL, and ICL, the equivalence concerns posterior objectives and first-order updates, not the source or granularity of supervision.SFT has dense token-level labels, RL begins with sparse trajectory rewards and explicit credit assignment, and ICL adapts in activations after dense pretraining supervision.

Part V: Implications for Reasoning Models and Training Recipes

The framework interprets RLHF/RLVR pipelines as Gibbs-posterior design followed by forward-KL projection, clarifying why cold starts support stable importance weighting. It also explains reasoning models as combining test-time posterior search with training-time amortization, while few-shot prompting can shift the targeted posterior.

  • Posterior design and projection: RLHF/RLVR can be viewed as designing a Gibbs posterior from a reference model and reward, then projecting it into a deployable policy.Projection may use reward- or advantage-weighted SFT or an optimizer such as PPO.
  • Design knobs: Reward shaping, temperature β, and sampling regime control posterior concentration, reference-policy movement, and the stability–exploration trade-off.Small β concentrates the posterior, while offline and on-policy methods trade data reuse and stability against exploration.
  • Cold start and on-policy distillation: When the sampling distribution has little overlap with the high-reward Gibbs region, importance weights collapse and gradients become unstable.
  • Cold start and on-policy distillation: A small supervised or behavior-cloning warm-up moves the policy toward the high-reward region, making Gibbs-posterior projection statistically feasible.
  • Compute allocation and reasoning models: Reasoning models combine test-time search over solution paths with training-time RL, OPD, or SFT that amortizes high-reward behavior into model weights.DeepSeek-R1 and o1-style systems are described as combining these inference-time and training-time regimes.
  • Few-shot prompting and posterior mismatch: Few-shot prompting can degrade RL-tuned reasoning models when its richer contextual posterior differs from the posterior trained under zero-shot or lightly structured prompts.RW-ICL suggests defining rewards at the contextual level and training the model directly on that Gibbs posterior.

Final Conclusions

The paper unifies ICL, SFT, and KL-regularized RL as posterior construction followed by forward-KL projection, while restricting the equivalence to distributional objectives and first-order updates. It distinguishes these shared forms from differences in learning-signal source and granularity.

  • Few-shot ICL approximates the Bayes posterior predictive through an amortized forward-KL/MLE projection.
  • KL-regularized RL defines a Gibbs posterior, while reward-weighted SFT, AWSFT, and RW-ICL project onto corresponding reward- or advantage-induced posteriors.
  • The equivalence holds at posterior-target and first-order-update levels, not as an identity between practical algorithms.
  • Learning processes remain different: SFT uses dense token supervision, RL uses sparse rewards with credit assignment, and ICL adapts through forward activations without test-time gradients.
  • The common mechanism is a Bayesian or generalized Bayesian update followed by forward-KL projection onto a parameterized policy family.

Appendix A: Bayesian ICL Results Used in the Main Text

The appendix summarizes prior theoretical and empirical work supporting the paper’s Bayesian account of in-context learning. It uses these results as foundations rather than reproducing their full technical development.

  • The appendix reviews Wakayama and Suzuki and Reuter et al. as the main prior foundations for the paper’s Bayesian ICL arguments.
  • Full technical details and proofs are deferred to the original papers.

A.1 Risk decomposition and non-asymptotic bounds (Wakayama & Suzuki)

Wakayama and Suzuki decompose ICL risk into a model-dependent Bayes Gap and irreducible Posterior Variance, establishing the Bayes predictor as the normative target. Their bounds characterize approximation and pretraining-generalization effects, while task mixtures and distribution shift alter the Bayes Gap.

  • The ICL risk decomposition separates Bayes Gap, measuring distance from the Bayes predictor, from Posterior Variance, determined by task uncertainty and context length.
  • The Bayes-optimal predictor uniquely minimizes risk, making its posterior predictive the normative target for in-context prediction.
  • The non-asymptotic bound contains approximation error controlled by feature dimension and pretraining generalization error depending jointly on context length and prompt count.
  • In mixed-task settings, posterior task identification concentrates quickly as context grows, enabling rapid specialization to the true task family.
  • Under input-distribution shift, the out-of-distribution penalty increases the Bayes Gap, while Posterior Variance remains intrinsic to the target domain.
  • The paper uses these results to assume that sufficiently expressive, many-task-trained Transformers can approximate the Bayes predictor with a small Bayes Gap.

A.2 Transformers as amortized Bayesian inference (Reuter et al.)

Reuter et al. test whether in-context Transformers can approximate Bayesian posterior inference over latent variables and responses. Across synthetic and real-world settings, suitably trained Transformers often approach Bayesian baselines, supporting the paper’s amortized-inference assumption.

  • Reuter et al. study whether in-context Transformers can approximate full Bayesian posterior inference over latent variables.
  • They train learners that map datasets to predictive distributions or posterior samples over latent variables and responses.
  • They compare Transformer samples with HMC, SGLD, and variational approximations using C2ST, MMD, Wasserstein distance, and RMSE.
  • The experiments investigate robustness under training–test distribution shift and report that the in-context learner often remains close to Bayesian baselines.
  • Suitably trained Transformers often match or approach fully Bayesian baselines and sometimes outperform classical approximations on distributional metrics.
  • These findings support approximating the Bayes posterior predictive through a single forward pass.

A.3 Relation to the main text (and to meta-learning views)

This appendix situates the Bayesian ICL and SFT framework within prior work on Bayes-optimal prediction, amortized inference, and meta-learning interpretations of Transformer computation.

  • Part I treats the Bayes posterior predictive as the canonical target for in-context prediction and decomposes ICL risk into Bayes Gap and posterior variance.
  • The amortized-inference view holds that a Transformer can approximate a dataset-level Bayesian posterior in a single forward pass.
  • Prior analyses show that, in linear and linearized regimes, Transformers can implement one-step gradient updates and attention can represent accumulated gradient corrections.
  • The appendix interprets these gradient-based results as algorithmic realizations of the Bayesian predictor rather than reproducing their derivations.

A.4 Task–prior representation under within-task exchangeability

The task–prior model underlying the Bayesian ICL analysis is justified by within-task exchangeability, rather than introduced as an independent neural-architecture assumption.

  • Within-task exchangeability is an assumption about the data-generating process, not about the neural architecture implementing in-context learning.
  • The appendix uses de Finetti–Hewitt–Savage representation to derive the hierarchical task–prior model from exchangeability.
  • The observations are modeled as sample pairs in X × Y generated within a single task episode.

A.4.1 Within-task exchangeability.

Within each task episode, the appendix assumes an exchangeable observation sequence: reordering examples leaves their joint distribution unchanged, while i.i.d. sampling is only a special case.

  • A finite sequence is exchangeable when its joint distribution is invariant under every permutation of its indices.
  • An infinite sequence is exchangeable when every finite prefix is exchangeable.
  • Exchangeability means that only the multiset of observations matters, and it is strictly weaker than the i.i.d. assumption.
  • Within-task exchangeability assumes that the conceptual infinite sequence of observation pairs (X_i, Y_i) is exchangeable.
  • The framework treats the observed dataset D_k as a finite prefix of an exchangeable sequence, with the infinite extension serving as a probability-theoretic idealization.
  • The standard Borel structure of typical supervised-learning spaces permits application of the general de Finetti–Hewitt–Savage theorem.

A.4.2 A de Finetti–Hewitt–Savage representation.

Under within-task exchangeability, de Finetti–Hewitt–Savage represents each task episode as a mixture of i.i.d. sequences, yielding the hierarchical task–prior model used for Bayesian ICL.

  • A de Finetti–Hewitt–Savage representation: The representation theorem states that an infinite exchangeable sequence is conditionally i.i.d. given a random probability measure, whose mixing law is unique.
  • A de Finetti–Hewitt–Savage representation: Under the exchangeability assumption, each episode therefore admits a mixture-of-i.i.d. representation.
  • From a random measure to a task prior: The random measure is identified with a task f, and its mixing distribution Π is renamed the task prior P(f).
  • From a random measure to a task prior: This task–prior representation is a canonical reparameterization of the exchangeable joint law, not an additional ad hoc modeling assumption.
  • From a random measure to a task prior: The construction identifies each task with a probability measure on X × Y and permits the conditional task-function view f: X → Δ(Y).
  • Remarks and scope: For finite datasets, the infinite-sequence formulation requires an exchangeable extension, while finite de Finetti results provide a conceptual approximate alternative.
  • Remarks and scope: The task–prior model is presented as an effective theory whose empirical plausibility for real pretraining corpora remains an open question, and practical transformers may exhibit order effects.
Loading 2609.05111v1…