Add composable distillation loss hooks

TRAINING_INTEGRATION.md

@@ -0,0 +1,89 @@
+# Custom training integration
+
+Use FastGen losses in your own training loop with your models, optimizers,
+and checkpoints. FastGen's trainer uses the default unmasked losses and
+iteration-based phase selection. Supply weights, select phases, or manage
+ratio statistics explicitly when your loop needs that control.
+
+## Mask and weight losses
+
+`variational_score_distillation_loss`, `denoising_score_matching_loss`, and
+`gan_loss_generator` accept a keyword-only `loss_weight`. The discriminator loss
+accepts `real_weight` and `fake_weight`. Weights must be finite, nonnegative, and
+broadcast-compatible with the predictions or discriminator scores. Positive
+weights scale squared residuals or score losses linearly. Zero weights exclude
+entries from the loss and from VSD target normalization. Means divide by the
+full element count, including masked entries. Omit
+weights to include every entry with unit weight.
+
+```python
+loss = variational_score_distillation_loss(
+    generated, teacher_x0, fake_x0, loss_weight=pixel_weight,
+)
+```
+
+The sCM and MeanFlow `_scm_pred_to_loss` and `_mf_pred_to_loss` extension points
+accept the same keyword to exclude masked entries from tangent norms and
+adaptive weights. Samples with no valid entries contribute zero loss,
+including the sCM log-variance term.
+
+MeanFlow clips `t - r` to the noise scheduler's time bounds by default. Set
+`clip_time_interval=False` when `t - r` already expresses the duration required
+by your custom time parameterization. For DSM flow prediction, supply
+`flow_target` to replace the default velocity `eps - x0` with a conditional
+velocity from your custom path.
+
+For adversarial losses, exclude corrupted inputs during feature extraction
+and pooling, then supply weights for the resulting discriminator scores.
+A pixel mask alone cannot identify valid scores. Masking final scores cannot
+undo corruption that has already spread through the backbone.
+
+## Phase selection
+
+```python
+losses, outputs = model.single_train_step(batch, iteration, phase="auxiliary")
+```
+
+DMD2 and f-distill accept `phase="student"` or `phase="auxiliary"`
+independently of iteration. Pass the actual training step as `iteration` for
+warmup schedules, and step the optimizer matching your selected phase. Omit
+`phase` to follow `student_update_freq`, as FastGen's trainer does.
+
+`single_train_step` prepares inputs, samples noise, and sets model modes and
+gradients for the selected phase. Accept the `phase` keyword when overriding
+noise generation in a subclass. Use `_dispatch_training_phase` with prepared
+inputs when your loop also manages sampling, model modes, and frozen parameters.
+
+## Ratio state and process groups
+
+Use `f_divergence_weights` to compute f-distill weights while keeping running
+ratio statistics in your own training state:
+
+```python
+h, new_bins = f_divergence_weights(
+    scores, t, weighting_fn, bins=bins,
+    min_t=min_t, max_t=max_t,
+    logit_weight=score_weight,
+    statistics_group=data_group,
+    update_state=True,
+)
+```
+
+Save the returned bins for the next call or in a checkpoint. The function does
+not change the input bins. Set `update_state=False` to normalize with saved
+statistics without updating them. Omit `bins` to skip bin normalization.
+Zero `logit_weight` entries exclude scores from ratio statistics. Positive
+magnitudes only mark scores as valid. Use `sample_mask` to exclude whole
+samples. Samples with no valid scores receive zero weight.
+
+Call the function on every rank in `statistics_group`, even if a rank has no
+valid samples. This group should join data replicas, not spatial partitions
+of one sample. Combine spatial features within your model when needed. During
+distributed execution, omitting `statistics_group` uses the
+default process group.
+
+Normalization uses the actual global valid count, supporting unequal batch
+sizes. Scalar divergence functions broadcast to the batch, including reverse KL.
+FastGen's f-distill method uses this function and stores bins on its student
+network when you enable ratio normalization. Its checkpoint format stays
+unchanged. External loops can instead checkpoint the returned bins.

fastgen/methods/common_loss.py

@@ -7,6 +7,7 @@
 
 from fastgen.networks.noise_schedule import BaseNoiseSchedule
 from fastgen.utils import expand_like
