File size: 5,167 Bytes
cfaf85b | 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 | """Two-trace quadratic objective from idea.pdf, sections 10 and 14."""
import torch
import torch.nn.functional as F
class TraceIMF:
def __init__(self, channels=3, image_size=32, num_classes=10, lambda_diag=1.0):
if lambda_diag <= 0:
raise ValueError("lambda_diag must be positive to train both traces")
self.channels = channels
self.image_size = image_size
self.num_classes = num_classes
self.lambda_diag = lambda_diag
def loss(self, model, x, labels=None, *, derivative_model=None, t=None, noise=None,
r=None, generator=None):
"""x is raw image data in [0,1]. Optional t/noise permit exact checks.
Use the unwrapped model for JVP under Accelerate/DDP. The single joint
grad-enabled call through model trains both edge and diagonal safely.
"""
b = x.shape[0]
x = x.float() * 2 - 1
t = torch.rand(b, device=x.device, generator=generator) if t is None else t.to(x.device).float()
noise = torch.randn(x.shape, device=x.device, generator=generator) if noise is None else noise.to(x.device).float()
if t.shape != (b,) or noise.shape != x.shape:
raise ValueError("Expected t with shape (B,) and noise with the image shape")
if self.num_classes is not None and labels is None:
raise ValueError("Class-conditioned training requires labels")
if self.num_classes is None:
labels = None
t_image = t[:, None, None, None]
z = (1 - t_image) * x + t_image * noise
target = noise - x
r = torch.zeros_like(t) if r is None else r.to(x.device).float()
if r.shape != t.shape or (r < 0).any() or (t > 1).any() or (r > t).any():
raise ValueError("Require 0 <= r <= t <= 1, with r and t shaped (B,)")
if getattr(self, "channels_last", False):
z = z.contiguous(memory_format=torch.channels_last)
# Both branches use the SAME u head, evaluated at r=0 and r=t.
pair_labels = None if labels is None else torch.cat((labels, labels))
predictions = model(torch.cat((z, z)), torch.cat((t, t)),
torch.cat((r, t)), y=pair_labels,
use_flash_attention=True)
edge, diagonal = predictions.chunk(2)
jvp_model = model if derivative_model is None else derivative_model
# FP32 JVP: native attention supports forward AD; SDPA is used only in
# the grad-enabled forward. No parameter gradient through the JVP.
use_bf16 = getattr(self, "jvp_precision", "fp32") == "bf16" and x.device.type == "cuda"
with torch.no_grad(), torch.autocast(device_type=x.device.type, enabled=use_bf16,
dtype=torch.bfloat16):
def edge_fn(z_value, t_value, r_value):
# Accelerate may autocast the instance's forward even after
# unwrap_model. The class forward bypasses that AMP wrapper
# without mutating the grad-enabled training forward.
return type(jvp_model).forward(jvp_model, z_value, t_value, r_value,
y=labels, use_flash_attention=False)
_, derivative = torch.func.jvp(
edge_fn, (z, t, r),
(diagonal.detach().float(), torch.ones_like(t), torch.zeros_like(r)),
)
compound = edge.float() + (t - r)[:, None, None, None] * derivative.detach().float()
edge_loss = F.mse_loss(compound, target)
diag_loss = F.mse_loss(diagonal.float(), target)
loss = edge_loss + self.lambda_diag * diag_loss
return loss, {"edge_loss": edge_loss.detach(), "diag_loss": diag_loss.detach(),
"jvp_rms": derivative.detach().square().mean().sqrt()}
@torch.no_grad()
def sample(self, model, n_samples=None, labels=None, *, device=None, noise=None):
"""Exactly one evaluation at (t,r)=(1,0), then map back to [0,1]."""
was_training = model.training
model.eval()
try:
if device is None:
device = next(model.parameters()).device
if labels is not None:
labels = labels.to(device)
n_samples = len(labels)
if self.num_classes is not None and labels is None:
raise ValueError("Class-conditioned sampling requires labels")
if self.num_classes is None:
labels = None
if noise is None:
if n_samples is None or n_samples <= 0:
raise ValueError("n_samples must be positive")
noise = torch.randn(n_samples, self.channels, self.image_size,
self.image_size, device=device)
else:
noise = noise.to(device)
t = torch.ones(noise.shape[0], device=device)
u = model(noise, t, torch.zeros_like(t), y=labels,
use_flash_attention=True)
return ((noise - u).clamp(-1, 1) + 1) * 0.5
finally:
model.train(was_training)
|