Rectified Flow: A Marginal Preserving Approach to Optimal Transport

import numpy as np

# Jensen's inequality (Section 2): between fixed end points, the straight line is the cheapest path.
# Path-wise cost with the quadratic cost c(v) = v^2: the integral over t of (dx/dt)^2.
x0, x1 = -1.0, 3.0
t = np.linspace(0.0, 1.0, 100_001)
dt = t[1] - t[0]

def path_cost(a):
    # x_t = straight line + a * sin(2 pi t): same end points for every a
    velocity = (x1 - x0) + a * 2 * np.pi * np.cos(2 * np.pi * t)
    return float(np.sum(velocity ** 2) * dt)

print("cost of the direct move c(x1 - x0) =", (x1 - x0) ** 2)
for a in (0.0, 0.1, 0.5, 1.0):
    print(f"wiggle amplitude {a:3.1f}: path-wise cost {path_cost(a):8.3f}   (extra {path_cost(a) - (x1 - x0) ** 2:.3f}, theory 2 pi^2 a^2 = {2 * np.pi ** 2 * a * a:.3f})")
print("the straight line (amplitude 0) attains the direct cost; every wiggle costs more")

Output:

cost of the direct move c(x1 - x0) = 16.0
wiggle amplitude 0.0: path-wise cost   16.000   (extra 0.000, theory 2 pi^2 a^2 = 0.000)
wiggle amplitude 0.1: path-wise cost   16.198   (extra 0.198, theory 2 pi^2 a^2 = 0.197)
wiggle amplitude 0.5: path-wise cost   20.935   (extra 4.935, theory 2 pi^2 a^2 = 4.935)
wiggle amplitude 1.0: path-wise cost   35.740   (extra 19.740, theory 2 pi^2 a^2 = 19.739)
the straight line (amplitude 0) attains the direct cost; every wiggle costs more

import numpy as np
rng = np.random.default_rng(0)

# One round of rectified flow in one dimension (Section 3 of the paper).
# pi0 = standard normal noise, pi1 = the two points -2 and +2 (equal probability).
N = 40_000
x0 = rng.normal(size=N)
x1 = rng.choice([-2.0, 2.0], N)             # start with the independent coupling: noise and data paired at random

def velocity(z, t):
    # v_t(z) = E[X1 - X0 | X_t = z] for the straight-line interpolation X_t = t X1 + (1 - t) X0 (exact for this toy).
    w_pos = np.exp(-0.5 * ((z - 2.0 * t) / (1 - t)) ** 2)
    w_neg = np.exp(-0.5 * ((z + 2.0 * t) / (1 - t)) ** 2)
    p = w_pos / (w_pos + w_neg)
    m = 2.0 * p - 2.0 * (1 - p)             # E[X1 | X_t = z]
    return m - (z - t * m) / (1 - t)        # E[X1 - X0 | X_t = z]

# Follow the ODE dZ = v(Z, t) dt from Z0 = X0 (midpoint method, many small steps), and record the path-wise cost.
steps = 600
dt = 1.0 / steps
z = x0.copy()
path_cost = 0.0                             # integral over time of E[ (dZ/dt)^2 ]
for k in range(steps):
    t = k * dt
    mid = z + 0.5 * dt * velocity(z, t)
    v_mid = velocity(mid, t + 0.5 * dt)
    path_cost += np.mean(v_mid ** 2) * dt
    z = z + dt * v_mid
z1 = z                                      # the rectified coupling is (Z0, Z1) = (x0, z1)

# The least-squares loss of the regression (Equation 3), by Monte Carlo over random times.
ts = (np.arange(200) + 0.5) / 200
loss = np.mean([np.mean(((x1 - x0) - velocity((1 - t) * x0 + t * x1, t)) ** 2) for t in ts])

print("marginal of Z1: share of samples above 0 =", round(float(np.mean(z1 > 0)), 3), "(target 0.5);",
      "mean |Z1| =", round(float(np.mean(np.abs(z1))), 3), "(target 2)")
