Reward-Weighted Sampling in ML & RL

• Reward-Weighted Sampling is a framework that uses explicit reward signals to bias the selection of samples, actions, or updates toward high-utility outcomes.
• It finds applications in policy optimization, variational inference, and large model decoding, thereby enhancing efficiency and alignment in various ML tasks.
• Its theoretical foundation is built on importance sampling, expectation-maximization, and policy search, which underpin its robust convergence and performance improvements.

Reward-Weighted Sampling (RWS) is a broad methodological framework in machine learning and reinforcement learning for selecting samples, actions, or model updates with probability proportional to their associated reward, quality, or utility signal. RWS has found applications in variational inference, policy search, distributed gradient aggregation, adaptive representation learning, dynamic graph sampling, LLM decoding, preference alignment, and offline reinforcement learning. Its central premise is to bias learning or inference toward regions of higher utility, using explicit reward scores either as sampling probabilities or as normalization weights in update rules.

1. Theoretical Underpinnings and Formulations

Reward-Weighted Sampling formalizes the principle of weighting candidate selections (such as actions, model samples, or algorithmic updates) by their reward or importance. Canonically, in the context of policy improvement and variational inference, this takes the form of:

• Self-normalized importance sampling estimator:

$$\mathbb{E}_q \left[ f(z) \right] \approx \sum_{k=1}^K \frac{w(z^k)}{\sum_{l=1}^K w(z^l)} f(z^k),$$

where $w(z^k)$ is a reward or importance score (often $w(z^k) = p(z^k)/q(z^k)$ or a function of reward/value).

• Reward-Weighted Policy Update in RL:

$$\pi_{n+1}(a|s) = \frac{Q^{(\pi_n)}(s, a) \, \pi_n(a|s)}{V^{(\pi_n)}(s)}, \quad V^{(\pi_n)}(s) = \int Q^{(\pi_n)}(s,a) \, \pi_n(a|s) da$$

as established for Reward-Weighted Regression (Štrupl et al., 2021), with monotonic improvement leading to global convergence.

• Gradient variants in variational inference (AISLE):

$$\nabla_\phi^{\text{AISLE}} = \sum_k \frac{w(z^k)}{\sum_l w(z^l)} \Delta_\psi(z^k)$$

where the estimator recovers reparameterized gradients such as those in IWAE-STL and IWAE-DREG (Finke et al., 2019).

RWS emphasizes the use of reward as a priority or probability, providing principled objectives and update rules that favor regions of high utility.

2. Connections to Expectation-Maximization, Importance Sampling, and Policy Search

RWS derives from foundational approaches in expectation-maximization (EM), importance sampling, and policy iteration:

• Expectation-Maximization: In RL, RWS updates can be seen as maximizing a reward-weighted log-likelihood, updating policies to assign higher probability to higher-reward actions (Štrupl et al., 2021).
• Importance Sampling: The principle of using unnormalized weights from a target distribution to reweight samples from a proposal (e.g., in variational inference) is central to RWS formulations (Finke et al., 2019).
• Policy Search: Iterative RWS-based policy updates yield monotonic improvement and theoretical convergence to optimal policies in both compact and finite state-action spaces (Štrupl et al., 2021).

This lineage underlies the robustness and theoretical guarantees associated with reward-weighted methods.

3. Variations, Generalizations, and Algorithmic Specializations

RWS supports several important variants and can be generalized through different choices of divergence, weighting, or sampling mechanisms:

Variant / Extension	Mechanism	Application Context
AISLE (Finke et al., 2019)	Adaptive divergence minimization	Variational inference, VI
IWAE-STL/DREG	Score-function-free path gradients	Multi-sample variational bounds
Reservoir Sampling	Weighted selection, skip algorithms	Streaming, distributed sampling
RL Weighted Sampling	RL-based weighting of dynamic data	Subgraph counting, graph streams
Reachability (Yang et al., 3 Jun 2025)	PU-learned classifier weights	Goal-conditioned RL
Cascade Sampling (Li et al., 24 Jun 2024)	Segment-level reward-based rejection	LLM decoding-time alignment

AISLE provides a unifying theory, showing that reweighted wake-sleep (classic RWS), IWAE-style multi-sample objectives, and divergence-based gradients (KL, $\chi^2$) are all instantiations of reward-weighted sampling. Weighted reservoir methods further generalize RWS, allowing efficient, streaming, and distributed sampling.

4. Practical Implementations: Reinforcement Learning, Sampling, and Large Model Decoding

RL and Policy Optimization: In RL, RWS underpins iterative algorithms that guarantee monotonic improvement and (in finite MDPs) R-linear convergence to the optimum (Štrupl et al., 2021). The update step uses the reward-weighted regression formula. Policy gradient mergers in distributed systems use episodic R-weighted aggregation to emphasize informative gradients (Holen et al., 2023).

Graph Streams: RL-enhanced weighted sampling (e.g., WSD (Wang et al., 2022)) employs continuous action RL to assign data-driven weights to streaming graph edges, using experience replay and actor-critic architectures to minimize estimation error in subgraph counting.

