Buckets:
| """Numerically stable distribution matching shared by SFT, OPD, and RL.""" | |
| from __future__ import annotations | |
| import math | |
| from typing import Literal | |
| import torch | |
| import torch.nn.functional as F | |
| def symmetric_jsd_from_log_probs( | |
| teacher_logp: torch.Tensor, student_logp: torch.Tensor | |
| ) -> torch.Tensor: | |
| """Symmetric JSD over the final class axis with teacher stop-gradient.""" | |
| if teacher_logp.shape != student_logp.shape or teacher_logp.ndim == 0: | |
| raise ValueError("teacher/student log-probabilities must have equal rank and shape") | |
| teacher = teacher_logp.detach().float() | |
| student = student_logp.float() | |
| log_mixture = torch.logaddexp(teacher, student) - math.log(2.0) | |
| teacher_kl = (teacher.exp() * (teacher - log_mixture)).sum(-1) | |
| student_kl = (student.exp() * (student - log_mixture)).sum(-1) | |
| # Roundoff can otherwise produce tiny negative values for identical | |
| # distributions or values just above the analytic log(2) bound. | |
| return (0.5 * (teacher_kl + student_kl)).clamp(0.0, math.log(2.0)) | |
| def symmetric_jsd_from_logits( | |
| teacher_logits: torch.Tensor, student_logits: torch.Tensor | |
| ) -> torch.Tensor: | |
| """Symmetric JSD over logits; only the student side remains differentiable.""" | |
| return symmetric_jsd_from_log_probs( | |
| F.log_softmax(teacher_logits.detach().float(), dim=-1), | |
| F.log_softmax(student_logits.float(), dim=-1), | |
| ) | |
| def contrast_shaped_log_target( | |
| full_teacher_logits: torch.Tensor, | |
| blind_teacher_logits: torch.Tensor, | |
| *, | |
| contrast_alpha: float, | |
| ratio_clip: float, | |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | |
| """Build the detached positive-ratio-shaped teacher distribution.""" | |
| if full_teacher_logits.shape != blind_teacher_logits.shape: | |
| raise ValueError("full/blind teacher logits must have equal shape") | |
| if contrast_alpha < 0 or ratio_clip <= 0: | |
| raise ValueError("invalid contrast alpha/ratio clip") | |
| full_logp = F.log_softmax(full_teacher_logits.detach().float(), dim=-1) | |
| blind_logp = F.log_softmax(blind_teacher_logits.detach().float(), dim=-1) | |
| ratio = full_logp - blind_logp | |
| shaped = F.log_softmax( | |
| full_logp + contrast_alpha * ratio.clamp(min=0.0, max=ratio_clip), | |
| dim=-1, | |
| ) | |
| return shaped.detach(), ratio.detach(), full_logp.detach(), blind_logp.detach() | |
| def realized_position_log_likelihood( | |
| log_probabilities: torch.Tensor, targets: torch.Tensor | |
| ) -> torch.Tensor: | |
| """Mean realized factor log-likelihood for each autoregressive position.""" | |
| if log_probabilities.ndim < 2: | |
| raise ValueError("log probabilities must include position and class axes") | |
| if log_probabilities.shape[:-1] != targets.shape: | |
| raise ValueError("target shape must match log probabilities without classes") | |
| realized = log_probabilities.gather(-1, targets.long().unsqueeze(-1)).squeeze(-1) | |
| if realized.ndim == 1: | |
| return realized | |
| return realized.flatten(1).mean(-1) | |
| def gt_positive_gate( | |
| full_logp: torch.Tensor, | |
| blind_logp: torch.Tensor, | |
| targets: torch.Tensor, | |
| *, | |
| delta: float = 0.0, | |
| ) -> torch.Tensor: | |
| """Return the strict-positive per-position GT necessity gate.""" | |
| if delta < 0: | |
| raise ValueError("GT-positive threshold must be non-negative") | |
| full = realized_position_log_likelihood(full_logp, targets) | |
| blind = realized_position_log_likelihood(blind_logp, targets) | |
| return (full - blind) > delta | |
| def sufficiency_gate( | |
| full_logp: torch.Tensor, | |
| local_logp: torch.Tensor, | |
| local_blind_logp: torch.Tensor, | |
| targets: torch.Tensor, | |
| *, | |
| delta: float = 0.0, | |
| ) -> torch.Tensor: | |
| """Return the per-position local sufficiency gate from the four views.""" | |
| if delta < 0: | |
| raise ValueError("sufficiency threshold must be non-negative") | |
| full = realized_position_log_likelihood(full_logp, targets) | |
| local = realized_position_log_likelihood(local_logp, targets) | |
| blind = realized_position_log_likelihood(local_blind_logp, targets) | |
| return ((local - blind) > 0.0) & (local >= full - delta) | |
| def event_jsd_loss( | |
| target_logp: torch.Tensor, | |
| student_logits: torch.Tensor, | |
| *, | |
| position_gate: torch.Tensor | None = None, | |
| factor_reduction: Literal["sum", "mean"] = "sum", | |
| ) -> tuple[torch.Tensor, torch.Tensor]: | |
| """Return all-position-normalized event JSD and per-position reductions.""" | |
| if target_logp.shape != student_logits.shape or target_logp.ndim < 2: | |
| raise ValueError("target/student distributions must have equal rank and shape") | |
| if factor_reduction not in {"sum", "mean"}: | |
| raise ValueError("factor_reduction must be sum or mean") | |
| factor_jsd = symmetric_jsd_from_log_probs( | |
| target_logp, F.log_softmax(student_logits.float(), dim=-1) | |
| ) | |
| if factor_jsd.ndim == 1: | |
| per_position = factor_jsd | |
| else: | |
| factors = factor_jsd.flatten(1) | |
| per_position = ( | |
| factors.sum(-1) if factor_reduction == "sum" else factors.mean(-1) | |
| ) | |
| if not per_position.numel(): | |
| raise ValueError("event must contain at least one target position") | |
| if position_gate is not None: | |
| if position_gate.shape != per_position.shape: | |
| raise ValueError("position gate must contain one value per target position") | |
| per_position = per_position * position_gate.to(per_position) | |
| return per_position.sum() / per_position.numel(), per_position | |
| __all__ = [ | |
| "contrast_shaped_log_target", | |
| "event_jsd_loss", | |
| "gt_positive_gate", | |
| "realized_position_log_likelihood", | |
| "sufficiency_gate", | |
| "symmetric_jsd_from_log_probs", | |
| "symmetric_jsd_from_logits", | |
| ] | |
Xet Storage Details
- Size:
- 5.71 kB
- Xet hash:
- 882b3c13c02191bd0423a7cdcbb46cd83ea004784e0e5c61954d53c1b0bd40d0
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.