for name, c in (("c(x) = x^2", lambda d: d ** 2), ("c(x) = |x|", np.abs), ("c(x) = x^4", lambda d: d ** 4)):
    before, after = np.mean(c(x1 - x0)), np.mean(c(z1 - x0))
    print(f"cost {name:<10} independent coupling {before:6.3f} -> rectified coupling {after:6.3f}")

# The optimal coupling in one dimension pairs the smaller half of the noise with -2 and the larger half with +2.
optimal_z1 = np.where(x0 > np.median(x0), 2.0, -2.0)
print("cost x^2 of the optimal (sorted) coupling:", round(float(np.mean((optimal_z1 - x0) ** 2)), 3))

# Equation 28 for the quadratic cost: cost drop = loss of the regression + straightness gap S.
drop = np.mean((x1 - x0) ** 2) - np.mean((z1 - x0) ** 2)
S = path_cost - np.mean((z1 - x0) ** 2)
print(f"cost drop {drop:.3f} against regression loss + S = {loss:.3f} + {S:.3f} = {loss + S:.3f}")

Output:

marginal of Z1: share of samples above 0 = 0.499 (target 0.5); mean |Z1| = 2.0 (target 2)
cost c(x) = x^2 independent coupling  4.996 -> rectified coupling  1.804
cost c(x) = |x| independent coupling  2.015 -> rectified coupling  1.236
cost c(x) = x^4 independent coupling 43.176 -> rectified coupling  4.673
cost x^2 of the optimal (sorted) coupling: 1.804
cost drop 3.192 against regression loss + S = 2.038 + 1.149 = 3.187

import numpy as np
rng = np.random.default_rng(0)

# Example 3.5: pi0 = pi1 = N(0, I) in two dimensions. The best coupling is the identity (X0, X0), with cost 0.
# Rotation coupling (X0, A X0), where A rotates by 90 degrees, is also a valid coupling of the two distributions.
theta = np.pi / 2
A = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
I = np.eye(2)

N = 200_000
x0 = rng.normal(size=(N, 2))
x1 = x0 @ A.T

print("marginal of X1 is still standard normal: mean", np.round(x1.mean(axis=0), 3).tolist(), ", covariance", np.round(np.cov(x1.T), 2).tolist())

# Cost with the quadratic cost c(x) = |x|^2 / 2.
cost_rot = np.mean(0.5 * np.sum((x1 - x0) ** 2, axis=1))
print("cost of the rotation coupling:", round(float(cost_rot), 3), "| theory: 2(1 - cos(theta)) =", round(2 * (1 - np.cos(theta)), 3), "| cost of the identity coupling: 0")

# Straightness: on the straight line interpolation, the velocity at X_t is fully determined by X_t,
# v_t(x) = (A - I)(tA + (1-t)I)^-1 x, so regression leaves no residual (Equation 32 of the appendix).
t = 0.3
xt = (1 - t) * x0 + t * x1
direction = x1 - x0
v = xt @ (((A - I) @ np.linalg.inv(t * A + (1 - t) * I)).T)
print("largest difference between the line direction and the regression field:", float(np.abs(direction - v).max()))
print("-> the coupling is straight (a fixed point of the rectification), yet its cost is", round(float(cost_rot), 3), "instead of 0")

Output:

marginal of X1 is still standard normal: mean [0.001, 0.001] , covariance [[1.0, 0.0], [0.0, 1.0]]
cost of the rotation coupling: 2.006 | theory: 2(1 - cos(theta)) = 2.0 | cost of the identity coupling: 0
largest difference between the line direction and the regression field: 1.7763568394002505e-15
-> the coupling is straight (a fixed point of the rectification), yet its cost is 2.006 instead of 0

import numpy as np

# The matching loss of Equation 23:  m_c(x; y) = c(x) - x*y + c*(y),  where c* is the convex conjugate of c.
# Example cost c(x) = x^4 / 4, whose conjugate is c*(y) = (3/4) |y|^(4/3) (in one dimension).
c = lambda x: x ** 4 / 4
c_star = lambda y: 0.75 * np.abs(y) ** (4 / 3)
grad_c = lambda x: x ** 3                       # derivative of c

