Add sCM and MeanFlow distillation losses

.importlinter

@@ -113,11 +113,13 @@ exhaustive = True
 containers=
    physicsnemo.diffusion
 layers = 
+   distillation
    generate
    samplers | metrics
-   multi_diffusion | preconditioners | guidance
+   multi_diffusion | preconditioners | guidance | _fastgen
    noise_schedulers
-   utils | base
+   utils
+   base
 
 [importlinter:contract:physicsnemo-sym]
 name = Control Internal Dependencies in PhysicsNeMo sym

CHANGELOG.md

@@ -15,6 +15,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Adds an experimental, independently implemented SMART operator with sampled
   geometry encoding, reusable latent contexts, and chunked field decoding.
 
+- Adds `DiffusionExtraTimes` and `DiffusionExtraOutputs` to
+  `physicsnemo.diffusion` for utilities that carry named times or auxiliary
+  outputs. Compatible model interfaces and preconditioners can accept
+  `(t, extra_times)`, where `extra_times` is a `DiffusionExtraTimes` container,
+  and return `(prediction, extra_outputs)`, where `extra_outputs` is a
+  `DiffusionExtraOutputs` container. The API reserves the named field
+  `extra_times.r` for a destination diffusion time, as in MeanFlow. It reserves
+  the named field `extra_outputs.logvar` for learned loss weighting, as in sCM
+  and EDM2. Existing single-time, single-prediction calls remain
+  unchanged. Existing DSM and flow-matching losses continue to consume
+  single tensors.
+- Adds `TrigFlowNoiseScheduler` and `TrigFlowPreconditioner` to the diffusion
+  noise schedulers and preconditioners in `physicsnemo.diffusion`. Pair them
+  with the same `sigma_data` to train and sample models along a trigonometric
+  diffusion path, with the backbone predicting normalized velocity. Built-in
+  linear Gaussian schedulers also provide differentiable `snr_inv()` methods
+  to convert signal-to-noise amplitude ratios back to diffusion times. Use
+  these inverses to match noise levels across schedules. The amplitude ratio
+  is `alpha(t) / sigma(t)`. The power signal-to-noise ratio is its square.
+  Custom schedulers can supply their own inverse. These additions are opt-in.
+  Existing schedulers, preconditioners, and training workflows work as before.
 - Adds standalone FLARE++ attention and model APIs, with input-conditioned
   dynamic routing, plus a ``GALE_FPP`` backend for using the same mixer inside
   GeoTransolver.

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/introduction.rst

@@ -290,6 +290,16 @@ API Reference
     :members:
     :exclude-members: __init__
 
+:code:`DiffusionExtraTimes`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: physicsnemo.diffusion.DiffusionExtraTimes
+
+:code:`DiffusionExtraOutputs`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: physicsnemo.diffusion.DiffusionExtraOutputs
+
 :code:`Predictor`
 ~~~~~~~~~~~~~~~~~
 

docs/api/diffusion/models.rst

@@ -80,7 +80,68 @@ expects ``(x, t, condition)``.  A short adapter class solves this:
         def forward(self, x, t, condition=None):
             return self.net(x, noise_labels=t, class_labels=condition)
 
+Most diffusion and flow-matching models take one time tensor and return one
+prediction tensor.
+A compatible :class:`~physicsnemo.diffusion.DiffusionModel` can instead
+accept ``(t, extra_times)`` with
+:class:`~physicsnemo.diffusion.DiffusionExtraTimes` and return
+``(prediction, extra_outputs)`` with
+:class:`~physicsnemo.diffusion.DiffusionExtraOutputs`. For example, MeanFlow
+uses a destination time ``r``, while sCM and EDM2 can use ``logvar`` for loss
+weighting. The primary time and prediction remain required tensors.
+
+The adapter below combines the destination time with the primary time for a
+DiT backbone. A small
+:class:`~physicsnemo.models.mlp.FullyConnected` head predicts a log-scale
+from the primary time. This example does not define a canonical MeanFlow
+architecture. A compatible loss decides how to use ``logvar`` during
+training. Inference code that needs only a prediction can ignore it.
 
