src/jolt/rlcd.py

"""Joint RLCD and cross-entropy loss aligned with pinned Laya references.

See docs/rlcd.md for upstream revisions, reduction semantics, and limitations.
"""

import math
from dataclasses import dataclass, field

import torch
from torch import Tensor


@dataclass(frozen=True)
class RLCDConfig:
    """The two RLCD controls; the scoring rule and estimator are fixed."""

    samples: int = field(default=4, init=False)
    sigma: float = 0.3
    log_weight: float = field(default=1.0, init=False)
    spherical_weight: float = field(default=0.75, init=False)
    ordinal_weight: float = field(default=1.0, init=False)
    ce_weight: float = 1.0
    log_floor: float = field(default=-9.21, init=False)

    def __post_init__(self):
        if not math.isfinite(self.sigma) or self.sigma <= 0:
            raise ValueError("sigma must be positive and finite")
        if not math.isfinite(self.ce_weight) or self.ce_weight < 0:
            raise ValueError("ce_weight must be nonnegative and finite")


def distribution_rewards(
    actions: Tensor,
    targets: Tensor,
    mask: Tensor,
    ordinal: Tensor,
    config: RLCDConfig,
) -> Tensor:
    """Score [G,B,K] sampled logits; padded candidates have zero probability.

    Counts use semantic ascending order. RPS excludes the final (always-one)
    cumulative probability and is divided by K-1, using each row's actual K.
    """
    log_probs = actions.float().masked_fill(~mask, -torch.inf).log_softmax(-1)
    probabilities = log_probs.exp()
    if targets.ndim == 1:
        targets = torch.nn.functional.one_hot(targets, actions.shape[-1]).float()
    log_score = probabilities.clamp_min(1e-12).log().clamp_min(config.log_floor)
    log_score = (log_score * targets).sum(-1)
    target_probability = (probabilities * targets).sum(-1)
    spherical = target_probability / probabilities.norm(dim=-1).clamp_min(1e-9)
    reward = config.log_weight * log_score + config.spherical_weight * spherical
    levels = torch.arange(actions.shape[-1], device=actions.device)
    target_cdf = targets.cumsum(-1)
    counts = mask.sum(-1)
    cdf_mask = levels[None, :] < (counts - 1)[:, None]
    rps = ((probabilities.cumsum(-1) - target_cdf).square() * cdf_mask).sum(-1)
    rps = rps / (counts - 1).clamp_min(1)
    return reward - config.ordinal_weight * rps * ordinal


def rlcd_loss(
    logits: Tensor,
    targets: Tensor,
    *,
    mask: Tensor | None = None,
    ordinal: Tensor | None = None,
    config: RLCDConfig | None = None,
) -> tuple[Tensor, dict[str, Tensor]]:
    """REINFORCE estimate; G perturbations never repeat the backbone forward.

    Actions AND rewards are detached. Without detaching sampled actions, the
    location parameter cancels in the Gaussian log density and its gradient
    vanishes. Match Laya's per-question group-mean baseline, then normalize with
    ONE sample standard deviation across G and B. This practical normalization,
    reward clipping, and auxiliary CE do not guarantee calibrated probabilities.
    """
    if config is None:
        config = RLCDConfig()
    if logits.ndim != 2 or targets.shape not in (logits.shape[:1], logits.shape):
        raise ValueError("Expected logits [B,K] and targets [B] or distributions [B,K]")
    if mask is None:
        mask = torch.ones_like(logits, dtype=torch.bool)
    if ordinal is None:
        ordinal = torch.zeros(logits.shape[0], dtype=torch.bool, device=logits.device)
    mean = logits.float()
    noise = torch.randn((config.samples, *mean.shape), device=mean.device) * config.sigma * mask
    noise = (noise - noise.sum(-1, keepdim=True) / mask.sum(-1, keepdim=True)) * mask
    actions = (mean.detach()[None] + noise).detach()
    with torch.no_grad():
        rewards = distribution_rewards(actions, targets, mask, ordinal, config)
        advantages = rewards - rewards.mean(0, keepdim=True)
        advantages = advantages / (advantages.std() + 1e-6)
    # Gaussian normalizing constants are independent of trainable parameters.
    log_density = -0.5 * ((actions - mean[None]) / config.sigma).square()
    log_density = log_density.masked_fill(~mask, 0).sum(-1)
    policy_loss = -(advantages * log_density).mean()
    ce_loss = -mean.masked_fill(~mask, -torch.inf).log_softmax(-1)
    if targets.ndim == 1:
        ce_loss = ce_loss.gather(-1, targets[:, None]).mean()
    else:
        ce_loss = (ce_loss.masked_fill(~mask, 0) * targets).sum(-1).mean()
    loss = policy_loss + config.ce_weight * ce_loss
    return loss, {
        "reward_mean": rewards.mean(),
        "reward_std": rewards.std(unbiased=False),
        "advantage_rms": advantages.square().mean().sqrt(),
        "policy_loss": policy_loss.detach(),
        "ce_loss": ce_loss.detach(),
    }


