Add sCM and MeanFlow distillation losses

.importlinter

@@ -113,9 +113,10 @@ exhaustive = True
 containers=
    physicsnemo.diffusion
 layers = 
+   distillation
    generate
    samplers | metrics
-   multi_diffusion | preconditioners | guidance
+   multi_diffusion | preconditioners | guidance | _fastgen
    noise_schedulers
    utils
    base

docs/api/diffusion/distillation.rst

@@ -0,0 +1,163 @@
+Distillation losses
+===================
+
+.. warning::
+
+   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.
+
+Use it in your training loop with initialized models and a noise scheduler.
+Your loop controls backpropagation, optimizer updates, and checkpoints.
+Use PhysicsNeMo samplers for inference.
+
+Install the backend
+-------------------
+
+Install the optional ``diffusion-fastgen`` extra from this checkout:
+
+.. code-block:: bash
+
+   pip install -e '.[diffusion-fastgen]'
+
+This extra installs the FastGen version required by the built-in losses.
+You do not need FastGen to write your own loss using ``BaseDistillationLoss``.
+
+Train a student
+---------------
+
+Give your student the :class:`~physicsnemo.diffusion.base.DiffusionModel`
+interface. For MeanFlow, it receives ``(t, DiffusionExtraTimes(r=r, batch_size=[B]))``
+and predicts average velocity between these times. A teacher supplies
+instantaneous velocity at ``t``. Set its ``prediction_type`` if it predicts
+clean data, noise, or score instead.
+
+.. code-block:: python
+
+   from physicsnemo.diffusion.distillation import DistillationLoss
+   from physicsnemo.diffusion.noise_schedulers import RectifiedFlowNoiseScheduler
+
+   loss_fn = DistillationLoss(
+       student,
+       RectifiedFlowNoiseScheduler(),
+       method="MeanFlow",
+       teacher=teacher,
+       model_options={"teacher": {"prediction_type": "flow"}},
+   )
+   optimizer.zero_grad()
+   losses = loss_fn(x0, condition=condition, step=step)
+   losses["total_loss"].backward()
+   optimizer.step()
+
+Omit ``teacher`` to learn from the conditional velocity of paired data and noise.
+Backpropagate through ``total_loss`` only. Other entries are for logging and
+may duplicate parts of the total. ``return_outputs=True`` also returns
+detached predictions and times without drawing extra samples.
+
+For sCM, select ``method="sCM"`` and return a prediction tensor paired with
+``DiffusionExtraOutputs(logvar=logvar, batch_size=[B])``. The prediction is
+always required. The ``logvar`` tensor has shape ``[B]`` and controls
+adaptive loss weighting. It does not describe diffusion posterior variance.
+The default student prediction type is ``"x0"``. A
+:class:`~physicsnemo.diffusion.preconditioners.TrigFlowPreconditioner` can provide
+this prediction from a normalized-velocity backbone.
+
+Control models and times
+------------------------
+
+``model_options`` accepts role-specific mappings for ``student`` and ``teacher``.
+Use ``prediction_type="x0"``, ``"epsilon"``, ``"score"``, or ``"flow"`` to
+describe the model output. ``model_kwargs`` forwards extra model controls.
+For example, a model may need ``{"return_logvar": True}`` to produce ``logvar``.
+
+MeanFlow students receive both times in a tuple. Other model calls use a
+single time tensor. The extra-time container supplies only the
+reserved ``r`` field. General custom-field propagation remains future work.
+
+To use a different noise schedule for the teacher, set
+``model_options["teacher"]["noise_scheduler"]``. The loss matches noise levels
+and converts predictions between schedules for you. The teacher's scheduler
+must provide ``snr_inv`` for this conversion.
+
+``method_options["guidance_scale"]`` enables classifier-free guidance. MeanFlow
+applies it at all diffusion times by default. Use
+``guidance_time_bounds=(lo, hi)`` to restrict it to an interval.
+
+By default, the loss uses the scheduler's ``sample_time`` distribution. Pass
+``t=...`` to choose diffusion times directly. For MeanFlow, pass
+``(t, DiffusionExtraTimes(r=r, batch_size=[B]))``; the loss uses this pair as given.
+You can also supply ``sample_t_fn(N, *, device, dtype)`` at call time. This
+callback returns the same time structure and applies only to that call.
+
+To request a bounded distribution, set
+``method_options={"sampling": {"time_dist_type": "uniform"}}``. Other choices
+are ``"lognormal"``, ``"logitnormal"``, and ``"log_t"``. Bounds follow live
+attributes of built-in schedulers. For custom schedulers, provide training
+limits through ``model_options["student"]["scheduler_options"]["time_bounds"]``
+when the selected method or sampling policy needs them. sCM also needs a differentiable
+``snr_inv`` implementation.
+
+Mask targets and weight residuals
+--------------------------------
+
+``WeightedDistillationLoss`` adds masks and element weights to the standard
+loss. Zero weights exclude entries from residuals, adaptive norms, and
+effective dimensions. Positive weights multiply the final squared residual
+once and leave the adaptive statistics unchanged. Mean reductions keep the
+full element count, matching ``WeightedMSEDSMLoss``. Samples with no positive
+weights contribute zero, including the sCM log-variance term.
+
+Use the same constructor as ``DistillationLoss`` and pass ``weight`` at call
+time. Weights must be finite, nonnegative, and broadcast to the data shape.
+
+.. code-block:: python
+
+   from physicsnemo.diffusion.distillation import WeightedDistillationLoss
+
+   loss_fn = WeightedDistillationLoss(student, scheduler, method="MeanFlow")
+   losses = loss_fn(x0, weight=valid_pixels, step=step)
+   losses["total_loss"].backward()
+
+Replace invalid input values before calling the model. A loss mask cannot
+prevent corrupted inputs from affecting predictions at other locations.
+
+Compilation and state
+---------------------
+
+To compile sCM or MeanFlow training, apply ``torch.compile`` to the enclosing
+loss calculation. Compiling only the model can fail when the loss computes
+an exact Jacobian-vector product (JVP). If you compile only the model, select
+``method_options={"use_jvp_finite_diff": True}`` to approximate the derivative
+with finite differences instead.
+
+With exact derivatives, MeanFlow updates stateful layers once per prediction
+and shares random draws between the prediction and its derivative. These calculations use
+the math attention backend because fused kernels may not support these
+derivatives. Check your combination of compilation, mixed precision, and
+distributed wrapping on your target hardware before a long training run.
+
+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.
+
+Custom implementations
+----------------------
+
+``BaseDistillationLoss`` defines the callable loss interface independently
+of the built-in objectives. A custom loss can be a function or callable class.
+Inheritance and FastGen are optional.
+
+Define ``__call__(x0, *, condition=None, **kwargs)`` and return a dictionary
+containing a differentiable ``total_loss``. Add named entries for logging or
+extra arguments for your method as needed. The training loop computes gradients
+through ``total_loss`` and updates the model.
+
+API reference
+-------------
+
+.. automodule:: physicsnemo.diffusion.distillation
+   :members:
+   :show-inheritance:

