torchdeltaflow 0.2.2__tar.gz → 0.2.3.dev23__tar.gz

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/.github/workflows/publish.yml

@@ -1,11 +1,18 @@
 name: publish
 
-# Publishes torchdeltaflow to PyPI whenever a version tag (e.g. v0.2.2) is
-# pushed. The version is derived from the tag by setuptools_scm, so the tag is
-# the single source of truth. To cut a release, run:
+# Publishes torchdeltaflow to PyPI in two cases:
 #
-#     git tag v0.2.2
-#     git push origin v0.2.2
+#   1. A version tag (e.g. v0.2.2) is pushed -> a stable release. The version is
+#      derived from the tag by setuptools_scm, so the tag is the single source
+#      of truth. To cut a release, run:
+#
+#          git tag v0.2.2
+#          git push origin v0.2.2
+#
+#   2. A pull request is merged into main -> a dev pre-release. setuptools_scm
+#      derives a clean PEP 440 version from the commit distance to the last tag
+#      (e.g. "0.2.3.dev4"), which is unique per merge. "skip-existing" below
+#      guards against the rare case where the same version was already uploaded.
 #
 # Authentication uses a PyPI API token stored as the repository secret
 # PYPI_API_TOKEN. Create the token at https://pypi.org/manage/account/token/
@@ -16,12 +23,18 @@ on:
   push:
     tags:
       - "v*"
+  pull_request:
+    branches: [main]
+    types: [closed]
 
 permissions:
   contents: read
 
 jobs:
   build:
+    # Run for pushed version tags, or when a PR is actually merged (not just
+    # closed). Plain "closed" without merge must not publish.
+    if: github.event_name == 'push' || github.event.pull_request.merged == true
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -60,3 +73,6 @@ jobs:
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           password: ${{ secrets.PYPI_API_TOKEN }}
+          # Merge builds may reproduce an existing dev version if no new commits
+          # changed the distance to the last tag. Skip rather than fail.
+          skip-existing: true

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/CHANGELOG.md

@@ -7,6 +7,26 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Added
 
+- `deltaflow.models.DiT`: a class-conditional Diffusion Transformer velocity
+  field with **adaLN-Zero** conditioning (Peebles & Xie, 2023). Ships a reusable
+  `TimestepEmbedding`, a `LabelEmbedding` with a learned null token for
+  classifier-free guidance, zero-initialized modulation gates (identity at init,
+  zero initial velocity), `cond["y"]` support through the existing
+  `BaseVelocityField` API, and a `forward_with_cfg` guided-sampling helper. See
+  the new `examples/20-training/03-conditional-dit` walkthrough.
+
+### Changed
+
+- `scipy` is now a core dependency (previously gated behind the `ot`
+  extra), so `OTInterpolant`/`OTCoupling` use the exact Hungarian-algorithm
+  mini-batch OT assignment by default. The `ot` extra is kept as a
+  backwards-compatible no-op, and the greedy nearest-neighbour matcher
+  remains as a defensive fallback if `scipy` is ever missing.
+
+## [0.2.2]
+
+### Added
+
 - API-reference docstrings now include the underlying mathematics in
   LaTeX (rendered via MathJax/arithmatex): probability paths and target
   velocities for the interpolants (`Linear`, `VariancePreserving`, `OT`,

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchdeltaflow
-Version: 0.2.2
+Version: 0.2.3.dev23
 Summary: Flow matching, mini-batch optimal transport coupling, and posterior sampling for inverse problems in PyTorch
 Author-email: Phrugsa Limbunlom <phrugsa.lim@gmail.com>
 Maintainer-email: Phrugsa Limbunlom <phrugsa.lim@gmail.com>
@@ -53,6 +53,7 @@ Requires-Dist: torch>=2.0
 Requires-Dist: numpy
 Requires-Dist: einops
 Requires-Dist: tqdm
+Requires-Dist: scipy>=1.10
 Provides-Extra: images
 Requires-Dist: pillow>=9.0; extra == "images"
 Provides-Extra: ot
@@ -85,6 +86,9 @@ Dynamic: license-file
     <a href="https://github.com/phrugsa-limbunlom/deltaflow/blob/main/LICENSE" target="_blank" title="License">
         <img alt="License" src="https://img.shields.io/github/license/phrugsa-limbunlom/deltaflow?style=flat-square&color=3f9e73">
     </a>
+    <a href="https://pepy.tech/project/torchdeltaflow" target="_blank" title="Downloads">
+        <img alt="Downloads" src="https://static.pepy.tech/badge/torchdeltaflow?style=flat-square">
+    </a>
     <a href="https://github.com/phrugsa-limbunlom/deltaflow" target="_blank" title="GitHub Repo Stars">
         <img alt="GitHub Stars" src="https://img.shields.io/github/stars/phrugsa-limbunlom/deltaflow?style=social">
     </a>
@@ -97,9 +101,6 @@ Dynamic: license-file
     <a href="https://github.com/phrugsa-limbunlom/deltaflow/actions/workflows/docs.yml" target="_blank" title="Documentation">
         <img alt="Docs" src="https://img.shields.io/github/actions/workflow/status/phrugsa-limbunlom/deltaflow/docs.yml?branch=main&style=flat-square&label=docs&color=0f9bab">
     </a>
-    <a href="https://pepy.tech/project/torchdeltaflow" target="_blank" title="Downloads">
-        <img alt="Downloads" src="https://static.pepy.tech/badge/torchdeltaflow?style=flat-square">
-    </a>
     <a href="https://pypi.org/project/torchdeltaflow/" target="_blank" title="Python Versions">
         <img alt="Python Versions" src="https://img.shields.io/pypi/pyversions/torchdeltaflow?style=flat-square&color=e0b13c">
     </a>

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/deltaflow/__init__.py

@@ -8,10 +8,13 @@ Module map:
   subclasses (`BaseVelocityField`, `BaseInterpolant`,
   `BaseSolver`, `BaseLoss`).
 - `deltaflow.interpolants`, probability paths (linear, mini-batch OT,
-  variance-preserving).
+  variance-preserving, Schrödinger bridge, and the Equilibrium Matching
+  energy-compatible target).
 - `deltaflow.losses`, conditional flow matching, plus the
   optional delta-alignment loss for guidance-representation pretraining.