def _direct_components(logits: Tensor, targets: Tensor, mask: Tensor, ordinal: Tensor):
    """Stable differentiable CE, spherical score, and ordered RPS components."""
    scores = logits.float().masked_fill(~mask, -torch.inf)
    log_probabilities = scores.log_softmax(-1)
    probabilities = log_probabilities.exp()
    if targets.ndim == 1:
        targets = torch.nn.functional.one_hot(targets, logits.shape[-1]).float()
    target = targets.float().masked_fill(~mask, 0)
    ce = -(target * log_probabilities.masked_fill(~mask, 0)).sum(-1)
    spherical = (probabilities * target).sum(-1) / probabilities.norm(dim=-1).clamp_min(1e-9)
    counts = mask.sum(-1)
    levels = torch.arange(logits.shape[-1], device=logits.device)
    cdf_mask = levels[None, :] < (counts - 1)[:, None]
    rps = ((probabilities.cumsum(-1) - target.cumsum(-1)).square() * cdf_mask).sum(-1) / (
        counts - 1
    ).clamp_min(1)
    return ce, spherical, rps * ordinal


def direct_loss(
    logits: Tensor,
    targets: Tensor,
    *,
    mask: Tensor | None = None,
    ordinal: Tensor | None = None,
    composite: bool,
) -> tuple[Tensor, dict[str, Tensor]]:
    """Soft-target CE, optionally with the plan's spherical and ordinal terms."""
    if mask is None:
        mask = torch.ones_like(logits, dtype=torch.bool)
    if ordinal is None:
        ordinal = torch.zeros(logits.shape[0], dtype=torch.bool, device=logits.device)
    ce, spherical, rps = _direct_components(logits, targets, mask, ordinal)
    ce_loss = ce.mean()
    spherical_mean = spherical.mean()
    rps_mean = rps.mean()
    loss = ce_loss
    if composite:
        loss = loss + 0.75 * (1 - spherical_mean) + rps_mean
    return loss, {
        "ce_loss": ce_loss.detach(),
        "spherical_score": spherical_mean.detach(),
        "ordinal_rps": rps_mean.detach(),
    }


def pathwise_noisy_loss(
    logits: Tensor,
    targets: Tensor,
    *,
    mask: Tensor | None = None,
    ordinal: Tensor | None = None,
    config: RLCDConfig | None = None,
) -> tuple[Tensor, dict[str, Tensor]]:
    """Minimize CE minus the differentiable mean reward of noisy logits."""
    config = config or RLCDConfig()
    if mask is None:
        mask = torch.ones_like(logits, dtype=torch.bool)
    if ordinal is None:
        ordinal = torch.zeros(logits.shape[0], dtype=torch.bool, device=logits.device)
    mean = logits.float()
    ce, _, _ = _direct_components(mean, targets, mask, ordinal)
    noise = torch.randn((config.samples, *mean.shape), device=mean.device) * config.sigma * mask
    noise = (noise - noise.sum(-1, keepdim=True) / mask.sum(-1, keepdim=True)) * mask
    rewards = distribution_rewards(mean[None] + noise, targets, mask, ordinal, config)
    ce_loss = ce.mean()
    reward = rewards.mean()
    return ce_loss - reward, {"ce_loss": ce_loss.detach(), "noisy_reward": reward.detach()}


def decision_loss(
    logits: Tensor,
    targets: Tensor,
    *,
    loss: str,
    mask: Tensor,
    ordinal: Tensor,
    config: RLCDConfig,
) -> tuple[Tensor, dict[str, Tensor]]:
    """Dispatch a versioned experiment objective; `rlcd` retains the release loss."""
    if loss == "rlcd":
        return rlcd_loss(logits, targets, mask=mask, ordinal=ordinal, config=config)
    if loss == "ce":
        return direct_loss(logits, targets, mask=mask, ordinal=ordinal, composite=False)
    if loss == "direct-composite":
        return direct_loss(logits, targets, mask=mask, ordinal=ordinal, composite=True)
    if loss == "pathwise-noisy":
        return pathwise_noisy_loss(logits, targets, mask=mask, ordinal=ordinal, config=config)
    raise ValueError(f"Unsupported loss: {loss}")