File size: 3,675 Bytes
8d31176 | 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 | from __future__ import annotations
from dataclasses import dataclass
import torch
from torch import Tensor, nn
import torch.nn.functional as F
from .model import SolPixOutput
@dataclass(frozen=True)
class FlowMatchingConfig:
"""Logit-normal timestep sampling used for the rectified-flow objective."""
logit_mean: float = 0.0
logit_std: float = 1.0
class REPAHead(nn.Module):
"""Optional train-only projection from SolPix features to teacher features."""
def __init__(self, model_width: int = 512, teacher_width: int = 768):
super().__init__()
self.projection = nn.Sequential(
nn.LayerNorm(model_width),
nn.Linear(model_width, teacher_width),
)
def forward(self, features: Tensor) -> Tensor:
return F.normalize(self.projection(features), dim=-1)
def sample_logit_normal_time(
batch_size: int,
device: torch.device,
config: FlowMatchingConfig | None = None,
) -> Tensor:
config = config or FlowMatchingConfig()
logits = torch.randn(batch_size, device=device) * config.logit_std + config.logit_mean
return logits.sigmoid()
def flow_matching_loss(
model: nn.Module,
clean_latents: Tensor,
text_embeddings: Tensor,
text_mask: Tensor | None = None,
*,
noise: Tensor | None = None,
time: Tensor | None = None,
config: FlowMatchingConfig | None = None,
) -> tuple[Tensor, Tensor, Tensor]:
"""Compute rectified-flow MSE for x_t=(1-t)x0+t*noise, v=noise-x0.
Returns ``(loss, predicted_velocity, time)`` so callers can log timesteps and
combine this term with auxiliary objectives such as REPA.
"""
if clean_latents.ndim != 4:
raise ValueError("clean_latents must have shape [batch, channels, height, width]")
batch = clean_latents.shape[0]
noise = torch.randn_like(clean_latents) if noise is None else noise
if noise.shape != clean_latents.shape:
raise ValueError("noise must have the same shape as clean_latents")
time = sample_logit_normal_time(batch, clean_latents.device, config) if time is None else time
if time.ndim == 0:
time = time.expand(batch)
if time.shape != (batch,):
raise ValueError("time must be a scalar or a [batch] tensor")
broadcast_time = time.to(clean_latents.dtype).reshape(batch, 1, 1, 1)
noisy_latents = (1 - broadcast_time) * clean_latents + broadcast_time * noise
target_velocity = noise - clean_latents
predicted_velocity = model(noisy_latents, time, text_embeddings, text_mask)
loss = F.mse_loss(predicted_velocity.float(), target_velocity.float())
return loss, predicted_velocity, time
def representation_alignment_loss(
output: SolPixOutput,
teacher_features: Tensor,
head: REPAHead,
token_mask: Tensor | None = None,
) -> Tensor:
"""Cosine alignment loss for clean-image teacher features (REPA-style)."""
if output.repa_features is None:
raise ValueError("call SolPix(..., return_features=True) to obtain REPA features")
prediction = head(output.repa_features)
target = F.normalize(teacher_features.detach(), dim=-1)
if prediction.shape != target.shape:
raise ValueError(
f"projected model features {tuple(prediction.shape)} do not match teacher features {tuple(target.shape)}"
)
per_token = 1 - (prediction * target).sum(dim=-1)
if token_mask is None:
return per_token.mean()
if token_mask.shape != per_token.shape:
raise ValueError("token_mask must have shape [batch, image_tokens]")
weights = token_mask.to(per_token.dtype)
return (per_token * weights).sum() / weights.sum().clamp_min(1)
|