| """Objective-preserving primitives for InfoFlow training. |
| |
| This module intentionally depends only on PyTorch. It contains no dataset or |
| model-specific policy, so its contracts can be tested before sampler wiring. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import torch |
|
|
|
|
| @dataclass(frozen=True) |
| class SafeMixture: |
| """A normalized safe proposal and its exact importance correction.""" |
|
|
| base_probability: torch.Tensor |
| proposal_probability: torch.Tensor |
| sampling_probability: torch.Tensor |
| importance_weight: torch.Tensor |
| max_importance_weight: float |
|
|
|
|
| def _normalized_probability(probability: torch.Tensor, *, name: str) -> torch.Tensor: |
| if not isinstance(probability, torch.Tensor): |
| raise TypeError(f"{name} must be a torch.Tensor") |
| if probability.ndim != 1 or probability.numel() == 0: |
| raise ValueError(f"{name} must be a non-empty one-dimensional tensor") |
| if probability.dtype == torch.bool or not probability.dtype.is_floating_point: |
| raise TypeError(f"{name} must have a floating-point dtype") |
|
|
| probability64 = probability.detach().to(dtype=torch.float64) |
| if not torch.isfinite(probability64).all(): |
| raise ValueError(f"{name} must contain only finite values") |
| if (probability64 < 0).any(): |
| raise ValueError(f"{name} must be non-negative") |
|
|
| total = probability64.sum() |
| if total <= 0: |
| raise ValueError(f"{name} must have a positive sum") |
| return probability64 / total |
|
|
|
|
| def build_safe_mixture( |
| base_probability: torch.Tensor, |
| proposal_probability: torch.Tensor, |
| *, |
| alpha: float, |
| ) -> SafeMixture: |
| """Build ``q = (1 - alpha) p + alpha r`` and exact ``p / q`` weights.""" |
|
|
| if base_probability.device != proposal_probability.device: |
| raise ValueError("base_probability and proposal_probability must be on the same device") |
| if base_probability.shape != proposal_probability.shape: |
| raise ValueError("base_probability and proposal_probability must have the same shape") |
| if not isinstance(alpha, (int, float)) or isinstance(alpha, bool) or not math.isfinite(float(alpha)): |
| raise TypeError("alpha must be a finite real number") |
| alpha = float(alpha) |
| if not 0.0 <= alpha < 1.0: |
| raise ValueError("alpha must satisfy 0 <= alpha < 1") |
|
|
| p = _normalized_probability(base_probability, name="base_probability") |
| r = _normalized_probability(proposal_probability, name="proposal_probability") |
| q = (1.0 - alpha) * p + alpha * r |
|
|
| positive_support = p > 0 |
| if not (q[positive_support] > 0).all(): |
| raise RuntimeError("safe mixture lost support for a positive-probability baseline item") |
|
|
| weight = torch.zeros_like(p) |
| weight[positive_support] = p[positive_support] / q[positive_support] |
| max_weight = 1.0 / (1.0 - alpha) |
| if (weight > max_weight + 1e-12).any(): |
| raise RuntimeError("safe-mixture importance bound was violated") |
|
|
| return SafeMixture( |
| base_probability=p, |
| proposal_probability=r, |
| sampling_probability=q, |
| importance_weight=weight.detach(), |
| max_importance_weight=max_weight, |
| ) |
|
|
|
|
| def reduce_flow_matching_loss(elementwise_loss: torch.Tensor, *, reduction: str = "mean") -> torch.Tensor: |
| """Reduce flow loss while retaining the original batch dimension.""" |
|
|
| if not isinstance(elementwise_loss, torch.Tensor): |
| raise TypeError("elementwise_loss must be a torch.Tensor") |
| if elementwise_loss.ndim == 0 or elementwise_loss.shape[0] == 0: |
| raise ValueError("elementwise_loss must have a non-empty batch dimension") |
| if reduction not in {"mean", "none"}: |
| raise ValueError("reduction must be either 'mean' or 'none'") |
|
|
| |
| |
| if reduction == "mean": |
| return elementwise_loss.mean() |
|
|
| if elementwise_loss.ndim == 1: |
| per_sample = elementwise_loss |
| else: |
| per_sample = elementwise_loss.mean(dim=tuple(range(1, elementwise_loss.ndim))) |
| return per_sample |
|
|
|
|
| def aggregate_repeated_losses( |
| repeated_per_sample_loss: torch.Tensor, |
| *, |
| repeat_count: int, |
| batch_size: int | None = None, |
| ) -> torch.Tensor: |
| """Average repeat-major stochastic flow draws back to shape ``[B]``.""" |
|
|
| if not isinstance(repeated_per_sample_loss, torch.Tensor): |
| raise TypeError("repeated_per_sample_loss must be a torch.Tensor") |
| if repeated_per_sample_loss.ndim != 1 or repeated_per_sample_loss.numel() == 0: |
| raise ValueError("repeated_per_sample_loss must be a non-empty one-dimensional tensor") |
| if not isinstance(repeat_count, int) or isinstance(repeat_count, bool) or repeat_count <= 0: |
| raise ValueError("repeat_count must be a positive integer") |
| if repeated_per_sample_loss.numel() % repeat_count != 0: |
| raise ValueError("loss count must be divisible by repeat_count") |
|
|
| inferred_batch_size = repeated_per_sample_loss.numel() // repeat_count |
| if batch_size is not None and batch_size != inferred_batch_size: |
| raise ValueError("batch_size does not match the repeated loss shape") |
| return repeated_per_sample_loss.reshape(repeat_count, inferred_batch_size).mean(dim=0) |
|
|
|
|
| def ordinary_importance_corrected_mean( |
| per_sample_loss: torch.Tensor, |
| importance_weight: torch.Tensor, |
| ) -> torch.Tensor: |
| """Return the ordinary (not self-normalized) importance estimator.""" |
|
|
| if not isinstance(per_sample_loss, torch.Tensor) or not isinstance(importance_weight, torch.Tensor): |
| raise TypeError("per_sample_loss and importance_weight must be torch.Tensor instances") |
| if per_sample_loss.ndim != 1 or per_sample_loss.numel() == 0: |
| raise ValueError("per_sample_loss must be a non-empty one-dimensional tensor") |
| if importance_weight.shape != per_sample_loss.shape: |
| raise ValueError("importance_weight must have the same shape as per_sample_loss") |
|
|
| detached_weight = importance_weight.detach() |
| if not torch.isfinite(detached_weight).all() or (detached_weight < 0).any(): |
| raise ValueError("importance_weight must be finite and non-negative") |
|
|
| work_dtype = torch.float64 if per_sample_loss.dtype == torch.float64 else torch.float32 |
| loss = per_sample_loss.to(dtype=work_dtype) |
| weight = detached_weight.to(device=per_sample_loss.device, dtype=work_dtype) |
| return (weight * loss).mean() |
|
|
