File size: 6,506 Bytes
8101696 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """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'")
# Preserve the legacy scalar path's reduction order exactly. Computing
# per-sample means first can introduce a tiny floating-point difference.
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()
|