-- `deltaflow.solvers`, Euler, Heun, and the
+- `deltaflow.solvers`, Euler, Heun, the `EquilibriumSolver` which samples an
+  Equilibrium Matching field by gradient descent on its implicit energy
+  landscape, and the
   `PosteriorSolver` which *wraps* a base solver
   and injects the measurement-likelihood gradient per step (FlowDPS /
   Flower style).

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/deltaflow/core/__init__.py

@@ -5,12 +5,18 @@ subclasses one of these bases, so new variants are drop-in and not
 rewrites of the surrounding machinery.
 """
 
+from .base_coupling import BaseCoupling
+from .base_equilibrium_field import BaseEquilibriumField
+from .base_equilibrium_interpolant import BaseEquilibriumInterpolant
 from .base_interpolant import BaseInterpolant
 from .base_loss import BaseLoss
 from .base_solver import BaseSolver
 from .base_velocity_field import BaseVelocityField
 
 __all__ = [
+    "BaseCoupling",
+    "BaseEquilibriumField",
+    "BaseEquilibriumInterpolant",
     "BaseInterpolant",
     "BaseLoss",
     "BaseSolver",

torchdeltaflow-0.2.3.dev23/deltaflow/core/base_coupling.py

@@ -0,0 +1,22 @@
+"""Base class for train-time coupling strategies between ``x0`` and ``x1``.
+
+Coupling is the choice of which ``(x0, x1)`` pairs the flow-matching
+regression is computed on, kept deliberately separate from the choice of
+probability path (`deltaflow.interpolants`). Concrete strategies live in
+`deltaflow.trainer.coupling`, this base sits in ``core`` alongside the other
+drop-in abstractions so a new coupling is a subclass rather than a rewrite of
+the surrounding training loop.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Tuple
+
+import torch
+
+
+class BaseCoupling(ABC):
+    """Given a batch ``x1`` of data samples, return a paired ``(x0, x1)``."""
+
+    @abstractmethod
+    def sample_pair(self, x1: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        raise NotImplementedError

torchdeltaflow-0.2.3.dev23/deltaflow/core/base_equilibrium_field.py

@@ -0,0 +1,31 @@
+"""Base class for time-invariant Equilibrium Matching fields ``f(x)``."""
+
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+
+
+class BaseEquilibriumField(nn.Module, ABC):
+    r"""Base class for a time-invariant Equilibrium Matching field \(f(x)\).
+
+    Unlike a flow-matching velocity field \(v_\theta(x, t)\), an EqM field
+    carries no time argument. It approximates the equilibrium gradient of an
+    implicit energy, \(f_\theta(x) \approx -\nabla E(x)\), pointing from noise
+    toward data. The interpolation coefficient \(\gamma\) is implicit and (per
+    the EqM paper) never seen by the model, so ``forward`` takes only ``x`` and
+    any extra conditioning, never a time or \(\gamma\).
+
+    Subclasses implement `forward` and return a tensor with the same
+    shape as ``x``. Additional conditioning is passed as keyword arguments and
+    forwarded unchanged by `EquilibriumMatchingLoss` and
+    `EquilibriumSolver`.
+
+    References:
+        Wang and Du, "Equilibrium Matching: Generative Modeling with Implicit
+        Energy-Based Models" (2025), https://arxiv.org/abs/2510.02300.
+    """
+
+    @abstractmethod
+    def forward(self, x: torch.Tensor, **cond) -> torch.Tensor:
+        raise NotImplementedError

torchdeltaflow-0.2.3.dev23/deltaflow/core/base_equilibrium_interpolant.py

@@ -0,0 +1,41 @@
+"""Base class for Equilibrium Matching interpolants (parametrised by gamma)."""
+
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple
+
+import torch
+
+
+class BaseEquilibriumInterpolant(ABC):
+    r"""Base class for an energy-compatible path indexed by a coefficient gamma.
+
+    This base is deliberately separate from `BaseInterpolant`. A flow-matching
+    interpolant is parametrised by a dynamical time ``t`` that a sampler
+    integrates over. An Equilibrium Matching interpolant is parametrised by an
+    interpolation coefficient \(\gamma \in [0, 1]\), a noise level rather than a
+    time. The learned field is trained to be time-invariant, so \(\gamma\) only
+    indexes where along the noise-to-data path a training point sits, it is
+    never integrated and (per the EqM paper) is not seen by the model.
+
+    Convention: \(\gamma = 0\) is noise (\(x_\gamma = x_0\)) and \(\gamma = 1\)
+    is data (\(x_\gamma = x_1\)).
+
+    References:
+        Wang and Du, "Equilibrium Matching: Generative Modeling with Implicit
+        Energy-Based Models" (2025), https://arxiv.org/abs/2510.02300.
+    """
+
+    @abstractmethod
+    def interpolate(
+        self, x1: torch.Tensor, gamma: torch.Tensor, x0: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        r"""Return ``(x_gamma, target)`` for the given data and coefficient.
+
+        Args:
+            x1: data sample (the \(\gamma = 1\) endpoint).
+            gamma: interpolation coefficient (noise level) in \([0, 1]\), not a
+                dynamical time.
+            x0: optional noise sample (the \(\gamma = 0\) endpoint). Drawn from
+                a standard normal when omitted.
+        """
+        raise NotImplementedError

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/deltaflow/interpolants/__init__.py

@@ -1,13 +1,17 @@
 """Probability paths connecting noise ``x0`` to data ``x1``."""
 
+from ..core.base_equilibrium_interpolant import BaseEquilibriumInterpolant
 from ..core.base_interpolant import BaseInterpolant
+from .equilibrium import EquilibriumInterpolant
 from .linear import LinearInterpolant
 from .ot import OTInterpolant
 from .schrodinger_bridge import SchrodingerBridgeInterpolant
 from .variance_preserving import VariancePreservingInterpolant
 
 __all__ = [
+    "BaseEquilibriumInterpolant",
     "BaseInterpolant",
+    "EquilibriumInterpolant",
     "LinearInterpolant",
     "OTInterpolant",
     "SchrodingerBridgeInterpolant",

torchdeltaflow-0.2.3.dev23/deltaflow/interpolants/equilibrium.py

@@ -0,0 +1,121 @@
+"""Equilibrium Matching probability path (energy-compatible target)."""
+
+from typing import Optional, Tuple
+
+import torch
+
+from ..core.base_equilibrium_interpolant import BaseEquilibriumInterpolant
+
+
+class EquilibriumInterpolant(BaseEquilibriumInterpolant):
+    r"""Straight-line path with an energy-compatible (equilibrium) target.
+
+    Equilibrium Matching (EqM) keeps the rectified-flow straight-line path but
+    parametrises it by an interpolation coefficient \(\gamma \in [0, 1]\), a
+    noise level rather than a dynamical time. The learned field is trained to be
+    *time-invariant* (it approximates an equilibrium gradient), so \(\gamma\)
+    only indexes where along the noise-to-data path a training point sits, it is
+    not integrated over. In the paper's notation, for data \(x\) and Gaussian
+    noise \(\epsilon\),
+
+    \[
+    x_\gamma = \gamma\,x + (1 - \gamma)\,\epsilon,
+    \]
+
+    which, with DeltaFlow's \((x_0, x_1)\) = (noise, data) naming, is the same
+    straight line
+
+    \[
+    x_\gamma = (1 - \gamma)\,x_0 + \gamma\,x_1, \qquad \gamma \in [0, 1].
+    \]
+
+    EqM reshapes the regression target so the learned field becomes the
+    gradient of an implicit energy landscape rather than a time-conditional
+    velocity. Instead of regressing onto the constant displacement
+    \(x_1 - x_0\), it regresses onto the scaled displacement (the paper's
+    \((x - \epsilon)\,c(\gamma)\)),
+
+    \[
+    u_\gamma = c(\gamma)\,(x_1 - x_0),
+    \]
+
+    where \(c(\gamma)\) is the equilibrium coefficient. The key design
+    constraint is \(c(1) = 0\): the target vanishes at data, so ground-truth
+    samples become stationary points (local minima) of the landscape whose
+    gradient the field learns. Away from data the coefficient is held on a
+    constant plateau, so the field points from noise toward data with roughly
+    constant magnitude.
+
+    Concretely the coefficient is the minimum of two lines, rescaled by
+    ``scale``,
+
+    \[
+    c(\gamma) = \text{scale}\cdot\min\!\Bigl(
+        \text{start} - \tfrac{\text{start} - 1}{p}\,\gamma,\;
+        \tfrac{1 - \gamma}{1 - p}
+    \Bigr),
+    \]
+
+    with plateau fraction \(p\) (``plateau``). With the defaults
+    (``start = 1``, ``plateau = 0.8``, ``scale = 4``) the first line is flat at
+    \(1\), so \(c(\gamma) = 4\,\min(1, 5(1 - \gamma))\): a plateau of \(4\) for
+    \(\gamma \le 0.8\) that then ramps linearly down to \(0\) at \(\gamma = 1\).
+
+    **Interpolation coefficient.** \(\gamma = 0\) is noise (\(x_\gamma = x_0\))
+    and \(\gamma = 1\) is data (\(x_\gamma = x_1\)). This interpolant subclasses
+    `BaseEquilibriumInterpolant`, whose ``interpolate`` takes ``gamma`` rather
+    than the dynamical time ``t`` of `BaseInterpolant`, because \(\gamma\) is a
+    noise level and not a time to be integrated.
+
+    **Coupling.** This interpolant defines only the *path* and *target*. It is
+    agnostic to how \((x_0, x_1)\) pairs are formed. If ``x0`` is not supplied
+    it is drawn from \(\mathcal{N}(0, I)\) independently of ``x1``. Pass an
+    `OTCoupling` on the training side for
+    straighter, OT-coupled displacements.
+
+    **Sampling.** A field trained with this target is *not* integrated over a
+    fixed time horizon. Because it approximates an equilibrium gradient it is
+    sampled by optimisation (gradient descent on the landscape), for which see
+    `EquilibriumSolver`.
+
+    References:
+        Wang and Du, "Equilibrium Matching: Generative Modeling with Implicit
+        Energy-Based Models" (2025), https://arxiv.org/abs/2510.02300.
+
+    Args:
+        plateau: fraction \(p \in (0, 1)\) of the path spent before the
+            coefficient starts ramping down to zero. The ramp occupies the
+            final ``1 - plateau`` of the \([0, 1]\) interval.
+        scale: overall magnitude applied to the coefficient. Sets the typical
+            norm of the learned equilibrium gradient away from data.
+        start: value of the first (upper) line at ``gamma=0``. Left at ``1``
+            this line stays flat so the coefficient plateaus at ``scale``,
+            values other than ``1`` tilt the plateau.
+    """
+
+    def __init__(self, plateau: float = 0.8, scale: float = 4.0, start: float = 1.0):
+        if not 0.0 < plateau < 1.0:
+            raise ValueError(f"plateau must lie in (0, 1), got {plateau!r}")
+        self.plateau = plateau
+        self.scale = scale
+        self.start = start
+
+    def equilibrium_coefficient(self, gamma: torch.Tensor) -> torch.Tensor:
+        r"""Return the scalar coefficient \(c(\gamma)\) applied to \(x_1 - x_0\).
+
+        ``gamma`` is the interpolation coefficient (noise level), \(0\) at noise
+        and \(1\) at data, not a dynamical time.
+        """
+        line_plateau = self.start - (self.start - 1.0) / self.plateau * gamma
+        line_ramp = (1.0 - gamma) / (1.0 - self.plateau)
+        return self.scale * torch.minimum(line_plateau, line_ramp)
+
+    def interpolate(
+        self, x1: torch.Tensor, gamma: torch.Tensor, x0: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if x0 is None:
+            x0 = torch.randn_like(x1)
+        gamma = gamma.view(-1, *([1] * (x1.dim() - 1)))
+        x_gamma = (1 - gamma) * x0 + gamma * x1
+        target_v = self.equilibrium_coefficient(gamma) * (x1 - x0)
+        return x_gamma, target_v

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/deltaflow/interpolants/linear.py

@@ -43,7 +43,7 @@ class LinearInterpolant(BaseInterpolant):
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         if x0 is None:
             x0 = torch.randn_like(x1)
-        t_ = t.view(-1, *([1] * (x1.dim() - 1)))
+        t_ = t.view(-1, *([1] * (x1.dim() - 1)))  # reshape t to match the dimensions of x1 for broadcasting e.g., (B, 1, 1, 1) -> (B, C, H, W) during operation
         x_t = (1 - t_) * x0 + t_ * x1
         target_v = x1 - x0
         return x_t, target_v

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/deltaflow/interpolants/ot.py

@@ -10,9 +10,10 @@ optimal transport" (arXiv:2302.00482), and the OT-vs-independent ablation
 in "Flower: A Flow-Matching Solver for Inverse Problems" (arXiv:2509.26287)
 for the sampling side.
 
-Exact optimal assignment (Hungarian algorithm) is used when ``scipy`` is
-installed, otherwise a deterministic greedy nearest-neighbour fallback is
-used. This is a coupling strategy, not a new probability path, the same
+Exact optimal assignment (Hungarian algorithm) is used by default, since
+``scipy`` is a core dependency. A deterministic greedy nearest-neighbour
+fallback is used only if ``scipy`` is unavailable in the environment. This
+is a coupling strategy, not a new probability path, the same
 straight-line ``x_t = (1-t) x0 + t x1`` interpolation is applied after
 permuting.
 """
@@ -22,43 +23,11 @@ from typing import Optional, Tuple
 import torch
 
 from ..core.base_interpolant import BaseInterpolant
+from ..utils.ot import batch_ot_permutation
 from .linear import LinearInterpolant
 
-
-def _batch_ot_permutation(x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor:
-    """Return a permutation ``perm`` such that ``x0[perm]`` is OT-coupled to ``x1``.
-
-    Costs are squared L2 distances on the flattened per-sample tensors.
-    """
-    b = x0.shape[0]
-    if b == 1:
-        return torch.zeros(1, dtype=torch.long, device=x0.device)
-
-    x0f = x0.reshape(b, -1).float()
-    x1f = x1.reshape(b, -1).float()
-    cost = torch.cdist(x0f, x1f) ** 2  # (B, B), cost[i, j] = |x0[i] - x1[j]|^2
-
-    try:
-        from scipy.optimize import linear_sum_assignment
-
-        row_ind, col_ind = linear_sum_assignment(cost.detach().cpu().numpy())
-        # linear_sum_assignment guarantees row_ind == 0..B-1 in sorted order;
-        # col_ind[i] is the x1 index paired with x0[i]. We want a permutation
-        # of x0 aligned to x1's original order: for each x1[j], take x0[i] where col_ind[i] == j.
-        col = torch.as_tensor(col_ind, dtype=torch.long, device=x0.device)
-        perm = torch.argsort(col)
-        return perm
-    except ImportError:
-        # Greedy fallback: for each x1[j] in order, pick the closest un-used x0[i].
-        used = torch.zeros(b, dtype=torch.bool, device=x0.device)
-        perm = torch.empty(b, dtype=torch.long, device=x0.device)
-        for j in range(b):
-            row = cost[:, j].clone()
-            row[used] = float("inf")
-            i = int(torch.argmin(row).item())
-            perm[j] = i
-            used[i] = True
-        return perm
+# Backward-compat alias: the implementation now lives in ``deltaflow.utils.ot``.
+_batch_ot_permutation = batch_ot_permutation
 
 
 class OTInterpolant(BaseInterpolant):
@@ -91,9 +60,9 @@ class OTInterpolant(BaseInterpolant):
     objective is identical to standard conditional flow matching and no other
     component (loss, solver, model) needs to change.
 
-    **Solver.** The exact assignment (Hungarian algorithm) is used when
-    ``scipy`` is installed, otherwise a deterministic greedy nearest-neighbour
-    fallback is used.
+    **Solver.** The exact assignment (Hungarian algorithm) is used by
+    default, since ``scipy`` is a core dependency. A deterministic greedy
+    nearest-neighbour fallback is used only if ``scipy`` is unavailable.
 
     References:
         Tong et al., "Improving and generalizing flow-based generative
@@ -110,6 +79,6 @@ class OTInterpolant(BaseInterpolant):
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         if x0 is None:
             x0 = torch.randn_like(x1)
-        perm = _batch_ot_permutation(x0, x1)
+        perm = batch_ot_permutation(x0, x1)
         x0 = x0[perm]
         return self._linear.interpolate(x1, t, x0=x0)

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/deltaflow/inverse/operators.py

@@ -82,9 +82,10 @@ class BlurOperator(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         pad = self.kernel_size // 2
+        kernel: torch.Tensor = self.get_buffer("kernel")
         return F.conv2d(
             x,
-            self.kernel.to(dtype=x.dtype, device=x.device),
+            kernel.to(dtype=x.dtype, device=x.device),
             padding=pad,
             groups=self.channels,
         )

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/deltaflow/losses/__init__.py

@@ -2,10 +2,12 @@
 
 from .conditional_flow_matching import ConditionalFlowMatchingLoss, FlowMatchingLoss
 from .delta_alignment import DeltaAlignmentLoss, delta_alignment_loss
+from .equilibrium_matching import EquilibriumMatchingLoss
 
 __all__ = [
     "ConditionalFlowMatchingLoss",
     "DeltaAlignmentLoss",
+    "EquilibriumMatchingLoss",
     "FlowMatchingLoss",
     "delta_alignment_loss",
 ]

torchdeltaflow-0.2.3.dev23/deltaflow/losses/equilibrium_matching.py

@@ -0,0 +1,97 @@
+"""Equilibrium Matching loss (energy-compatible target, time-invariant field)."""
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import torch.nn.functional as F
+
+from ..core.base_equilibrium_interpolant import BaseEquilibriumInterpolant
+from ..core.base_loss import BaseLoss
+from ..interpolants.equilibrium import EquilibriumInterpolant
+
+if TYPE_CHECKING:
+    from ..trainer.coupling import BaseCoupling
+
+
+class EquilibriumMatchingLoss(BaseLoss):
+    r"""Regress a time-invariant field onto the Equilibrium Matching target.
+
+    Equilibrium Matching (EqM) trains a field \(f_\theta\) to match an
+    energy-compatible target along the straight noise-to-data path. For an
+    interpolation coefficient \(\gamma \in [0, 1]\) (a noise level), noise
+    \(x_0\) and data \(x_1\), the corrupted sample and target are produced by
+    an `EquilibriumInterpolant`,
+
+    \[
+    x_\gamma = (1 - \gamma)\,x_0 + \gamma\,x_1, \qquad
+    u_\gamma = c(\gamma)\,(x_1 - x_0),
+    \]
+
+    and the objective is (the paper's Eq. 3, up to the target-sign convention
+    below)
+
+    \[
+    \mathcal{L}_{\text{EqM}} =
+        \mathbb{E}_{\gamma,\, x_0,\, x_1}
+        \bigl\| f_\theta(x_\gamma) - u_\gamma \bigr\|^2 .
+    \]
+
+    **\(\gamma\) is implicit, and there is no time at all.** Unlike flow
+    matching, where the model is conditioned on time \(t\), an EqM field is
+    time-invariant and noise-unconditional, \(f_\theta(x)\). The coefficient
+    \(\gamma\) is *not seen by the model*, and neither is any surrogate time.
+    This loss therefore queries the model as ``model(x, **cond)``, passing no
+    time (and no \(\gamma\)) at all, matching `BaseEquilibriumField`.
+
+    **Target sign.** The interpolant returns \(u_\gamma = c(\gamma)(x_1 - x_0)\)
+    (data minus noise), so the field points from noise toward data and is
+    sampled by gradient *ascent* \(x \leftarrow x + \eta f(x)\) in
+    `EquilibriumSolver`. This matches the official EqM code. The
+    paper writes the mirror-image \((\epsilon - x)c(\gamma)\) with descent, the
+    two conventions are equivalent under a global sign flip.
+
+    References:
+        Wang and Du, "Equilibrium Matching: Generative Modeling with Implicit
+        Energy-Based Models" (2025), https://arxiv.org/abs/2510.02300.
+
+    Args:
+        interpolant: the energy-compatible path. Defaults to
+            `EquilibriumInterpolant`.
+        coupling: optional train-time coupling that produces \((x_0, x_1)\)
+            pairs from a batch of \(x_1\). See
+            `deltaflow.trainer.coupling`.
+        loss_type: one of ``"l2"``, ``"l1"``, ``"huber"``.
+    """
+
+    def __init__(
+        self,
+        interpolant: Optional[BaseEquilibriumInterpolant] = None,
+        coupling: Optional["BaseCoupling"] = None,
+        loss_type: str = "l2",
+    ):
+        self.interpolant = interpolant or EquilibriumInterpolant()
+        self.coupling = coupling
+        self.loss_type = loss_type
+
+    def _reduce(self, target: torch.Tensor, pred: torch.Tensor) -> torch.Tensor:
+        if self.loss_type == "l1":
+            return F.l1_loss(target, pred)
+        if self.loss_type == "l2":
+            return F.mse_loss(target, pred)
+        if self.loss_type == "huber":
+            return F.smooth_l1_loss(target, pred)
+        raise NotImplementedError(f"Unknown loss_type: {self.loss_type!r}")
+
+    def __call__(self, model, x1: torch.Tensor, **cond) -> torch.Tensor:
+        if self.coupling is not None:
+            x0, x1 = self.coupling.sample_pair(x1)
+        else:
+            x0 = None
+        gamma = torch.rand(x1.shape[0], device=x1.device)
+        x_gamma, target = self.interpolant.interpolate(x1, gamma, x0=x0)
+        # gamma is implicit and there is no time: the field is queried as f(x).
+        pred = model(x_gamma, **cond)
+        return self._reduce(target, pred)
+
+
+__all__ = ["EquilibriumMatchingLoss"]

{torchdeltaflow-0.2.2 → torchdeltaflow-0.2.3.dev23}/deltaflow/models/__init__.py

@@ -1,13 +1,21 @@
-"""Model components: velocity-field backbone wrappers, projector heads, EMA."""
+"""Model components: velocity-field backbone wrappers, projector heads, EMA,
+and the DiT transformer with adaLN-Zero conditioning."""
 
 from .backbone import TinyVelocityField, WrappedBackbone
+from .dit import DiT, DiTBlock, FinalLayer, LabelEmbedding, TimestepEmbedding, modulate
 from .ema import EMA
 from .projector import MultiScaleProjector, ProjectorHead
 
 __all__ = [
+    "DiT",
+    "DiTBlock",
     "EMA",
+    "FinalLayer",
+    "LabelEmbedding",
     "MultiScaleProjector",
     "ProjectorHead",
+    "TimestepEmbedding",
     "TinyVelocityField",
     "WrappedBackbone",
+    "modulate",
 ]