Spaces:
Running on Zero
Running on Zero
| """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) | |
| 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()) | |