File size: 3,234 Bytes
b2cb4a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | #
# For licensing see accompanying LICENSE file.
# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
#
def right_pad_dims_to(x, t):
padding_dims = x.ndim - t.ndim
if padding_dims <= 0:
return t
return t.reshape(*t.shape, *((1,) * padding_dims))
class BasePath:
"""base class for flow matching path"""
def __init__(self):
return
def compute_alpha_t(self, t):
"""Compute the data coefficient along the path"""
return None, None
def compute_sigma_t(self, t):
"""Compute the noise coefficient along the path"""
return None, None
def compute_d_alpha_alpha_ratio_t(self, t):
"""Compute the ratio between d_alpha and alpha"""
alpha_t, d_alpha_t = self.compute_alpha_t(t)
return d_alpha_t / alpha_t
def compute_mu_t(self, t, x0, x1):
"""Compute the mean of time-dependent density p_t"""
alpha_t, _ = self.compute_alpha_t(t)
sigma_t, _ = self.compute_sigma_t(t)
return alpha_t * x1 + sigma_t * x0
def compute_xt(self, t, x0, x1):
"""Sample xt from time-dependent density p_t; rng is required"""
xt = self.compute_mu_t(t, x0, x1)
return xt
def compute_ut(self, t, x0, x1):
"""Compute the vector field corresponding to p_t"""
_, d_alpha_t = self.compute_alpha_t(t)
_, d_sigma_t = self.compute_sigma_t(t)
return d_alpha_t * x1 + d_sigma_t * x0
def interpolant(self, t, x0, x1):
t = right_pad_dims_to(x0, t)
xt = self.compute_xt(t, x0, x1)
ut = self.compute_ut(t, x0, x1)
return t, xt, ut
def compute_drift(self, x, t):
"""We always output sde according to score parametrization; """
t = right_pad_dims_to(x, t)
alpha_ratio = self.compute_d_alpha_alpha_ratio_t(t)
sigma_t, d_sigma_t = self.compute_sigma_t(t)
drift_mean = alpha_ratio * x
drift_var = alpha_ratio * (sigma_t ** 2) - sigma_t * d_sigma_t
return -drift_mean, drift_var
def compute_score_from_velocity(self, v_t, y_t, t):
t = right_pad_dims_to(y_t, t)
alpha_t, d_alpha_t = self.compute_alpha_t(t)
sigma_t, d_sigma_t = self.compute_sigma_t(t)
mean = y_t
reverse_alpha_ratio = alpha_t / d_alpha_t
var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
score = (reverse_alpha_ratio * v_t - mean) / var
return score
def compute_velocity_from_score(self, s_t, y_t, t):
t = right_pad_dims_to(y_t, t)
drift_mean, drift_var = self.compute_drift(y_t, t)
velocity = -drift_mean + drift_var * s_t
return velocity
class LinearPath(BasePath):
"""
Linear flow process:
x0: noise, x1: data
In inference, we sample data from 0 -> 1
"""
def __init__(self):
super().__init__()
def compute_alpha_t(self, t):
"""Compute the data coefficient along the path"""
return t, 1
def compute_sigma_t(self, t):
"""Compute the noise coefficient along the path"""
return 1 - t, -1
def compute_d_alpha_alpha_ratio_t(self, t):
"""Compute the ratio between d_alpha and alpha"""
return 1 / t
|