Spaces:
Running on Zero
Running on Zero
File size: 4,095 Bytes
36cdb93 | 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 | """Distribution Matching Distillation for the blockcausal student.
The generator loss is the DMD2 form. Given a generated clean latent `x0`, noise
it to some level t and ask two score networks where they think it came from:
x0_real = teacher(x_t, t) with CFG -- the distribution we want
x0_fake = critic(x_t, t) -- the distribution we have
grad = (x0_fake - x0_real) / normalizer
L_gen = 0.5 * || x0 - stopgrad(x0 - grad) ||^2
so d L_gen / d x0 == grad, which is the KL gradient between the two
distributions. The normalizer is DMD2 scale invariant one. Which is mean |x0 - x0_real|
. Without which the loss magnitude swings by orders of magnitude across t and
the generator LR cannot be set at all.
The critic is trained in the ordinary way, a flow-matching loss on the student's
own samples, so it tracks a moving target. That is the whole reason DMD needs
two time scales: the critic must stay ahead of the generator.
Why the teacher is allowed to be bidirectional: it only ever scores a finished
clip, never generates one. Causality is a constraint on the students sampling
procedure, not on the reference distribution.
"""
import torch
import torch.nn.functional as tnnF
from . import blockcausal as bc
from .data import add_noise
def bidirectional_velocity(model, z, t, rope, ctx, ctx_lens, dtype,
time_scale=1000.0, grad_checkpoint=False):
"""Full bidirectional forward over a whole clip.
`block_forward` with no K/V prefix and the entire clip as one block *is* the
bidirectional forward -- attention runs over the clip and nothing else -- so
teacher, critic and student all go through one code path and one RoPE
implementation.
"""
return bc.block_forward(model, z, t, 0, rope, ctx, ctx_lens, kv=None,
dtype=dtype, time_scale=time_scale,
grad_checkpoint=grad_checkpoint)
def cfg_velocity(model, z, t, rope, ctx_pos, ctx_neg, ctx_lens, dtype,
guidance, time_scale=1000.0):
"""Teacher velocity with classifier-free guidance, cond and uncond batched."""
if not guidance or guidance == 1.0:
return bidirectional_velocity(model, z, t, rope, ctx_pos, ctx_lens,
dtype, time_scale)
z2 = torch.cat([z, z], 0)
ctx2 = torch.cat([ctx_pos, ctx_neg], 0)
lens2 = ctx_lens.expand(2)
v = bidirectional_velocity(model, z2, t, rope, ctx2, lens2, dtype, time_scale)
vc, vu = v[0:1], v[1:2]
return vu + guidance * (vc - vu)
@torch.no_grad()
def dmd_gradient(x0, t, z_t, v_real, v_fake):
"""DMD2's KL gradient wrt the generated clip, for the whole clip at once.
One normalizer for the clip rather than one per block: per-block
normalisation would rescale each block's gradient by its own error and so
quietly up-weight whichever block is currently worst.
"""
while t.dim() < z_t.dim():
t = t.unsqueeze(-1)
x0_real = (z_t - t * v_real).float()
x0_fake = (z_t - t * v_fake).float()
normalizer = (x0.float() - x0_real).abs().mean().clamp_min(1e-4)
grad = torch.nan_to_num((x0_fake - x0_real) / normalizer)
return grad, {'dmd_grad_norm': float(grad.norm()),
'dmd_normalizer': float(normalizer),
'dmd_real_fake_gap': float((x0_real - x0_fake).abs().mean())}
def dmd_surrogate(x0_hat, grad, scale=1.0):
"""Scalar whose gradient wrt x0_hat is exactly `grad` (times scale)."""
return scale * 0.5 * tnnF.mse_loss(x0_hat.float(),
(x0_hat.detach() - grad).float())
def critic_loss(model, x0, t, rope, ctx, ctx_lens, dtype, time_scale=1000.0,
grad_checkpoint=True, noise=None):
"""Flow-matching loss for the fake-score network on student samples."""
z_t, target, _ = add_noise(x0, t, noise)
v = bidirectional_velocity(model, z_t, t, rope, ctx, ctx_lens, dtype,
time_scale, grad_checkpoint=grad_checkpoint)
return tnnF.mse_loss(v.float(), target.float())
|