| """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) |
| |
| 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}") |
|
|
