File size: 7,479 Bytes
da3babb | 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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | """Flow matching, finite-map identities, Shortcut, MeanFlow, and latent maps.
# %% 1. Learn local motion before learning an interval.
All times run noise -> data except MeanFlow, whose original backward clock
r <= t runs data at zero -> noise at one during training.
"""
import torch
from torch import nn
from torch.nn import functional as F
from common import (MapNet, draw_batch, ema_copy, finite_map, interpolate, mlp,
optimize, ordered_times, time_like)
def diagonal_loss(model, data):
t = torch.rand(len(data), 1, device=data.device)
x, displacement = interpolate(data, t)
return F.mse_loss(model(x, t, t), displacement)
def lagrangian_residual(model, velocity, x, s, t):
# d_t F(s,t,x) = b(t,F(s,t,x)). JVP differentiates only arrival time.
y, dt = torch.func.jvp(lambda end: finite_map(model, x, s, end),
(t,), (torch.ones_like(t),))
return dt - velocity(y.detach(), t).detach()
def eulerian_residual(model, velocity, x, s, t):
# (d_s + b_s dot grad_x)F = 0. No full Jacobian is materialized.
direction = velocity(x, s).detach()
_, residual = torch.func.jvp(lambda z, start: finite_map(model, z, start, t),
(x, s), (direction, torch.ones_like(s)))
return residual
def semigroup_loss(model, teacher, x, s, t, context=None):
u = (s + t) / 2
with torch.no_grad():
mid = finite_map(teacher, x, s, u, context)
target = finite_map(teacher, mid, u, t, context)
prediction = finite_map(model, x, s, t, context)
return ((prediction - target) / (t - s).clamp_min(.05)).square().mean()
def meanflow_loss(model, data):
# %% 2. Average backward velocity satisfies u = v - (t-r) D_t u.
noise = torch.randn_like(data)
r, t, _ = ordered_times(data, 1.)
# Include diagonal examples with nonzero probability, as in the paper.
r = torch.where(torch.rand_like(r) < .75, t, r)
z = (1 - t) * data + t * noise
v = noise - data
average, derivative = torch.func.jvp(model, (z, r, t),
(v, torch.zeros_like(r), torch.ones_like(t)))
target = (v - (t - r) * derivative).detach()
residual = (average - target).square().mean(-1)
# Detached adaptive weighting controls large self-distillation residuals.
weight = (residual.detach() + .01).pow(-.5)
return (weight * residual).mean(), {'mse': residual.mean()}
def shortcut_loss(model, teacher, data):
t = torch.rand(len(data), 1, device=data.device)
x, v = interpolate(data, t)
diag = F.mse_loss(model(x, t, t), v)
# d is the half-step. Learn the 2d shortcut from two d shortcuts.
powers = torch.randint(1, 6, (len(data), 1), device=data.device)
d = 2. ** (-powers)
start = torch.rand_like(d) * (1 - 2 * d)
x, _ = interpolate(data, start)
with torch.no_grad():
first = teacher(x, start, start + d)
second = teacher(x + d * first, start + d, start + 2 * d)
target = .5 * (first + second)
finite = F.mse_loss(model(x, start, start + 2 * d), target)
return diag + finite, {'diagonal': diag, 'shortcut': finite}
def train_continuous(method, data, args):
model = MapNet(data.shape[1], args.width).to(data)
teacher_logs = []
teacher = None
if method in {'fmm-lagrangian', 'fmm-eulerian', 'consistency'}:
teacher = MapNet(data.shape[1], args.width).to(data)
teacher_logs = optimize(teacher, lambda _: (diagonal_loss(teacher, draw_batch(data, args.batch_size)), {}),
args.teacher_steps, args.lr)
teacher.requires_grad_(False)
model.load_state_dict(teacher.state_dict())
ema = ema_copy(model)
def objective(step):
batch = draw_batch(data, args.batch_size)
if method == 'meanflow':
return meanflow_loss(model, batch)
if method == 'shortcut':
return shortcut_loss(model, ema, batch)
diag = diagonal_loss(model, batch)
if method == 'flow-matching':
return diag, {'diagonal': diag}
s, t, _ = ordered_times(batch)
x, _ = interpolate(batch, s)
if method == 'fmm-lagrangian':
finite = lagrangian_residual(model, lambda z, a: teacher(z, a, a), x, s, t).square().mean()
elif method == 'fmm-eulerian':
finite = eulerian_residual(model, lambda z, a: teacher(z, a, a), x, s, t).square().mean()
elif method == 'consistency':
# Endpoint consistency distillation on teacher-solved short intervals.
# At t=1 the residual endpoint parametrization is exactly identity.
t = (s + .1).clamp_max(1.)
with torch.no_grad():
y = integrate_velocity(lambda z, a: teacher(z, a, a), x, s, t, 4)
target = finite_map(ema, y, t, torch.ones_like(t))
finite = F.mse_loss(finite_map(model, x, s, torch.ones_like(s)), target)
else:
finite = semigroup_loss(model, ema, x, s, t)
weight = min(1., (step + 1) / max(1, args.train_steps // 5))
return diag + weight * finite, {'diagonal': diag, 'finite': finite}
logs = optimize(model, objective, args.train_steps, args.lr, ema)
# Use the trained student; EMA supplies fixed bootstrap targets during training.
state = {'model': model.state_dict(), 'dim': data.shape[1]}
if teacher is not None:
state['teacher'] = teacher.state_dict()
return model, state, logs, teacher_logs
@torch.no_grad()
def integrate_velocity(velocity, x, s, t, steps):
"""Heun integration of a velocity, also supports batch-specific time bounds."""
s, t = time_like(s, x), time_like(t, x)
h = (t - s) / steps
for i in range(steps):
a = s + i * h
k1 = velocity(x, a)
k2 = velocity(x + h * k1, a + h)
x = x + h * .5 * (k1 + k2)
return x
@torch.no_grad()
def sample_continuous(model, method, noise, steps):
if method == 'flow-matching':
return integrate_velocity(lambda z, a: model(z, a, a), noise, 0., 1., steps)
if method == 'consistency':
return finite_map(model, noise, 0., 1.)
x = noise
for i in range(steps):
if method == 'meanflow':
t, r = 1 - i / steps, 1 - (i + 1) / steps
x = x - (t - r) * model(x, r, t)
else:
x = finite_map(model, x, i / steps, (i + 1) / steps)
return x
class Autoencoder(nn.Module):
def __init__(self, width):
super().__init__()
self.encoder = mlp(3, 2, width)
self.decoder = mlp(2, 3, width)
def embed_surface(x):
return torch.cat([x, .3 * (x[:, :1].square() - x[:, 1:].square())], -1)
def train_latent(data, args):
# %% 3. First fit the representation, then freeze it while fitting its flow.
surface = embed_surface(data)
ae = Autoencoder(args.width).to(data)
def objective(_):
batch = draw_batch(surface, args.batch_size)
reconstruction = ae.decoder(ae.encoder(batch))
return F.mse_loss(reconstruction, batch), {}
ae_logs = optimize(ae, objective, args.teacher_steps, args.lr)
ae.requires_grad_(False)
with torch.no_grad():
z = ae.encoder(surface)
mean, std = z.mean(0), z.std(0).clamp_min(.05)
z = (z - mean) / std
model, state, logs, _ = train_continuous('self-distill', z, args)
state.update({'autoencoder': ae.state_dict(), 'latent_mean': mean, 'latent_std': std})
return model, ae, state, logs, ae_logs
|