Variational Inference: AISLE (Finke et al., 2019) generalizes RWS for adaptive proposal learning, minimizing divergences with importance weights, and incorporating reparameterization for efficient score-function-free gradients.

LLMs and Diffusion Models:

• Reward-weighted decoding in Masked Diffusion Models guides token selection by externally provided reward signals, enhancing non-autoregressive characteristics and global coherence (Gwak et al., 31 Aug 2025).
• Cascade Reward Sampling (CARDS) introduces segment-level rejection sampling for decoding-time alignment, using reward models and uncertainty-based segmentation to ensure efficient, high-quality text generation (Li et al., 24 Jun 2024).

5. Empirical Performance, Advantages, and Limitations

RWS is empirically validated across a range of tasks:

• Variance Reduction: AISLE-based score-function-free variants (IWAE-STL, IWAE-DREG) exhibit near-zero variance when the proposal matches the target and avoid gradient breakdown for large sample sizes (Finke et al., 2019).
• Efficiency: Reward-weighted graph sampling yields up to 47% lower estimation error and often faster runtime compared to uniform sampling (Wang et al., 2022).
• Alignment Quality: RewardSDS produces state-of-the-art text-to-image and text-to-3D generation via reward-weighted loss, improving CLIPScore and user-rated aesthetics (Chachy et al., 12 Mar 2025).
• LLM Decoding: RWS decoding in MDMs increases generation order deviation (less sequential order), and improves win rate and coherence on multiple benchmarks (Gwak et al., 31 Aug 2025).
• Distributed RL: R-weighted gradient aggregation shows modest but consistent improvement and stability in heterogeneous environments (Holen et al., 2023).
• Offline RL: Reachability-weighted sampling yields up to 50% improvement in manipulation tasks by prioritizing actionable goals (Yang et al., 3 Jun 2025).

Limitations include:

• Sensitivity to reward distributions (extreme weights can cause sample domination or stability issues).
• Numerical challenges in high-throughput scenarios, especially with weighted reservoir or skip algorithms (Meligrana, 29 Mar 2024).
• Dependence on reward model quality—suboptimal reward signals can misguide sampling or optimization (as observed in large-scale LLM alignment and representation tasks).

6. Modern Extensions: Preference Alignment, Adaptive and Reinforcement-Learning-Based Sampling

Recent research generalizes RWS toward optimal human alignment and adaptive sampling strategies:

• Maximum Preference Optimization (MPO) (Jiang et al., 2023) reframes preference learning as reward maximization via importance sampling, enabling effective off-policy training with KL regularization and memory efficiency.
• PILAF (Feng et al., 6 Feb 2025) proposes policy-interpolated sampling distributions, aligning gradients of the reward model with the oracle objective, theoretically and empirically improving RLHF sample efficiency and reward maximization.
• Adaptive Sampling with Reward (ASR) (Dou et al., 2022) leverages RL to adjust sampling distributions dynamically in representation learning tasks, outperforming heuristic and static methods and exhibiting phenomena such as “ASR gravity well.”
• ELHSR (Guo et al., 18 May 2025) introduces ultra-efficient hidden-state-based reward models for LLMs, dramatically increasing sampling efficiency and enabling hybrid reward-weighted selection.

These developments reflect a shift from static or heuristically weighted sampling toward highly adaptive, optimization-aligned, and even RL-driven sample selection.

7. Impact, Applications, and Future Directions

RWS and its extensions have led to substantive impact in probabilistic inference, RL, structured data sampling, representation learning, preference alignment, and generative modeling. Key application areas include:

• Variational Autoencoders and Generative Models: RWS generalizes importance sampling for proposal distribution learning, variance reduction, and robust inference (Finke et al., 2019).
• Policy Optimization and Distributed RL: Gradient-merging strategies and reward-weighted updates enhance learning stability and convergence in multi-agent systems (Štrupl et al., 2021, Holen et al., 2023).
• Preference Learning in LLMs: Off-policy algorithms (MPO, PILAF) and reward-weighted decoding provide scalable, stable, and efficient alignment with human values (Jiang et al., 2023, Feng et al., 6 Feb 2025).
• Diffusion-Based Image and Text Generation: Reward-weighted score distillation directly incorporates reward signals during gradient updates for fine-grained alignment (Chachy et al., 12 Mar 2025).
• Representation Learning and Clustering: RL-driven adaptive sampling achieves optimal sample selection and overcomes static method limitations (Dou et al., 2022).
• Offline RL and Goal-Conditioned Learning: Reachability-weighted sampling filters out unreachable goal syntheses, boosting performance on complex manipulation tasks (Yang et al., 3 Jun 2025).

Future directions include tighter integration with reinforcement learning, adaptive reward weighting schemes, hybrid and ensemble reward models, advanced handling of weight extrema, scalable distributed procedures, and theoretical guarantees connecting sample-induced gradients to global optimization objectives.

---

Reward-Weighted Sampling provides a principled, theoretically justified, and highly versatile toolkit for leveraging reward signals to bias learning, inference, and data selection toward optimal or high-utility outcomes, with robust applications across the spectrum of modern machine learning research.