Reinforce Ada Adaptive Sampling Rl Llms

Reinforce-Ada: Adaptive Sampling Framework for Non-Linear RL Objectives

Core Concept

Standard group-based RL (GRPO) suffers signal loss when small groups produce identical rewards—advantages normalize to zero, causing gradient collapse. The root cause is undersampling difficult prompts, not model limitations. Reinforce-Ada recovers learning signals by dynamically allocating more samples to harder prompts, weighted by an implicit importance function derived from optimizing non-linear reward objectives.

Architecture Overview

• Non-Linear Objective Optimization: Frame RL as maximizing f(p_θ(x)) where f is log, power, or other monotonic function and p_θ(x) is per-prompt pass rate
• Implicit Weighting: Optimal gradient estimator naturally prioritizes difficult prompts with weight proportional to f'(p)
• Two Implementation Strategies: (1) Estimation-based: explicit reweighting with value network estimates; (2) Sequential: successive elimination of solved prompts
• Budget Efficiency: Allocate inference compute dynamically rather than uniformly across all prompts

Implementation Steps

1. Problem Formulation and Theoretical Framework

Define weighted RL objective:

J_f(θ) = E_x[f(p_θ(x))]

where f is a non-linear function of per-prompt pass rate. The gradient naturally acquires prompt-dependent weights:

# Non-linear objective framework
def weighted_rl_objective(pass_rates, f_func, f_prime_func):
    """
    J_f(θ) = E_x[f(p_θ(x))]
    ∇J_f = E_x[f'(p_θ(x)) · ∇p(x)]
    """

    # Log-objective: f(p) = log(p), f'(p) = 1/p
    # Weights difficult prompts (low p) more heavily
    def log_objective():
        weights = 1.0 / (pass_rates + 1e-8)  # Avoid division by zero
        return weights

    # Power objective: f(p) = p^α, f'(p) = α · p^(α-1)
    # Softer weighting scheme
    def power_objective(alpha=0.5):
        weights = alpha * (pass_rates ** (alpha - 1))
        return weights

    # Use log-objective as primary weighting strategy
    return log_objective()

Key insight: Log-objective weighting (f'(p) = 1/p) is variance-optimal under squared-loss constraints and naturally prioritizes difficult prompts.

2. Algorithm 1: Reinforce-Ada-Estimation

Estimation-based approach using value network to predict per-prompt difficulty, then allocate samples accordingly.

def reinforce_ada_estimation(
    prompts, model, value_network,
    total_budget=512, n_min=4, n_max=32, num_phases=2
):
    """
    Phase 1: Estimate difficulty → allocate sample budget
    Phase 2: Train with explicit weighting
    """

    # Phase 1: Difficulty estimation via value network
    pass_rates = []
    for x in prompts:
        # Estimate: p̂_t = (N_pos^(t) + α) / (N_total^(t) + α + β)
        # Use exponential moving average (EMA) with Bayesian prior
        v = value_network(x)
        p_hat = torch.sigmoid(v)  # Convert to probability
        pass_rates.append(p_hat)

    pass_rates = torch.stack(pass_rates)

    # Allocate samples: N_x ∝ 1/√(p̂_x + ε)
    # Difficult prompts get more samples
    weights = 1.0 / torch.sqrt(pass_rates + 1e-8)
    sample_allocation = (weights / weights.sum()) * total_budget

    # Clip to [N_min, N_max] range
    sample_counts = torch.clamp(sample_allocation, n_min, n_max).int()

    # Phase 2: Training with explicit weighting
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-6)

    for prompt_idx, num_samples in enumerate(sample_counts):
        # Generate num_samples responses
        responses = [model.generate(prompts[prompt_idx]) for _ in range(num_samples)]

        # Evaluate responses
        rewards = evaluate_responses(responses)

        # Compute baseline per prompt
        baseline = rewards.mean()

        # Explicit weight: α_x = 1/√(p̂_x + ε)
        prompt_weight = 1.0 / torch.sqrt(pass_rates[prompt_idx] + 1e-8)

        # Compute weighted advantage
        advantages = (rewards - baseline) * prompt_weight

        # Policy gradient update
        loss = -torch.log(model.policy(responses)).mean() * advantages.mean()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    return model

Configuration: Sample bounds [4, 32], EMA decay λ ∈ (0, 1), Bayesian prior Beta(α, β).

3. Algorithm 2: Reinforce-Ada-Sequential

Model-free approach using successive elimination—generate multiple samples per round and progressively remove solved prompts.

