Add DMD2 and f-distill training losses

docs/api/diffusion/distillation.rst

@@ -6,9 +6,9 @@ Distillation losses
    This API is alpha. Imports emit ``ExperimentalFeatureWarning`` and future
    releases may introduce changes that break backward compatibility.
 
-``DistillationLoss`` computes sCM or MeanFlow objectives for PhysicsNeMo
-models. Each call constructs training targets and returns a differentiable
-``total_loss`` with detached values for logging.
+``DistillationLoss`` computes sCM, MeanFlow, DMD2, or f-distill losses for
+PhysicsNeMo models. Each call constructs training targets and returns a
+differentiable ``total_loss`` with detached values for logging.
 
 Use it in your training loop with initialized models and a noise scheduler.
 Your loop controls backpropagation, optimizer updates, and checkpoints.
@@ -143,6 +143,128 @@ Save your PhysicsNeMo models and optimizers as usual. Recreate the loss with
 the same noise schedules and settings when resuming. sCM and MeanFlow have no
 extra running statistics, so their ``state_dict`` results are empty.
 
+Alternate distribution-matching objectives
+-----------------------------------------
+
+DMD2 and f-distill match a student's generated distribution to a fixed
+diffusion teacher. Unlike an ordinary denoising loss that trains a model on
+clean data and noise, these objectives also train a fake-score model on
+student-generated samples.
+
+Initialize the models and pass the fake-score model in ``auxiliary_models``.
+Choose ``phase="student"`` or ``phase="auxiliary"`` on each call and step
+the matching optimizer. Pass the training step separately as ``step``.
+
+.. code-block:: python
+
+   loss_fn = DistillationLoss(
+       student, scheduler, method="DMD2", teacher=teacher,
+       auxiliary_models={"fake_score": fake_score},
+   )
+   for phase, optimizer in (("auxiliary", fake_optimizer),
+                            ("student", student_optimizer)):
+       optimizer.zero_grad()
+       loss_fn(x0, phase=phase, step=step)["total_loss"].backward()
+       optimizer.step()
+
+The student phase trains only the student. The auxiliary phase trains the
+fake-score model and any discriminator. The teacher stays fixed in both phases.
+Use separate parameters for each model. The loss preserves your choices of
+training mode and frozen parameters.
+
+DMD2 trains the student using teacher and fake-score predictions, with an
+optional adversarial loss. f-distill also weights this score
+matching by a selected f-divergence and requires a discriminator.
+
+Supply ``discriminator`` in ``auxiliary_models`` to enable adversarial training.
+The generator GAN coefficient defaults to ``0.001``. Set
+``gan_loss_weight_gen`` to change it. f-distill requires a positive coefficient.
+The discriminator accepts a TensorDict of named teacher features and returns
+scores with shape :math:`(B, L)`, where :math:`L` counts scores per sample.
+Supply a ``feature_fn(prediction, extras)`` callback to select those features
+from your teacher output. The teacher must still return its prediction.
+Feature fields are model-specific, not reserved diffusion outputs. The callback
+interprets them. For example, if your teacher adds a
+``bottleneck`` field to a subclass of ``DiffusionExtraOutputs``:
+
+.. code-block:: python
+
+   from tensordict import TensorDict
+
+   def teacher_features(prediction, extras):
+       return TensorDict({"bottleneck": extras.bottleneck}, batch_size=extras.batch_size)
+
+   loss_fn = DistillationLoss(
+       student, scheduler, method="f-distill", teacher=teacher,
+       auxiliary_models={"fake_score": fake_score,
+                         "discriminator": discriminator},
+       model_options={"teacher": {"feature_names": ("bottleneck",),
+                                  "feature_fn": teacher_features}},
+       method_options={"f_distill": {"f_div": "js"}},
+   )
+
+Supported divergence names are ``rkl``, ``kl``, ``js``, ``sf``, ``neyman``,
+``sh``, and ``jf``. f-distill tracks running ratio statistics in the loss object.
+Save and restore ``loss_fn.state_dict()`` alongside model and optimizer state.
+Pass ``update_state=False`` to compute a loss without changing these statistics.
+
+``model_options["fake_score"]`` configures its prediction type or scheduler.
+The loss converts predictions when these models use different noise schedules.
+Use ``fake_score_pred_type`` to select the fake-score training target and
+``fake_score_sampling`` to choose its training-time distribution.
+
+One-step training draws student inputs at the scheduler's upper training time.
+Custom schedulers without that bound need explicit ``input_student`` and
+``t_student`` on the loss call, or configured training bounds. Multi-step training
+uses ``student_sample_steps`` and an explicit ``training_times`` grid under the
+student's ``scheduler_options``. Specify training times independently of your
+inference grid.
+Pass ``noise`` to control the score perturbation independently.
+
+Adversarial masking
+~~~~~~~~~~~~~~~~~~~
+
+Adversarial masks exclude corrupted inputs during feature extraction and
+discriminator pooling. The loss needs a callback that maps the input mask to
+valid discriminator scores.
+
+For masked adversarial training, supply a callback with this signature:
+
+.. code-block:: python
+
+   def adversarial_fn(x_t, t, condition, weight):
+       ...
+       return logits, logit_weight
+
+Return ``logits`` with shape :math:`(B, L)` and finite,
+nonnegative ``logit_weight`` values that broadcast to that shape.
+
+Use the same teacher and discriminator you passed to the loss, and keep
+gradients through their inputs for student training. You may need a masked
+backbone or feature extractor. Masking only the final scores cannot remove
+corruption that has already spread through features or pooled statistics.
+
+The loss applies the returned weights to generator, discriminator, and approximate
+R1 residuals. Zero weights also exclude scores and empty samples from f-distill
+ratio statistics. Positive residual weights remain linear. The callback receives
+both real and generated perturbations, including the R1 query when enabled.
+
+Distributed statistics and compilation
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+f-distill reduces ratio statistics over the default process group for ordinary
+data parallelism. For a domain-parallel scheduler, explicitly supply the data
+group through ``method_options["f_distill"]["statistics_group"]``. It must
+group corresponding spatial partitions across independent data replicas, not
+different partitions of the same sample. Each rank participates even when its
+local batch has no valid entries. The callback must produce consistent global scores
+across spatial partitions when the architecture needs spatial aggregation.
+
+For DMD2, you can compile the individual models with ``torch.compile``.
+Keep phase selection and f-distill statistics updates outside compilation.
+Check combined domain parallelism and FSDP on your target hardware before
+a long training run.
+
 Custom implementations
 ----------------------
 