def m(x, y):
    return c(x) - x * y + c_star(y)

xs = np.linspace(-2, 2, 9)
ys = np.linspace(-8, 8, 4001)
print("x     | smallest value of m_c(x; y) over y | the y where it is smallest | derivative of c at x")
for x in xs:
    values = m(x, ys)
    j = values.argmin()
    print(f"{x:5.1f} | {round(float(values[j]), 6) + 0.0:34.6f} | {ys[j]:26.3f} | {grad_c(x):>8.3f}")

print()
print("m_c is never negative on the grid:", bool(all(m(x, ys).min() > -1e-9 for x in xs)))

# For the quadratic cost c(x) = x^2 / 2 (so c* = c) it reduces to half the squared error.
cq = lambda x: 0.5 * x ** 2
mq = lambda x, y: cq(x) - x * y + cq(y)
x, y = 1.3, -0.4
print("quadratic cost: m_c(1.3; -0.4) =", round(float(mq(x, y)), 4), "and (x - y)^2 / 2 =", round(0.5 * (x - y) ** 2, 4))

Output:

x     | smallest value of m_c(x; y) over y | the y where it is smallest | derivative of c at x
 -2.0 |                           0.000000 |                     -8.000 |   -8.000
 -1.5 |                           0.000000 |                     -3.376 |   -3.375
 -1.0 |                           0.000000 |                     -1.000 |   -1.000
 -0.5 |                           0.000001 |                     -0.124 |   -0.125
  0.0 |                           0.000000 |                      0.000 |    0.000
  0.5 |                           0.000001 |                      0.124 |    0.125
  1.0 |                           0.000000 |                      1.000 |    1.000
  1.5 |                           0.000000 |                      3.376 |    3.375
  2.0 |                           0.000000 |                      8.000 |    8.000

m_c is never negative on the grid: True
quadratic cost: m_c(1.3; -0.4) = 1.445 and (x - y)^2 / 2 = 1.445

# Corollary 5.7: summing the cost decreases of K + 1 rectification steps cannot exceed the cost you started with,
#   sum over k of ( loss_k + S_k ) <= E c(X1 - X0)  =>  min over k of ( loss_k + S_k ) <= E c(X1 - X0) / (K + 1).
initial_cost = 5.0                  # a made-up starting transport cost
print("steps K+1 | guaranteed bound on the smallest (loss + S) among the steps")
for steps in (1, 2, 5, 10, 50, 100, 1000):
    print(f"{steps:>9} | at most {initial_cost / steps:.4f}")

# A made-up sequence of costs after each rectification, with its decreases (each equals loss_k + S_k).
costs = [5.0, 2.0, 1.2, 1.0, 0.95, 0.93]
drops = [a - b for a, b in zip(costs[:-1], costs[1:])]
print()
print("decreases:", [round(d, 2) for d in drops], "| they add up to", round(sum(drops), 2), "<= initial cost", initial_cost)
print("smallest decrease so far:", round(min(drops), 2), "<=", round(initial_cost / len(drops), 2), "= initial cost / number of steps")
print("note: the bound is on loss + S, a surrogate of optimality; it does not say how far the cost is from the optimal cost")

Output:

steps K+1 | guaranteed bound on the smallest (loss + S) among the steps
        1 | at most 5.0000
        2 | at most 2.5000
        5 | at most 1.0000
       10 | at most 0.5000
       50 | at most 0.1000
      100 | at most 0.0500
     1000 | at most 0.0050

decreases: [3.0, 0.8, 0.2, 0.05, 0.02] | they add up to 4.07 <= initial cost 5.0
smallest decrease so far: 0.02 <= 1.0 = initial cost / number of steps
note: the bound is on loss + S, a surrogate of optimality; it does not say how far the cost is from the optimal cost