ParkSeongHyunandClaude Opus 4.6
ca612bfff0
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
65 lines
2.4 KiB
Python
65 lines
2.4 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def compute_repulsion_gradient(embeddings, eps_noise=1e-4):
|
|
"""Compute the normalized gradient of the Vendi Score diversity loss.
|
|
|
|
Based on "On-the-fly Repulsion in the Contextual Space for Rich Diversity
|
|
in Diffusion Transformers" (arXiv:2603.28762).
|
|
|
|
The gradient is normalized per-sample so that repulsion_scale has a
|
|
consistent effect regardless of embedding dimensionality.
|
|
|
|
Args:
|
|
embeddings: Tensor of shape [B, N, D] — batch of token sequences.
|
|
eps_noise: Small noise added to break symmetry when embeddings are identical.
|
|
|
|
Returns:
|
|
Normalized gradient tensor of shape [B, N, D]. Subtract (scale * grad)
|
|
from embeddings to increase diversity.
|
|
"""
|
|
B, N, D = embeddings.shape
|
|
if B < 2:
|
|
return torch.zeros_like(embeddings)
|
|
|
|
with torch.inference_mode(False):
|
|
with torch.enable_grad():
|
|
# Clone into float32 with grad tracking
|
|
emb = embeddings.detach().float().clone().requires_grad_(True)
|
|
|
|
# Break symmetry for initially identical embeddings
|
|
noise = torch.randn_like(emb) * eps_noise
|
|
emb_noisy = emb + noise
|
|
|
|
# Flatten: [B, N, D] -> [B, N*D]
|
|
flat = emb_noisy.reshape(B, -1)
|
|
|
|
# Cosine similarity kernel: [B, B]
|
|
flat_norm = F.normalize(flat, dim=-1)
|
|
K = flat_norm @ flat_norm.T
|
|
|
|
# Normalized kernel
|
|
K_tilde = K / B
|
|
|
|
# Eigenvalues (K_tilde is real symmetric positive semi-definite)
|
|
eigenvalues = torch.linalg.eigvalsh(K_tilde)
|
|
|
|
# Von Neumann entropy loss (negative entropy — minimize to maximize diversity)
|
|
eigenvalues_clamped = eigenvalues.clamp(min=1e-8)
|
|
loss = -(eigenvalues_clamped * eigenvalues_clamped.log()).sum()
|
|
|
|
loss.backward()
|
|
|
|
grad = emb.grad # [B, N, D]
|
|
|
|
# Normalize: scale each sample's gradient to have unit norm relative
|
|
# to the embedding norm, so repulsion_scale is interpretable as a
|
|
# fraction of embedding magnitude.
|
|
emb_norms = embeddings.float().reshape(B, -1).norm(dim=-1, keepdim=True).clamp(min=1e-8) # [B, 1]
|
|
grad_norms = grad.reshape(B, -1).norm(dim=-1, keepdim=True).clamp(min=1e-8) # [B, 1]
|
|
grad = grad.reshape(B, -1) * (emb_norms / grad_norms)
|
|
grad = grad.reshape(B, N, D)
|
|
|
|
return grad.to(embeddings.dtype)
|