""" Flow-CoPD loss: Contrastive On-Policy Distillation for flow-matching T2I. See refine-logs/FINAL_PROPOSAL.md. One-line method: Flow-OPD / DiffusionOPD are POSITIVE-ONLY (pull student velocity toward a teacher velocity). Flow-CoPD adds the missing CONTRASTIVE NEGATIVE: on the student's own LOW-reward on-policy trajectories, repel the student velocity from the (constant rectified-flow) velocity that regenerates those low-reward samples. Asymmetric, data-source dependent (DistiLLM-2 port). All functions are pure tensor ops so they can be unit-tested / reviewed in isolation. Shapes: v_* are (B, C, H, W) velocity-field predictions at one denoising step; weight_t is a scalar or (B,1,1,1) per-step weight; advantages is (B,) per-trajectory (already group-normalized, as in flow_grpo). DESIGN NOTES / open questions flagged for cross-model review: [R1] Positive term = reverse-KL -> time-weighted velocity-L2 (DiffusionOPD ODE form). This is NOT advantage-weighted (advantage-weighting the L2 would collapse to plain loss-weighting; see TAD-OPD post-mortem). [R2] Negative term is SIGN-FLIPPED (repulsion), so it is NOT expressible as a positive reweighting of the positive L2 -> genuinely new signal, not a collapse. It is CLAMPED (neg_clamp) so the unbounded -||.||^2 cannot diverge. [R3] v_neg_target default = constant rectified-flow velocity (noise - x0) that regenerates the self-generated low-reward sample. Scheduler-convention sensitive -> verify against SD3 FlowMatchEulerDiscreteScheduler. """ import torch def _mse(a, b): # per-sample mean squared error over all non-batch dims -> (B,) return ((a.float() - b.float()) ** 2).flatten(1).mean(dim=1) def opd_positive_loss(v_student, v_teacher, weight_t=1.0): """A0 = positive-only OPD (Flow-OPD / DiffusionOPD reproduction). L_pos = E[ weight_t * || v_student - v_teacher ||^2 ]. v_teacher must be detached (computed under the frozen teacher adapter / no_grad). """ per = _mse(v_student, v_teacher.detach()) # (B,) if torch.is_tensor(weight_t): per = weight_t.flatten() * per else: per = weight_t * per return per.mean() def copd_loss( v_student, v_teacher, v_neg_target, advantages, weight_t=1.0, lambda_neg=1.0, neg_clamp=1.0, adv_thresh=0.0, detach_neg_target=True, ): """A2 = Flow-CoPD = positive OPD + contrastive negative. Args: v_student: (B,C,H,W) trainable student velocity at this step. v_teacher: (B,C,H,W) frozen teacher velocity at this step (detached). v_neg_target: (B,C,H,W) velocity that REGENERATES the self-generated sample (constant rectified-flow target); we repel from it on low-reward trajectories. Detached by default. advantages: (B,) group-normalized per-trajectory advantage (reward signal). >adv_thresh = "good", <-adv_thresh treated as "bad". weight_t: per-step time weight (scalar or (B,1,1,1)). lambda_neg: strength of the contrastive negative term. neg_clamp: cap on the per-sample negative MSE (stability; prevents the repulsive -||.||^2 from diverging). adv_thresh: deadband; |adv|<=thresh contributes only to the positive term. Returns: (total_loss, dict_of_components_for_logging) """ if detach_neg_target: v_neg_target = v_neg_target.detach() w = weight_t.flatten() if torch.is_tensor(weight_t) else weight_t # ---- positive term: pull every trajectory's velocity toward the teacher ---- pos_per = _mse(v_student, v_teacher.detach()) # (B,) pos_loss = (w * pos_per).mean() if torch.is_tensor(w) else (w * pos_per).mean() # ---- contrastive negative: repel ONLY low-reward trajectories from the # velocity that regenerates them (clamped, sign-flipped) ---- bad_mask = (advantages < -adv_thresh).float() # (B,) neg_per = _mse(v_student, v_neg_target) # (B,) neg_per = torch.clamp(neg_per, max=neg_clamp) # bound it # maximize distance on bad samples => subtract; weight by |advantage| so the # worst samples are repelled hardest (sign already handled by bad_mask). neg_weight = bad_mask * advantages.abs() if torch.is_tensor(w): neg_term = (w * neg_weight * neg_per).mean() else: neg_term = (w * neg_weight * neg_per).mean() neg_loss = -lambda_neg * neg_term total = pos_loss + neg_loss info = { "loss_pos": pos_loss.detach(), "loss_neg": neg_loss.detach(), "neg_frac": bad_mask.mean().detach(), "neg_mse_mean": neg_per.mean().detach(), } return total, info def regen_velocity_per_step(x_t, x0, sigma_t, eps=1e-3): """Per-step model-output velocity that regenerates endpoint x0 from state x_t. SD3 FlowMatchEuler: x0 = x_t - sigma_t * v_model => v_model = (x_t - x0)/sigma_t. So the velocity the student USED (in expectation) to drive x_t toward the (low-reward) endpoint x0 at this step is (x_t - x0)/sigma_t. Repelling the student from THIS per-step target unlearns the path to bad samples. Fixes CRITICAL#1 from review: the old constant (noise - x0) target is only valid for a straight deterministic path, NOT the Flow-GRPO SDE rollout. Args: x_t: (B,C,H,W) state before this step (sample["latents"][:, j]). x0: (B,C,H,W) low-reward endpoint (sample["next_latents"][:, -1]). sigma_t: scalar or (B,) noise level at this step (from scheduler.sigmas[j]). """ if not torch.is_tensor(sigma_t): sigma_t = torch.as_tensor(sigma_t, dtype=x_t.dtype, device=x_t.device) sigma_t = sigma_t.clamp(min=eps).reshape(-1, 1, 1, 1) return ((x_t - x0) / sigma_t).detach() def rectified_flow_regen_velocity(latents_first, latents_last): """DEPRECATED (review CRITICAL#1): constant noise->x0 velocity. Wrong for SDE rollouts; kept only so old configs import-error loudly. Use regen_velocity_per_step.""" raise DeprecationWarning("use regen_velocity_per_step(x_t, x0, sigma_t)") def sde_time_weight(sigma_t, scheme="snr", eps=1e-3): """Per-step OPD weight (review MAJOR#5). 'uniform'->1; 'snr'-> 1/sigma^2-ish so low-noise (late) steps that decide detail are weighted up. Returns (B,1,1,1) or scalar.""" if scheme == "uniform": return 1.0 if not torch.is_tensor(sigma_t): sigma_t = torch.as_tensor(sigma_t) sigma_t = sigma_t.clamp(min=eps) w = 1.0 / (sigma_t ** 2 + 1.0) # bounded SNR-like weight in (0,1] return w.reshape(-1, 1, 1, 1)