llds-grpo-collapse

name	llds-grpo-collapse
title	On GRPO Collapse in Search-R1: Lazy Likelihood Displacement Suppression
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2512.04220
keywords	["reinforcement-learning","grpo-training","rl-stability","tool-integration","token-level-optimization"]
description	Identifies Lazy Likelihood Displacement (LLD) as GRPO failure mechanism in tool-integrated LLMs and proposes lightweight LLDS regularization that penalizes only tokens responsible for likelihood decrease, achieving up to 45.2% performance gains.

Summary

On GRPO Collapse identifies Lazy Likelihood Displacement (LLD) as the fundamental failure mechanism when training tool-integrated LLMs with GRPO, and proposes LLDS (LLD Suppression), a lightweight regularization that prevents unintended confidence reduction. The method activates only when a response's likelihood decreases and penalizes only responsible tokens.

Core Technique

Lazy Likelihood Displacement (LLD): During GRPO training, the policy can unintentionally reduce confidence in correct tool calls while improving on other aspects. This "lazy" displacement occurs because the RL objective optimizes for reward, not confidence preservation.

Fine-Grained Token Penalties: Rather than penalizing entire sequences, identify which specific tokens caused likelihood decrease and penalize only those:

likelihood_per_token_t = log p(token_t | context)
if likelihood_t < baseline_likelihood_t:
    penalty_t = -λ * (baseline_likelihood_t - likelihood_t)

Selective Activation: Only apply penalty when overall sequence likelihood decreases:

if likelihood_total < baseline_likelihood_total:
    apply_llds_penalty()

Implementation

Likelihood tracking: Compute baseline likelihood from reference model:

def compute_baseline_likelihood(sequence, reference_model):
    with torch.no_grad():
        logits = reference_model(sequence)
        likelihood = F.log_softmax(logits, dim=-1)
    return likelihood

Token-wise penalty computation:

def compute_llds_penalty(current_logits, baseline_likelihood, sequence):
    current_likelihood = F.log_softmax(current_logits, dim=-1)

    # Per-token likelihood change
    likelihood_change = current_likelihood - baseline_likelihood

    # Penalize tokens that decreased likelihood
    penalties = torch.where(
        likelihood_change < 0,
        -likelihood_change,  # Penalty proportional to decrease
        torch.zeros_like(likelihood_change)
    )

    return penalties.sum() / sequence.numel()