Upload sinkhorn_flow.py

sinkhorn_flow.py

@@ -0,0 +1,163 @@
+"""sinkhorn_flow.py — Sinkhorn gradient flow and W_ε potential computation.
+
+Core implementation of:
+  - Sinkhorn divergence computation via GeomLoss
+  - W_ε-potential gradients (∇f_{μ,μ} and ∇f_{μ,μ*})
+  - Velocity field: v(x) = ∇f_{μ,μ}(x) - ∇f_{μ,μ*}(x)  (Theorem 1, Eq. 10)
+  - Euler discretization of the Sinkhorn WGF (Algorithm 1)
+  - Trajectory pool construction for velocity field matching
+
+Reference: arXiv:2401.14069, Section 4.1, 4.3, Appendix A
+"""
+
+import torch
+import torch.nn as nn
+from typing import List, Tuple, Optional
+from geomloss import SamplesLoss
+
+
+class SinkhornPotentialComputer:
+    """Computes W_ε-potentials and their gradients using GeomLoss.
+
+    The velocity field of the Sinkhorn WGF is (Theorem 1):
+        v(x) = ∇f_{μ,μ}(x) - ∇f_{μ,μ*}(x)
+
+    Args:
+        blur: GeomLoss blur parameter (related to ε: ε = blur^p).
+        scaling: Multiscale scaling parameter for Sinkhorn iterations.
+        p: Cost exponent (default 2 for squared Euclidean).
+        backend: GeomLoss backend ('auto', 'tensorized', 'online').
+    """
+
+    def __init__(self, blur: float = 0.5, scaling: float = 0.80,
+                 p: int = 2, backend: str = "tensorized"):
+        self.blur = blur
+        self.scaling = scaling
+        self.p = p
+        self.backend = backend
+
+        self.loss_fn = SamplesLoss(
+            loss="sinkhorn", p=p, blur=blur, scaling=scaling,
+            backend=backend, potentials=True,
+        )
+        self.loss_monitor = SamplesLoss(
+            loss="sinkhorn", p=p, blur=blur, scaling=scaling,
+            backend=backend, potentials=False,
+        )
+
+    def compute_velocity(self, X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
+        """Compute the Sinkhorn WGF velocity field at particles X.
+
+        v(X_i) = ∇f_{μ,μ}(X_i) - ∇f_{μ,μ*}(X_i)
+        """
+        X_grad = X.detach().clone().requires_grad_(True)
+        Y_det = Y.detach()
+
+        F_self, _ = self.loss_fn(X_grad, X_grad.detach().clone())
+        grad_self = torch.autograd.grad(
+            F_self.sum(), X_grad, create_graph=False, retain_graph=False
+        )[0]
+
+        X_grad2 = X.detach().clone().requires_grad_(True)
+        F_cross, _ = self.loss_fn(X_grad2, Y_det)
+        grad_cross = torch.autograd.grad(
+            F_cross.sum(), X_grad2, create_graph=False, retain_graph=False
+        )[0]
+
+        velocity = grad_self.detach() - grad_cross.detach()
+        return velocity
+
+    def compute_sinkhorn_divergence(self, X: torch.Tensor, Y: torch.Tensor) -> float:
+        with torch.no_grad():
+            return self.loss_monitor(X, Y).item()
+
+
+class SinkhornGradientFlow:
+    """Implements the discrete Sinkhorn Wasserstein Gradient Flow.
+
+    Evolves particles via Euler steps:
+        X^{t+1} = X^t + η * v(X^t)
+    """
+
+    def __init__(self, potential_computer: SinkhornPotentialComputer,
+                 eta: float = 1.0, num_steps: int = 5):
+        self.potential_computer = potential_computer
+        self.eta = eta
+        self.num_steps = num_steps
+
+    def run_flow(self, X0: torch.Tensor, Y: torch.Tensor,
+                 store_trajectory: bool = True
+                 ) -> Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor, int]]]:
+        trajectory = []
+        X_t = X0.clone()
+
+        for t in range(self.num_steps):
+            v_t = self.potential_computer.compute_velocity(X_t, Y)
+            if store_trajectory:
+                trajectory.append((
+                    X_t.detach().cpu().clone(),
+                    v_t.detach().cpu().clone(),
+                    t,
+                ))
+            X_t = X_t.detach() + self.eta * v_t.detach()
+
+        return X_t, trajectory
+
+    def run_flow_no_store(self, X0: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
+        X_T, _ = self.run_flow(X0, Y, store_trajectory=False)
+        return X_T
+
+
+class TrajectoryPool:
+    """Stores (x, v, t) tuples from Sinkhorn gradient flow trajectories."""
+
+    def __init__(self, max_size: int = 1_000_000):
+        self.max_size = max_size
+        self.x_pool: List[torch.Tensor] = []
+        self.v_pool: List[torch.Tensor] = []
+        self.t_pool: List[int] = []
+        self._size = 0
+
+    def add_trajectory(self, trajectory: List[Tuple[torch.Tensor, torch.Tensor, int]]):
+        for x, v, t in trajectory:
+            n = x.shape[0]
+            if self._size + n > self.max_size:
+                excess = (self._size + n) - self.max_size
+                self._drop_oldest(excess)
+            self.x_pool.append(x)
+            self.v_pool.append(v)
+            self.t_pool.extend([t] * n)
+            self._size += n
+
+    def _drop_oldest(self, n: int):
+        removed = 0
+        while removed < n and len(self.x_pool) > 0:
+            batch_size = self.x_pool[0].shape[0]
+            if removed + batch_size <= n:
+                self.x_pool.pop(0)
+                self.v_pool.pop(0)
+                self.t_pool = self.t_pool[batch_size:]
+                removed += batch_size
+                self._size -= batch_size
+            else:
+                keep = batch_size - (n - removed)
+                self.x_pool[0] = self.x_pool[0][-keep:]
+                self.v_pool[0] = self.v_pool[0][-keep:]
+                self.t_pool = self.t_pool[(batch_size - keep):]
+                self._size -= (batch_size - keep)
+                removed = n
+
+    def sample(self, batch_size: int, device: str = "cpu"
+               ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        all_x = torch.cat(self.x_pool, dim=0)
+        all_v = torch.cat(self.v_pool, dim=0)
+        all_t = torch.tensor(self.t_pool, dtype=torch.float32)
+        idx = torch.randint(0, all_x.shape[0], (batch_size,))
+        return all_x[idx].to(device), all_v[idx].to(device), all_t[idx].to(device)
+
+    @property
+    def size(self) -> int:
+        return self._size
+
+    def __len__(self) -> int:
+        return self._size