bridge.models.diffusion_gemma.loss#

Selectable uniform-state diffusion SFT losses for DiffusionGemma (PyTorch).

This objective is a derivative work of the Apache-2.0 licensed DiffGemma SFT reference from Rhine-AI-Lab/diffusion-gemma-sft. The formulas are reimplemented here against Megatron’s loss contract; no reference repository code is imported at runtime.

Port of the JAX/kauldron objectives in gemma/diffusion/hackable_diffusion_adapter/hd/sft_model.py. All variants are a masked denoising cross-entropy on the completion tokens, differing only in a scalar time-weight w(t) and in which distribution the CE scores:

L_SFT = -E_t[ w(t) * sum_l M_l * log <x0_l, p_theta_l(x_t, t)> ]

variant             w(t)                prediction p_theta
------------------  ------------------  ------------------------------------------
base-sft            1                   softmax(f_theta)            (raw denoiser)
reweighted-ce       -alpha'(t)/alpha(t) softmax(f_theta)            (raw denoiser)
loo-ce              1                   softmax(f_theta + delta*xt) (LOO->denoiser)
reweighted-loo-ce   -alpha'(t)/alpha(t) softmax(f_theta + delta*xt) (LOO->denoiser)

where delta = log(1 + K*alpha/(1-alpha)) is the leave-one-out -> denoiser logit shift (Gourevitch et al., arXiv:2605.22765), K the vocab size, and the rectified-flow schedule gives alpha(t) = 1 - t, alpha'(t) = -1 so reweighted-ce’s weight is 1/(1-t).

Note: reweighted-ce uses -alpha'/alpha with no 1/K factor. The 1/K only belongs to the exact ELBO’s coeff = alpha'/(K*alpha) where it cancels in the term structure; in a plain weighted CE it just shrinks the loss by K (the original bug that drove eval to ~0.4%).

Module Contents#

Functions#

normalize_variant

Return the canonical objective name, accepting the early branch alias.

rf_alpha

Rectified-flow signal-retention schedule alpha(t) = 1 - t.

rf_alpha_dot

Time derivative d alpha / dt = -1 for the RF schedule.

loo_to_denoiser_logits

Map LOO logits -> standard denoiser logits (the +delta observed-token shift).

reweight

Per-example time weight w(t) for the given variant. Shape matches t ([B]).

diffusion_sft_loss

Masked, time-weighted denoising cross-entropy. Returns per-example loss [B].

diffusion_sft_logps

Per-example summed denoise log-prob surrogate -( w(t) * sum_l M_l CE_l ).

Data#

VARIANTS

_ALIASES

_LOO_VARIANTS

_REWEIGHTED_VARIANTS

_EPS

API#

bridge.models.diffusion_gemma.loss.VARIANTS#

(‘base-sft’, ‘reweighted-ce’, ‘loo-ce’, ‘reweighted-loo-ce’)

bridge.models.diffusion_gemma.loss._ALIASES#

None

bridge.models.diffusion_gemma.loss._LOO_VARIANTS#

(‘loo-ce’, ‘reweighted-loo-ce’)

bridge.models.diffusion_gemma.loss._REWEIGHTED_VARIANTS#

(‘reweighted-ce’, ‘reweighted-loo-ce’)

bridge.models.diffusion_gemma.loss._EPS#

1e-12

bridge.models.diffusion_gemma.loss.normalize_variant(variant: str) → str#

Return the canonical objective name, accepting the early branch alias.

bridge.models.diffusion_gemma.loss.rf_alpha(t: torch.Tensor) → torch.Tensor#

Rectified-flow signal-retention schedule alpha(t) = 1 - t.

bridge.models.diffusion_gemma.loss.rf_alpha_dot(t: torch.Tensor) → torch.Tensor#

Time derivative d alpha / dt = -1 for the RF schedule.

bridge.models.diffusion_gemma.loss.loo_to_denoiser_logits(
logits: torch.Tensor,
xt: torch.Tensor,
alpha: torch.Tensor,
vocab_size: int,
) → torch.Tensor#

Map LOO logits -> standard denoiser logits (the +delta observed-token shift).

delta = log(1 + K*alpha/(1-alpha)) added to the logit of the observed noisy token xt at each position. Inverse of the denoiser->LOO map.

Parameters:

  • logits – [B, L, K] raw (LOO) logits.

  • xt – [B, L] observed noisy token ids.

  • alpha – [B, 1] or [B, L] signal-retention probability.

  • vocab_size – K.

bridge.models.diffusion_gemma.loss.reweight(
variant: str,
t: torch.Tensor,
alpha: torch.Tensor,
alpha_dot: torch.Tensor,
) → torch.Tensor#

Per-example time weight w(t) for the given variant. Shape matches t ([B]).

bridge.models.diffusion_gemma.loss.diffusion_sft_loss(
logits: torch.Tensor,
x0: torch.Tensor,
xt: torch.Tensor,
t: torch.Tensor,
loss_mask: torch.Tensor,
*,
variant: str,
vocab_size: int,
diffusion_weight_clip: float | None = None,
alpha_fn=rf_alpha,
alpha_dot_fn=rf_alpha_dot,
) → torch.Tensor#

Masked, time-weighted denoising cross-entropy. Returns per-example loss [B].

Parameters:

  • logits – [B, L, K] network logits over the canvas.

  • x0 – [B, L] clean (target) completion tokens.

  • xt – [B, L] noisy canvas tokens.

  • t – [B] per-sequence diffusion time.

  • loss_mask – [B, L] 1.0 on completion tokens to score, 0.0 elsewhere.

  • variant – one of VARIANTS.

  • vocab_size – K.

  • diffusion_weight_clip – optional upper bound for the scalar time weight. This is mainly useful for reweighted-ce, whose RF weight is 1 / (1 - t) and can explode near t=1.

  • alpha_dot_fn (alpha_fn /) – schedule callables of t.

bridge.models.diffusion_gemma.loss.diffusion_sft_logps(
logits: torch.Tensor,
x0: torch.Tensor,
xt: torch.Tensor,
t: torch.Tensor,
loss_mask: torch.Tensor,
*,
variant: str,
vocab_size: int,
diffusion_weight_clip: float | None = None,
alpha_fn=rf_alpha,
alpha_dot_fn=rf_alpha_dot,
) → torch.Tensor#

Per-example summed denoise log-prob surrogate -( w(t) * sum_l M_l CE_l ).

Same masked, time-weighted denoising cross-entropy as :func:diffusion_sft_loss, but returned in summed (un-normalized) form as a log-probability (sign flipped). This is the convention the GRPO/GDSD objective expects: it divides the per-sequence logp by norm = number of scored tokens, so

diffusion_sft_logps(...) / n == -diffusion_sft_loss(...)

when norm equals the SFT mask count n = loss_mask.sum(-1). Args match

Func:

diffusion_sft_loss. Returns [B].