physicsnemo/diffusion/_fastgen.py

@@ -17,7 +17,7 @@
 r"""Adapt PhysicsNeMo model calls and predictions for FastGen losses."""
 
 import importlib
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from typing import Any, Literal
 
 import torch
@@ -59,6 +59,11 @@ def __init__(
         prediction_type: Literal["x0", "epsilon", "score", "flow"],
         prediction_semantics: Literal["instantaneous", "average"] = "instantaneous",
         model_kwargs: Mapping[str, Any] | None = None,
+        feature_names: tuple[str, ...] = (),
+        feature_fn: Callable[
+            [Float[Tensor, " B *dims"], DiffusionExtraOutputs | None], TensorDict
+        ]
+        | None = None,
     ) -> None:
         if not _AVAILABLE:
             raise ImportError(
@@ -90,6 +95,8 @@ def __init__(
         self._prediction_semantics = prediction_semantics
         self._model_kwargs = dict(model_kwargs or {})
         self._input_noise_scheduler = None
+        self._feature_names = tuple(feature_names)
+        self._feature_fn = feature_fn
 
     def set_noise_schedule(
         self, schedule_type: str | None = None, **noise_schedule_kwargs: Any
@@ -131,10 +138,6 @@ def forward(
             raise ValueError(
                 "An average flow cannot use instantaneous prediction conversions."
             )
-        if return_features_early or feature_indices:
-            raise NotImplementedError(
-                "This adapter does not define feature extraction."
-            )
         input_x, input_t = x_t, t
         source = self._input_noise_scheduler
         if source is not None:
@@ -168,6 +171,21 @@ def forward(
         prediction, extras = output if isinstance(output, tuple) else (output, None)
         if not isinstance(prediction, Tensor):
             raise TypeError("A diffusion model must return a prediction tensor.")
+        features = None
+        if return_features_early or feature_indices:
+            if self._feature_fn is None:
+                raise ValueError(
+                    "Adversarial feature requests require an explicit feature_fn."
+                )
+            named_features = self._feature_fn(prediction, extras)
+            indices = (
+                sorted(feature_indices)
+                if feature_indices
+                else range(len(self._feature_names))
+            )
+            features = [named_features[self._feature_names[index]] for index in indices]
+            if return_features_early:
+                return features
         if self._prediction_type == "score":
             prediction = self.noise_scheduler._scheduler.score_to_epsilon(prediction, t)
         prediction = self.noise_scheduler.convert_model_output(
@@ -194,4 +212,4 @@ def forward(
                     f"Expected logvar shape {(x_t.shape[0],)}, got {logvar.shape}."
                 )
             return prediction, logvar[:, None]  # (B, 1)
-        return prediction
+        return (prediction, features) if features is not None else prediction

physicsnemo/diffusion/distillation/losses.py

@@ -36,20 +36,22 @@
 _METHODS = {
     "sCM": ("scm._fastgen", "_SCMLoss"),
     "MeanFlow": ("meanflow._fastgen", "_MeanFlowLoss"),
+    "DMD2": ("dmd2._fastgen", "_DMD2Loss"),
+    "f-distill": ("f_distill._fastgen", "_FDistillLoss"),
 }
 
 
 class DistillationLoss:
-    r"""Compute sCM or MeanFlow objectives for a diffusion student.
+    r"""Compute sCM, MeanFlow, DMD2, or f-distill losses for a diffusion student.
 
-    The selected method defines the training loss and constructs its
-    targets, with or without a teacher. The loss returns ``total_loss`` for
-    backpropagation and detached values for logging. Your training loop
-    controls optimization.
+    The selected method defines the training loss and constructs its targets.
+    Each call returns ``total_loss`` for backpropagation and detached values
+    for logging. Your training loop controls optimization.
 
     Initialize and wrap your models before constructing the loss. Call it on
-    each batch of clean data, backpropagate through ``losses["total_loss"]``,
-    and step your optimizer. See the distillation guide for compilation options.
+    each batch, backpropagate through ``losses["total_loss"]``, and step the
+    optimizer for the selected phase. See the distillation guide for
+    compilation options.
 
     Parameters
     ----------
@@ -59,18 +61,27 @@ class DistillationLoss:
     noise_scheduler : LinearGaussianNoiseScheduler or DomainParallelNoiseScheduler
         Noise schedule for the student. Its ``sample_time`` method supplies
         training times unless you pass times to the loss call.
-    method : {"sCM", "MeanFlow"}
+    method : {"sCM", "MeanFlow", "DMD2", "f-distill"}
         Training loss. MeanFlow students receive ``(t, DiffusionExtraTimes)``
         and predict average velocity over the interval from ``t`` to ``r``.
     teacher : torch.nn.Module, optional
-        Model providing training targets. Omit it to train from clean data
+        Model providing training targets. DMD2 and f-distill require a
+        teacher. For sCM and MeanFlow, omit it to train from clean data
         and noise without a pretrained teacher.
+    auxiliary_models : mapping, optional
+        Models trained during the auxiliary phase. Supply ``fake_score`` for
+        DMD2 and f-distill, and ``discriminator`` for adversarial training.
+        Initialize these models before passing them.
     method_options : mapping, optional
         Method-specific loss settings. Unknown settings raise an error.
     model_options : mapping, optional
-        Settings for ``student`` or ``teacher``. Set ``prediction_type`` to
+        Settings for ``student``, ``teacher``, or ``fake_score``. Set ``prediction_type`` to
         describe what each model predicts. ``model_kwargs`` forwards model-specific options,
         and ``noise_scheduler`` selects a different schedule for that model.
+        Adversarial teachers also need ``feature_names`` and a
+        ``feature_fn(prediction, extras)`` callback that returns named features
+        as a ``TensorDict``. The teacher always returns its prediction; the
+        callback selects model-specific extras for the discriminator.
 
     Examples
     --------
@@ -100,6 +111,7 @@ def __init__(
         *,
         method: str,
         teacher: torch.nn.Module | None = None,
+        auxiliary_models: Mapping[str, torch.nn.Module] | None = None,
         method_options: Mapping[str, Any] | None = None,
         model_options: Mapping[str, Mapping[str, Any]] | None = None,
     ) -> None:
@@ -111,12 +123,18 @@ def __init__(
         implementation = getattr(
             importlib.import_module(f"{__package__}.methods.{module}"), name
         )
+        auxiliary_options = {}
+        if method in ("DMD2", "f-distill"):
+            auxiliary_options["auxiliary_models"] = dict(auxiliary_models or {})
+        elif auxiliary_models:
+            raise ValueError(f"{method} does not use auxiliary models.")
         self._implementation = implementation(
             student,
             noise_scheduler,
             teacher=teacher,
             method_options=dict(method_options or {}),
             model_options=dict(model_options or {}),
+            **auxiliary_options,
         )
 
     def __call__(
@@ -152,7 +170,9 @@ def __call__(
             Conditioning forwarded to the supplied models.
         phase : str, optional
             Training phase. sCM and MeanFlow accept only ``"student"``
-            or ``None``.
+            or ``None``. For DMD2 and f-distill, choose ``"student"``
+            to train the student or ``"auxiliary"`` to train the fake-score
+            model and optional discriminator. Step the matching optimizers.
         step : int, optional, default=0
             Training step, used to adjust loss warmup.
         return_outputs : bool, optional, default=False

physicsnemo/diffusion/distillation/methods/dmd2/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.