+from fastgen.methods.loss_support import LossReduction
 
 
 def denoising_score_matching_loss(
@@ -16,13 +17,17 @@ def denoising_score_matching_loss(
     eps: torch.Tensor = None,
     noise_scheduler: Optional[BaseNoiseSchedule] = None,
     t: Optional[torch.Tensor] = None,
+    *,
+    loss_weight: Optional[torch.Tensor] = None,
+    flow_target: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    """Compute the denoising diffusion objective.
-        Forward process:
-            x_t = alpha_t * x_0 + sigma_t * eps
+    """Compute mean squared error against the selected denoising target.
 
-        Note: We currently don't add any loss weighting for simplicity. In the future, we may include
-            time-dependent weighting (e.g., by SNR, variance schedule).
+    The forward process is ``x_t = alpha_t * x_0 + sigma_t * eps``.
+    By default, every residual has unit weight. Supply ``loss_weight`` for
+    masking or time-dependent weighting. The mean divides by the full element
+    count. For ``pred_type="flow"``, ``flow_target``
+    replaces the default velocity ``eps - x0`` for a custom diffusion path.
 
     Args:
         pred_type (str): Prediction type: 'x0', 'eps', 'v', 'flow'.
@@ -31,6 +36,10 @@ def denoising_score_matching_loss(
         x0 (torch.Tensor): The clean data x_0.
         eps (torch.Tensor): The epsilon used to compute noised data.
         t (torch.Tensor): The target data t.
+        loss_weight (Optional[torch.Tensor]): Finite nonnegative weights that broadcast
+            to net_pred. Zero entries exclude residuals. Positive values scale squared errors.
+        flow_target (Optional[torch.Tensor]): Conditional velocity target used only
+            for flow prediction. Defaults to eps - x0.
 
     Raises:
         NotImplementedError: If an unknown pred_type is used.
@@ -40,44 +49,60 @@ def denoising_score_matching_loss(
     """
     if pred_type == "x0":
         assert x0 is not None, "x0 cannot be None"
-        loss = F.mse_loss(x0, net_pred, reduction="mean")
+        target = x0
     elif pred_type == "eps":
         assert eps is not None, "eps cannot be None"
-        loss = F.mse_loss(eps, net_pred, reduction="mean")
+        target = eps
     elif pred_type == "v":
         assert x0 is not None and eps is not None and t is not None, "x0, eps, t should not be None"
         assert noise_scheduler is not None, "noise_scheduler should not be None"
         alpha_t = expand_like(noise_scheduler.alpha(t), x0).to(device=x0.device, dtype=x0.dtype)
         sigma_t = expand_like(noise_scheduler.sigma(t), x0).to(device=x0.device, dtype=x0.dtype)
         v = alpha_t * eps - sigma_t * x0
-        loss = F.mse_loss(v, net_pred, reduction="mean")
+        target = v
     elif pred_type == "flow":
-        assert x0 is not None and eps is not None, "x0 and eps cannot be None"
-        flow_velocity = eps - x0
-        loss = F.mse_loss(flow_velocity, net_pred, reduction="mean")
+        if flow_target is None:
+            assert x0 is not None and eps is not None, "x0 and eps cannot be None"
+            target = eps - x0
+        else:
+            target = flow_target
     else:
         raise NotImplementedError(f"Unknown prediction type {pred_type}")
-    return loss
+    if loss_weight is None:
+        return F.mse_loss(target, net_pred, reduction="mean")
+    return LossReduction(net_pred, loss_weight).residual(net_pred - target, mean=True).mean()
 
 
 def variational_score_distillation_loss(
     gen_data: torch.Tensor,
     teacher_x0: torch.Tensor,
     fake_score_x0: torch.Tensor,
     additional_scale: Optional[torch.Tensor] = None,
+    *,
+    loss_weight: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """
     Compute the variational score distillation loss.
+
+    By default, all entries contribute to the target normalization and loss.
+    Supply ``loss_weight`` to mask entries or scale squared residuals. Zero
+    weights also exclude entries from the normalization statistics. Positive
+    magnitudes affect only the residual. The mean divides by the full element
+    count.
+
     Args:
         gen_data (torch.Tensor): generated data
         teacher_x0 (torch.Tensor): x0-prediction from the teacher
         fake_score_x0 (torch.Tensor): x0-prediction from the fake score
         additional_scale (Optional[torch.Tensor]): Additional scale parameter for the VSD loss.
+        loss_weight (Optional[torch.Tensor]): Finite nonnegative weights that broadcast
+            to gen_data. Zero entries exclude residuals and target normalization statistics.
 
     Returns:
         loss (torch.Tensor): The variational score distillation loss.
     """
     dims = tuple(range(1, teacher_x0.ndim))
+    reduction = LossReduction(gen_data, loss_weight) if loss_weight is not None else None
 
     with torch.no_grad():
         # Perform weight calculation in fp32 for numerical stability
@@ -86,7 +111,12 @@ def variational_score_distillation_loss(
         teacher_x0_fp32 = teacher_x0.float()
 
         # Compute weight in fp32 to avoid numerical instability
-        diff_abs_mean = (gen_data_fp32 - teacher_x0_fp32).abs().mean(dim=dims, keepdim=True)
+        if reduction is None:
+            diff_abs_mean = (gen_data_fp32 - teacher_x0_fp32).abs().mean(dim=dims, keepdim=True)
+        else:
+            difference = reduction.clean(gen_data_fp32 - teacher_x0_fp32)
+            diff_abs_mean = difference.abs().sum(dims) / reduction.count.clamp_min(1)
+            diff_abs_mean = expand_like(diff_abs_mean, gen_data)
         w_fp32 = 1 / (diff_abs_mean + 1e-6)
 
         # Apply additional scale if provided
@@ -97,40 +127,73 @@ def variational_score_distillation_loss(
         w = w_fp32.to(dtype=original_dtype)
 
         vsd_grad = (fake_score_x0 - teacher_x0) * w
+        if reduction is not None:
+            vsd_grad = reduction.clean(vsd_grad)
         pseudo_target = gen_data - vsd_grad
 
-    loss = 0.5 * F.mse_loss(gen_data, pseudo_target, reduction="mean")
+    loss = 0.5 * (
+        F.mse_loss(gen_data, pseudo_target, reduction="mean")
+        if reduction is None
+        else reduction.residual(gen_data - pseudo_target, mean=True).mean()
+    )
     return loss
 
 
-def gan_loss_generator(fake_logits: torch.Tensor) -> torch.Tensor:
+def gan_loss_generator(fake_logits: torch.Tensor, *, loss_weight: Optional[torch.Tensor] = None) -> torch.Tensor:
     """
-    Compute the GAN loss for the generator
+    Compute the generator's mean softplus loss over discriminator scores.
+
+    By default, every score contributes equally. Supply ``loss_weight`` to
+    exclude or scale score losses while retaining the full score count in the
+    mean. Determine valid scores during feature extraction and pooling. A
+    score mask cannot undo corruption that has spread through the backbone.
+
     Args:
         fake_logits (torch.Tensor): The logits for the fake data.
+        loss_weight (Optional[torch.Tensor]): Finite nonnegative weights that broadcast
+            to fake_logits. Zero entries exclude scores. Positive values scale their losses.
 
     Returns:
         gan_loss (torch.Tensor): The GAN loss for the generator.
 
     """
 
     assert fake_logits.ndim == 2, f"fake_logits has shape {fake_logits.shape}"
-    gan_loss = F.softplus(-fake_logits).mean()
+    if loss_weight is None:
+        gan_loss = F.softplus(-fake_logits).mean()
+    else:
+        reduction = LossReduction(fake_logits, loss_weight)
+        gan_loss = (reduction.clean(F.softplus(-reduction.clean(fake_logits))) * reduction.weight).mean()
     return gan_loss
 
 
-def gan_loss_discriminator(real_logits: torch.Tensor, fake_logits: torch.Tensor) -> torch.Tensor:
+def gan_loss_discriminator(
+    real_logits: torch.Tensor,
+    fake_logits: torch.Tensor,
+    *,
+    real_weight: Optional[torch.Tensor] = None,
+    fake_weight: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
     """
-    Compute the GAN loss for the discriminator
+    Sum the discriminator's mean softplus losses for real and fake scores.
+
+    By default, every score contributes equally. Use ``real_weight`` and
+    ``fake_weight`` to exclude or scale each set of score losses, with the
+    same reduction and feature-validity requirements as ``gan_loss_generator``.
+
     Args:
         real_logits (torch.Tensor): The logits for the real data.
         fake_logits (torch.Tensor): The logits for the fake data.
+        real_weight, fake_weight (Optional[torch.Tensor]): Finite nonnegative weights
+            that broadcast to the corresponding logits. Zero entries exclude scores.
 
     Returns:
         gan_loss (torch.Tensor): The GAN loss for the discriminator.
     """
     assert fake_logits.ndim == 2, f"fake_logits has shape {fake_logits.shape}"
     assert real_logits.ndim == 2, f"real_logits has shape {real_logits.shape}"
-    gan_loss = F.softplus(fake_logits).mean() + F.softplus(-real_logits).mean()
+    gan_loss = gan_loss_generator(-fake_logits, loss_weight=fake_weight) + gan_loss_generator(
+        real_logits, loss_weight=real_weight
+    )
 
     return gan_loss