def reinforce_ada_sequential(
    prompts, model,
    budget_per_round=512, samples_per_round=16,
    k_pos=16, k_neg=8, n_downsampled=16
):
    """
    Successive elimination approach:
    - Multiple rounds of generation
    - Exit when K_pos correct responses collected (or K_pos correct + K_neg incorrect for balanced)
    - Downsample to fixed size, reweight explicitly
    """

    all_results = {idx: {'correct': [], 'incorrect': []} for idx in range(len(prompts))}
    active_prompts = set(range(len(prompts)))

    # Phase 1: Adaptive collection (multiple rounds)
    round_num = 0
    while active_prompts and round_num < 100:
        for prompt_idx in list(active_prompts):
            # Generate M responses per round
            for _ in range(samples_per_round):
                response = model.generate(prompts[prompt_idx])
                is_correct = evaluate_response(response)

                if is_correct:
                    all_results[prompt_idx]['correct'].append(response)
                else:
                    all_results[prompt_idx]['incorrect'].append(response)

            # Exit conditions
            n_correct = len(all_results[prompt_idx]['correct'])
            n_incorrect = len(all_results[prompt_idx]['incorrect'])

            # Option 1: Positive-focused (stop after K_pos successes)
            if n_correct >= k_pos:
                active_prompts.discard(prompt_idx)

            # Option 2: Balanced (stop after K_pos correct AND K_neg incorrect)
            elif n_correct >= k_pos and n_incorrect >= k_neg:
                active_prompts.discard(prompt_idx)

        round_num += 1

    # Phase 2: Static batch construction with reweighting
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-6)

    for prompt_idx, results in all_results.items():
        # High-fidelity pass rate estimate
        n_total = len(results['correct']) + len(results['incorrect'])
        p_hat = len(results['correct']) / max(n_total, 1)

        # Downsample to fixed size (e.g., n=16)
        # Balanced sampling: n/2 correct, n/2 incorrect
        sampled_correct = results['correct'][:n_downsampled//2]
        sampled_incorrect = results['incorrect'][:n_downsampled//2]

        # Compute explicit weight for reweighting
        weight = 1.0 / (p_hat + 1e-8)

        # Policy gradient with reweighting
        for response in sampled_correct:
            log_prob = torch.log(model.policy(response))
            advantage = (1.0 - p_hat) * weight  # Positive advantage
            loss = -log_prob * advantage
            optimizer.zero_grad()
            loss.backward()

        for response in sampled_incorrect:
            log_prob = torch.log(model.policy(response))
            advantage = (-p_hat) * weight  # Negative advantage
            loss = -log_prob * advantage
            optimizer.zero_grad()
            loss.backward()

        optimizer.step()

    return model

Sequential elimination achieves K_pos=16 (positive) or K_pos=K_neg=8 (balanced) before moving to next prompt.

4. Hyperparameter Configuration

Tuning is minimal—learning rate fixed at 1e-6 across model families. Key parameters control budget allocation and sampling strategy.

# Experimental configuration
training_config = {
    'learning_rate': 1e-6,  # Fixed across all settings
    'entropy_coefficient': 1e-4,
    'clipping_range': (0.2, 0.28),
    'training_steps': 600,
    'no_kl_penalty': True,

    # Sampling strategy (Ada-Est)
    'sample_bounds': (4, 32),  # N_min, N_max
    'total_budget_per_iter': 512,
    'ema_decay': 0.99,  # Bayesian EMA

    # Sampling strategy (Ada-Seq)
    'samples_per_round': 16,
    'k_pos': 16,  # Stop after 16 correct
    'k_pos_balanced': 8,  # Balanced: 8 correct + 8 incorrect
    'n_downsampled': 16,  # Downsample to 16 for training
}

# Models tested: Qwen2.5-Math (1.5B, 7B), Qwen3-4B, Llama-3.2-3B
# Benchmarks: MATH500, Minerva Math, OlympiadBench, AIME-like

Practical Guidance

Signal Loss Root Cause: When all group members get identical rewards, standardized advantages collapse to zero. This is statistical undersampling, not model failure—fix by allocating more samples to difficult prompts.

Weighting Strategy: Log-objective (weight = 1/p) is both theoretically variance-optimal and empirically robust. Power functions (p^α) offer softer weighting if needed.

Algorithm Selection: Ada-Seq (successive elimination) avoids requiring value network estimates and is model-free. Ada-Est scales better for very large prompt sets but requires training value network.

Convergence: Expect 2× faster convergence to comparable performance using identical total compute. Wall-clock overhead: 1.3-2.8× depending on implementation, but offset by training speedup.

When to Use / When NOT to Use

Use When:

• Training reasoning models (math, coding, STEM) with verifiable rewards
• Small-to-medium group sizes (n=4-8) causing gradient collapse in GRPO
• You have computational budget for adaptive allocation (1.3-2.8× overhead acceptable)
• Task difficulty varies substantially across examples

NOT For:

• When all examples have similar difficulty (uniform sampling sufficient)
• Very large-scale distributed training where centralized difficulty tracking is infeasible
• Tasks where pass rate estimation is unreliable or undefined

Reference

This skill encodes techniques from "Reinforce-Ada: An Adaptive Sampling Framework under Non-linear RL Objectives" (arXiv:2510.04996). Code available at https://github.com/RLHFlow/Reinforce-Ada.