File size: 6,513 Bytes
3c13fb9 | 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 | """
Flow matching and score matching algorithms for generative modeling
This module implements three training algorithms for ablation:
1. ScoreMatchingDiffusion - Baseline (predict score, VP-SDE diffusion path)
2. FlowMatchingDiffusion - Flow matching with diffusion path (predict velocity)
3. FlowMatchingOT - Flow matching with optimal transport path (predict velocity)
Each algorithm encapsulates both the probability path and loss computation.
"""
import math
import torch
import torch.nn.functional as F
from abc import ABC, abstractmethod
# SD3-style resolution-dependent time shift.
# SD3 (Esser et al., 2024, Eq. 23) derives the shift from matching the
# uncertainty about the clean image across resolutions:
# alpha = sqrt(m / n),
# where n and m are the pixel counts of the reference and target resolutions.
# With a fixed VAE and patch size, token count is proportional to pixel count,
# so the same ratio applies to token counts. Reference: 256 tokens (256×256
# image → 32×32 latent, patch_size=2 → 16×16 patches) gets shift 1.0.
# SD3 empirically found little quality difference among shift values above
# 1.5 and used 3.0 at 1024×1024, where the formula gives 4.0 — so the formula
# values used here sit inside the range they found acceptable.
_SHIFT_BASE_TOKENS = 256 # 256×256 → 32×32 latent → 16×16 patches
def resolution_time_shift(z: torch.Tensor, patch_size: int = 2) -> float:
"""Compute SD3-style time shift from a noise/latent tensor [B, C, H, W]."""
_, _, h, w = z.shape
n_tokens = (h * w) / (patch_size ** 2)
if n_tokens <= _SHIFT_BASE_TOKENS:
return 1.0
return math.sqrt(n_tokens / _SHIFT_BASE_TOKENS)
def apply_time_shift(t: torch.Tensor, shift: float) -> torch.Tensor:
"""Apply the SD3 shift transform in this repo's timestep convention.
SD3 (Esser et al., 2024, Eq. 23) uses t=0 for data and t=1 for noise and
pushes timesteps toward the noisy end at higher resolutions:
u' = (s * u) / (1 + (s - 1) * u).
This repo uses the opposite convention (t=0 noise, t=1 data), so the same
transform applied to u = 1 - t yields:
t' = t / (s - (s - 1) * t),
which decreases t (moves toward noise) when shift > 1.
"""
return t / (shift - (shift - 1) * t)
def shift_timesteps(
t: torch.Tensor,
z_like: torch.Tensor,
*,
patch_size: int = 2,
time_shift: float | None = None,
) -> torch.Tensor:
"""Shift timesteps using the repo's resolution-dependent convention.
This is the shared entry point for both training and inference to ensure
the model always sees the same "t_used" that is implied by the resolution.
Args:
t: Base timesteps in [0, 1], shape [B] or scalar tensor.
z_like: A latent/noise tensor with shape [B, C, H, W] used to infer
the resolution-dependent shift when time_shift is None.
patch_size: Patch size used to convert latent HxW into token count.
time_shift: Optional explicit shift scalar to override auto-compute.
"""
if time_shift is None:
time_shift = resolution_time_shift(z_like, patch_size=patch_size)
return apply_time_shift(t, time_shift)
def sample_weighted_mse(
error: torch.Tensor, sample_weights: torch.Tensor
) -> torch.Tensor:
"""Weighted mean of the per-sample mean squared errors of ``error``.
The samples are combined as ``sum(w_i * L_i) / sum(w_i)``: a weighted mean,
not a weighted sum. A sum would scale a step's gradient with the weights
themselves, which is a learning-rate change wearing a reweighting's
clothes, and two runs would no longer be comparable at equal step counts.
The reduction runs in float32 whatever the compute dtype is, because the
weight sum is a normalization: a denominator rounded to the compute dtype
would bias every loss the weights touch.
"""
if sample_weights.shape != (error.shape[0],):
raise ValueError(
"sample weights must hold one multiplier per sample, got "
f"{tuple(sample_weights.shape)} for a batch of {error.shape[0]}"
)
weights = sample_weights.float()
per_sample = (error.float() ** 2).flatten(1).mean(dim=1)
return (per_sample * weights).sum() / weights.sum()
class BaseAlgorithm(ABC):
@abstractmethod
def sample_zt(
self, z0: torch.Tensor, z1: torch.Tensor, t: torch.Tensor
) -> torch.Tensor:
"""
Sample z_t from the probability path p_t(z_t | z0, z1).
Args:
z0: Source samples (noise), [B, C, H, W]
z1: Target samples (data), [B, C, H, W]
t: Timesteps, [B] or [B, 1, 1, 1]
Returns:
Interpolated samples z_t, [B, C, H, W]
"""
pass
@abstractmethod
def compute_loss(
self,
model_output: torch.Tensor,
z0: torch.Tensor,
z1: torch.Tensor,
t: torch.Tensor,
sample_weights: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Compute the loss for the algorithm.
Args:
model_output: Model predictions, [B, C, H, W]
z0: Source samples (noise), [B, C, H, W]
z1: Target samples (data), [B, C, H, W]
t: Timesteps, [B] or [B, 1, 1, 1]
sample_weights: Optional per-sample multipliers, one per row of the
batch. With them the returned scalar is the weighted mean of
the per-sample losses instead of their plain mean, which is how
a run raises the share of its gradient that long captions
receive; without them the computation is unchanged.
Returns:
Scalar loss value
"""
pass
class FlowMatchingOT(BaseAlgorithm):
"""
Flow matching with Optimal Transport path.
"""
def sample_zt(
self, z0: torch.Tensor, z1: torch.Tensor, t: torch.Tensor
) -> torch.Tensor:
if t.dim() == 1:
t = t.view(-1, 1, 1, 1)
z_t = (1.0 - t) * z0 + t * z1
return z_t
def compute_loss(
self,
model_output: torch.Tensor,
z0: torch.Tensor,
z1: torch.Tensor,
t: torch.Tensor,
sample_weights: torch.Tensor | None = None,
) -> torch.Tensor:
velocity_target = z1 - z0
if sample_weights is not None:
return sample_weighted_mse(model_output - velocity_target, sample_weights)
return F.mse_loss(model_output, velocity_target)
|