From 831ccfa7aa2e116d81848011dcb6a32791521737 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 15 Oct 2026 08:26:12 -0600 Subject: [PATCH] Added experimental band loss --- extensions_built_in/sd_trainer/SDTrainer.py | 5 +- toolkit/config_modules.py | 2 +- toolkit/util/band_loss.svg | 94 +++++++++++++++++++++ toolkit/util/losses.py | 61 +++++++++++++ 4 files changed, 160 insertions(+), 2 deletions(-) create mode 100644 toolkit/util/band_loss.svg diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index d464d3d..957918d 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -36,7 +36,7 @@ from diffusers import EMAModel import math from toolkit.train_tools import precondition_model_outputs_flow_match from toolkit.models.diffusion_feature_extraction import DiffusionFeatureExtractor, load_dfe -from toolkit.util.losses import wavelet_loss, stepped_loss +from toolkit.util.losses import wavelet_loss, stepped_loss, band_loss import torch.nn.functional as F from toolkit.unloader import unload_text_encoder from PIL import Image @@ -952,6 +952,9 @@ class SDTrainer(BaseSDTrainProcess): loss = torch.nn.functional.l1_loss(pred.float(), target.float(), reduction="none") elif self.train_config.loss_type == "wavelet": loss = wavelet_loss(pred, batch.latents, noise) + elif self.train_config.loss_type == "band": + # per-sample frequency-weighted MSE: bands the noise has buried at this t drop out + loss = band_loss(pred, target, batch.latents, noise, timesteps, self.sd.is_flow_matching, self.sd.noise_scheduler) elif self.train_config.loss_type == "stepped": loss = stepped_loss(pred, batch.latents, noise, noisy_latents, timesteps, self.sd.noise_scheduler) # the way this loss works, it is low, increase it to match predictable LR effects diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 4f05d74..c45a3ab 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -522,7 +522,7 @@ class TrainConfig: self.correct_pred_norm = kwargs.get('correct_pred_norm', False) self.correct_pred_norm_multiplier = kwargs.get('correct_pred_norm_multiplier', 1.0) - self.loss_type = kwargs.get('loss_type', 'mse') # mse, mae, wavelet, pixelspace, mean_flow, pseudo_huber + self.loss_type = kwargs.get('loss_type', 'mse') # mse, mae, wavelet, band, pixelspace, mean_flow, pseudo_huber # do the loss on a timestep to 0 prediction self.t0_loss_target = kwargs.get('t0_loss_target', False) diff --git a/toolkit/util/band_loss.svg b/toolkit/util/band_loss.svg new file mode 100644 index 0000000..f71ce95 --- /dev/null +++ b/toolkit/util/band_loss.svg @@ -0,0 +1,94 @@ + + + + + + + + + + +band loss: frequency-weighted residual (toolkit/util/losses.py: band_loss) +The plain MSE penalizes detail the noise has already destroyed. This loss keeps each frequency band of the residual +only as far as that band is still recoverable at the sample's noise level t. + + + +1. residual +pred - target +x0, velocity or epsilon: +one picture scaled by t, +no conversion needed + + + + +2. split into Laplacian bands + +fine +mid +coarse +low (8x8 pool) +band_k = low_k - low_{k+1} +low_k = residual pooled over 2^k x 2^k cells, +resampled to full size; bands sum to the residual +pred and target split separately (linear) +video: per frame over H, W + + + + +3. weight each band by its recoverability at t +flow: snr(t) = ((1-t)/t)^2, t = steps/1000 +ddpm: snr(t) = alpha_bar / (1 - alpha_bar) +snr_k = snr(t) * P_clean[k] / P_noise[k] +w_k = snr_k / (1 + snr_k) +P_clean[k], P_noise[k]: mean squared band k of the +batch's clean latents and of its noise, measured per +batch: no 1/f prior, any VAE, pyramid noise just works. +t -> 0: every w_k = 1, the loss IS the plain MSE. +t -> 1: only the low band keeps a weight. + + + + +4. filtered residual, squared +filtered = sum_k w_k * band_k +loss = filtered ** 2 +elementwise, pred's shape: masks, timestep weights, +LossScaler and prior loss apply downstream unchanged +config: train.loss_type = "band" + + +w_k(t) for typical latent band powers + + + + + +fine band +mid band +coarse band +low band +1 +0 +t = 0 (clean) +t = 1 (noise) +fine detail drops out first, global layout last + + + +why +Images carry most power at low frequencies and white noise is flat, so at noise level t each band has its own +signal-to-noise ratio and the fine ones go under first. Above that cutoff the best possible estimate is a blur, and an +MSE gradient on those bands is a coin flip. Dropping them removes gradient noise at high t; at low t nothing changes. +The depth-wise code heads give coarse-to-fine supervision across layers; this gives it across timesteps. They compose. + diff --git a/toolkit/util/losses.py b/toolkit/util/losses.py index 5328674..0b02d17 100644 --- a/toolkit/util/losses.py +++ b/toolkit/util/losses.py @@ -1,4 +1,5 @@ import torch +import torch.nn.functional as F _dwt = None @@ -91,3 +92,63 @@ def stepped_loss(model_pred, latents, noise, noisy_latents, timesteps, scheduler latents.float().to(device=predicted_images.device), reduction="none", ) + + +def _band_lows(x: torch.Tensor, levels: int) -> list: + """Gaussian-pyramid low-passes of (N, C, H, W) at the input size: lows[0] = x, lows[k] = x + averaged over 2^k x 2^k cells and bilinearly resampled back.""" + lows = [x] + cur = x + for _ in range(levels): + cur = F.avg_pool2d(cur, 2, ceil_mode=True) + lows.append(F.interpolate(cur, size=x.shape[-2:], mode="bilinear", align_corners=False)) + return lows + + +def _bands(x: torch.Tensor, levels: int) -> list: + """Laplacian bands at the input size, fine to coarse; they sum back to x.""" + lows = _band_lows(x, levels) + return [lows[k] - lows[k + 1] for k in range(levels)] + [lows[levels]] + + +def band_loss(pred, target, latents, noise, timesteps, is_flow_matching, scheduler, levels: int = 3): + """MSE on a per-sample frequency-filtered residual: each Laplacian band of the residual is + weighted by how recoverable that band is at the sample's noise level, snr_k / (1 + snr_k) with + snr_k = snr(t) * band power of the clean latents / band power of the noise (both measured on + the batch). At t -> 0 every weight is 1 and this is the plain MSE; at high t the bands the + noise has buried drop out instead of adding gradient noise. The filter is linear, so it + applies to pred / target as they are (x0, velocity or epsilon residuals differ only by a + scalar in t). Returns an elementwise loss, pred's shape.""" + pred, target = pred.float(), target.float() + shape = pred.shape + if pred.dim() not in (4, 5): + return F.mse_loss(pred, target, reduction="none") + n = shape[0] + if pred.dim() == 5: # video: bands over (H, W) of every frame + flat = lambda x: x.float().transpose(1, 2).flatten(0, 1) + else: + flat = lambda x: x.float() + with torch.no_grad(): + clean_b = _bands(flat(latents.to(pred.device)), levels) + noise_b = _bands(flat(noise.to(pred.device)), levels) + # band power over the whole batch: a prior-level quantity, not per image + p_clean = torch.stack([b.pow(2).mean() for b in clean_b]) + p_noise = torch.stack([b.pow(2).mean() for b in noise_b]).clamp_min(1e-12) + ts = timesteps.to(pred.device).float().flatten() + if is_flow_matching: + t = (ts / 1000.0).clamp(1e-4, 1 - 1e-4) + snr = ((1 - t) / t) ** 2 + else: + ac = scheduler.alphas_cumprod.to(pred.device, torch.float32)[ts.long().clamp(0, scheduler.alphas_cumprod.shape[0] - 1)] + snr = ac / (1 - ac).clamp_min(1e-8) + s = snr.view(n, 1) * (p_clean / p_noise).view(1, -1) # (N, levels + 1) + w = s / (1 + s) + if pred.dim() == 5: + w = w.repeat_interleave(shape[2], dim=0) # per frame row + w = w.t().reshape(levels + 1, -1, 1, 1, 1) + pb, tb = _bands(flat(pred), levels), _bands(flat(target), levels) + filtered = sum(w[k] * (pb[k] - tb[k]) for k in range(levels + 1)) + loss = filtered.pow(2) + if pred.dim() == 5: + loss = loss.view(n, shape[2], shape[1], *shape[-2:]).transpose(1, 2) + return loss