+.. code-block:: python
+
+    import torch
+    from physicsnemo.core import Module
+    from physicsnemo.diffusion import DiffusionExtraOutputs, DiffusionExtraTimes
+    from physicsnemo.models.dit import DiT
+    from physicsnemo.models.mlp import FullyConnected
+
+    class ExtendedDiTAdapter(Module):
+        def __init__(self):
+            super().__init__()
+            self.backbone = DiT(
+                input_size=8, in_channels=2, out_channels=2,
+                patch_size=4, hidden_size=32, depth=1, num_heads=4,
+            )
+            self.r_embed = FullyConnected(
+                in_features=1, layer_size=8, out_features=1, num_layers=1,
+            )
+            self.logvar_head = FullyConnected(
+                in_features=1, layer_size=8, out_features=1, num_layers=1,
+            )
+
+        def forward(self, x, t, condition=None):
+            t, extra_times = t
+            # DiT takes one time tensor, so combine the two times in the adapter.
+            r_embed = self.r_embed(extra_times.r[:, None]).squeeze(-1)
+            prediction = self.backbone(x, t + r_embed, condition)
+            logvar = self.logvar_head(t[:, None]).squeeze(-1)
+            return prediction, DiffusionExtraOutputs(
+                logvar=logvar, batch_size=[x.shape[0]],
+            )
+
+    model = ExtendedDiTAdapter()
+    x = torch.randn(2, 2, 8, 8)
+    t = torch.ones(2)
+    times = (t, DiffusionExtraTimes(r=torch.zeros(2), batch_size=[2]))
+    prediction, extra_outputs = model(x, times)
+    assert prediction.shape == x.shape
+    assert extra_outputs.logvar.shape == t.shape
+
+This adapter illustrates one way to embed the reserved ``r`` time and compute
+the reserved ``logvar`` output. Other implementations can use different
+architectures. Its ``forward`` method implements
+:class:`~physicsnemo.diffusion.DiffusionModel` and works with utilities that
+support both extra times and extra outputs.
 
 .. _backbone_external_models:
 
@@ -100,4 +161,3 @@ You can bring in models from external libraries, such as:
 * by using :meth:`~physicsnemo.core.module.Module.from_torch` to convert a
 ``torch.nn.Module`` subclass into a :class:`~physicsnemo.core.module.Module`
 with PhysicsNeMo serialization support
-

docs/api/diffusion/noise_schedulers.rst

@@ -23,6 +23,46 @@ the coupling between noise schedule and reverse process in one place and avoids
 the need for schedule-specific sampler code.
 
 
+Amplitude ratio and inverse
+---------------------------
+
+The amplitude ratio :math:`\alpha(t)/\sigma(t)` measures signal strength
+relative to noise strength. The method ``scheduler.snr(t)`` returns this
+ratio. The usual power signal-to-noise ratio is its square.
+
+Use ``scheduler.snr_inv(q)`` to recover a diffusion time from an amplitude
+ratio or to match noise levels between schedules. EDM, VE, VP, rectified-flow,
+continuous iDDPM, and TrigFlow provide differentiable inverses. Use finite
+positive ratios on the schedule's invertible domain. Inverses do not clip
+to training bounds. Add ``snr_inv`` for custom schedulers that need this
+conversion, and keep it differentiable for training losses such as sCM.
+
+.. code-block:: python
+
+    import torch
+    from physicsnemo.diffusion.noise_schedulers import TrigFlowNoiseScheduler
+
+    scheduler = TrigFlowNoiseScheduler(sigma_data=0.5)
+    t = torch.tensor([0.2, 0.7])
+    recovered = scheduler.snr_inv(scheduler.snr(t))
+    torch.testing.assert_close(recovered, t)
+
+TrigFlow interpolates between data and Gaussian noise with trigonometric
+coefficients. Its diffusion time is an angle in radians, whereas EDM uses
+the noise standard deviation.
+
+Pair ``TrigFlowNoiseScheduler`` with ``TrigFlowPreconditioner`` and the same
+``sigma_data``. Training samples diffusion times uniformly. Inference uses
+an evenly spaced grid ending at zero. The denoising weight
+:math:`1/(\sigma_d\sin(t))^2` gives unit weight to the normalized-velocity
+residual with this preconditioner. Keep training diffusion times above zero
+for denoising losses and below :math:`\pi/2` for reverse-process sampling.
+
+.. autoclass:: TrigFlowNoiseScheduler
+    :show-inheritance:
+    :members:
+
+
 Role in Training
 ----------------
 

docs/api/diffusion/preconditioners.rst

@@ -140,6 +140,14 @@ API Reference
     :members:
     :exclude-members: forward
 
+:code:`TrigFlowPreconditioner`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: physicsnemo.diffusion.preconditioners.TrigFlowPreconditioner
+    :show-inheritance:
+    :members:
+    :exclude-members: forward
+
 :code:`VEPreconditioner`
 ~~~~~~~~~~~~~~~~~~~~~~~~
 

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/__init__.py

@@ -14,4 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .base import Denoiser, DiffusionModel, Predictor, PredictorType  # noqa: F401
+from .base import (  # noqa: F401
+    Denoiser,
+    DiffusionExtraOutputs,
+    DiffusionExtraTimes,
+    DiffusionModel,
+    Predictor,
+    PredictorType,
+)