docs/api_diffusion.rst

@@ -13,6 +13,7 @@ PhysicsNeMo Diffusion
    api/diffusion/noise_schedulers.rst
    api/diffusion/preconditioners.rst
    api/diffusion/metrics.rst
+   api/diffusion/distillation.rst
    api/diffusion/samplers.rst
    api/diffusion/multi_diffusion.rst
-   api/diffusion/utils.rst
+   api/diffusion/utils.rst

physicsnemo/diffusion/_fastgen.py

@@ -0,0 +1,197 @@
+# 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.
+
+r"""Adapt PhysicsNeMo model calls and predictions for FastGen losses."""
+
+import importlib
+from collections.abc import Mapping
+from typing import Any, Literal
+
+import torch
+from jaxtyping import Float
+from tensordict import TensorDict
+from torch import Tensor
+
+from physicsnemo.core import Module
+from physicsnemo.core.version_check import check_version_spec
+
+from .base import (
+    DiffusionExtraOutputs,
+    DiffusionExtraTimes,
+)
+from .noise_schedulers._fastgen import _FastGenNoiseScheduler
+
+_AVAILABLE = check_version_spec("fastgen", "0.1.0", hard_fail=False)
+if _AVAILABLE:
+    _FastGenNetworkBase = importlib.import_module(
+        "fastgen.networks.network"
+    ).FastGenNetwork
+else:
+    _FastGenNetworkBase = Module
+
+
+class _FastGenNetwork(_FastGenNetworkBase):
+    r"""Adapt model time inputs and output fields to FastGen conventions.
+
+    Wrap an initialized PhysicsNeMo model, preserving its parameters and
+    distributed configuration. Save the original model in checkpoints and
+    reconstruct this adapter when resuming training.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        noise_scheduler: _FastGenNoiseScheduler,
+        *,
+        prediction_type: Literal["x0", "epsilon", "score", "flow"],
+        prediction_semantics: Literal["instantaneous", "average"] = "instantaneous",
+        model_kwargs: Mapping[str, Any] | None = None,
+    ) -> None:
+        if not _AVAILABLE:
+            raise ImportError(
+                "PhysicsNeMo distillation requires the pinned FastGen backend. "
+                "Install nvidia-physicsnemo[diffusion-fastgen]."
+            )
+        if not isinstance(model, torch.nn.Module):
+            raise TypeError(
+                "Training requires a parameter-owning module, not a bare callable."
+            )
+        if prediction_type not in ("x0", "epsilon", "score", "flow"):
+            raise ValueError(
+                f"Expected x0, epsilon, score, or flow, got {prediction_type!r}."
+            )
+        if prediction_semantics not in ("instantaneous", "average"):
+            raise ValueError(
+                f"Expected instantaneous or average semantics, got {prediction_semantics!r}."
+            )
+        if prediction_semantics == "average" and prediction_type != "flow":
+            raise ValueError("Average flow requires flow prediction.")
+        torch.nn.Module.__init__(self)
+        self.model = model
+        self.noise_scheduler = noise_scheduler
+        self.net_pred_type = (
+            "eps" if prediction_type in ("epsilon", "score") else prediction_type
+        )
+        self.schedule_type = "physicsnemo"
+        self._prediction_type = prediction_type
+        self._prediction_semantics = prediction_semantics
+        self._model_kwargs = dict(model_kwargs or {})
+        self._input_noise_scheduler = None
+
+    def set_noise_schedule(
+        self, schedule_type: str | None = None, **noise_schedule_kwargs: Any
+    ) -> None:
+        r"""Preserve the configured source across FastGen's reset lifecycle."""
+        if schedule_type is not None or noise_schedule_kwargs:
+            raise ValueError(
+                "Configure the PhysicsNeMo noise scheduler before creating this bridge."
+            )
+
+    def fully_shard(self, **kwargs: Any) -> None:
+        r"""Keep distributed wrapping under the caller's PhysicsNeMo runtime."""
+        raise NotImplementedError(
+            "Shard the PhysicsNeMo model before constructing the distillation loss."
+        )
+
+    def forward(
+        self,
+        x_t: Float[Tensor, " B *dims"],
+        t: Float[Tensor, " B"],
+        condition: Float[Tensor, " B *cond_dims"] | TensorDict | None = None,
+        r: Float[Tensor, " B"] | None = None,
+        return_features_early: bool = False,
+        feature_indices: set[int] | None = None,
+        return_logvar: bool = False,
+        fwd_pred_type: str | None = None,
+        **fwd_kwargs: Any,
+    ) -> (
+        Float[Tensor, " B *dims"]
+        | list[Float[Tensor, " B *feature_dims"]]
+        | tuple[Float[Tensor, " B *dims"], Float[Tensor, "B 1"]]
+        | tuple[Float[Tensor, " B *dims"], list[Float[Tensor, " B *feature_dims"]]]
+    ):
+        r"""Evaluate once, extract requested fields, and convert only predictions."""
+        target = self.net_pred_type if fwd_pred_type is None else fwd_pred_type
+        if target not in ("x0", "eps", "flow"):
+            raise ValueError(f"Expected x0, eps, or flow prediction, got {target!r}.")
+        if self._prediction_semantics == "average" and target != "flow":
+            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:
+            if r is not None:
+                raise ValueError(
+                    "Two-time predictions cannot use instantaneous coordinate mapping."
+                )
+            native_t = self.noise_scheduler.sqrt_snr_to_t(source.sqrt_snr(t.double()))
+            source_scale = torch.hypot(
+                source.alpha(t.double()), source.sigma(t.double())
+            )
+            target_scale = torch.hypot(
+                self.noise_scheduler.alpha(native_t),
+                self.noise_scheduler.sigma(native_t),
+            )
+            scale = target_scale / source_scale
+            if not torch.compiler.is_compiling() and not (
+                torch.isfinite(native_t).all() and torch.isfinite(scale).all()
+            ):
+                raise ValueError(
+                    "Teacher coordinate conversion needs finite native times and state scales."
+                )
+            scale = scale.reshape(-1, *([1] * (x_t.ndim - 1)))  # (B, 1, ...)
+            x_t, t = (x_t.double() * scale).to(x_t.dtype), native_t.to(t.dtype)
+        time = (
+            t if r is None else (t, DiffusionExtraTimes(r=r, batch_size=[x_t.shape[0]]))
+        )
+        output = self.model(
+            x_t, time, condition=condition, **(self._model_kwargs | fwd_kwargs)
+        )
+        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.")
+        if self._prediction_type == "score":
+            prediction = self.noise_scheduler._scheduler.score_to_epsilon(prediction, t)
+        prediction = self.noise_scheduler.convert_model_output(
+            x_t,
+            prediction,
+            t,
+            self.net_pred_type,
+            "x0" if source is not None else target,
+        )
+        if source is not None:
+            prediction = source.convert_model_output(
+                input_x, prediction, input_t, "x0", target
+            )
+        if return_logvar:
+            if not isinstance(extras, DiffusionExtraOutputs):
+                raise ValueError(
+                    "This method requires (prediction, DiffusionExtraOutputs(logvar=...))."
+                )
+            logvar = extras.logvar
+            if logvar is None:
+                raise ValueError("This method requires a logvar tensor; got None.")
+            if logvar.shape != (x_t.shape[0],):
+                raise ValueError(
+                    f"Expected logvar shape {(x_t.shape[0],)}, got {logvar.shape}."
+                )
+            return prediction, logvar[:, None]  # (B, 1)
+        return prediction