File size: 1,746 Bytes
617d016 | 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 | """HypDAE-style feature-space diffusion decoder (auxiliary loss only)."""
from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class HyperbolicDiffusionDecoder(nn.Module):
"""
Small DDPM-style denoiser in Euclidean feature space, conditioned on the
hierarchy code. Trained to reconstruct the CLS visual feature; not used
at inference time.
"""
def __init__(self, dim: int, code_dim: int, steps: int = 6, hidden: int = 512):
super().__init__()
self.dim = dim
self.steps = steps
betas = torch.linspace(1e-4, 2e-2, steps)
alphas = 1.0 - betas
alpha_bar = torch.cumprod(alphas, dim=0)
self.register_buffer("alpha_bar", alpha_bar)
t_dim = hidden // 4
self.time_mlp = nn.Sequential(
nn.Linear(1, t_dim),
nn.SiLU(),
nn.Linear(t_dim, t_dim),
)
self.eps_net = nn.Sequential(
nn.Linear(dim + code_dim + t_dim, hidden),
nn.SiLU(),
nn.Linear(hidden, hidden),
nn.SiLU(),
nn.Linear(hidden, dim),
)
def forward(self, hierarchy_code: torch.Tensor, cls_tok: torch.Tensor) -> torch.Tensor:
b = cls_tok.shape[0]
device = cls_tok.device
t = torch.randint(0, self.steps, (b,), device=device)
eps = torch.randn_like(cls_tok)
ab = self.alpha_bar[t].view(b, 1)
x_t = ab.sqrt() * cls_tok + (1.0 - ab).sqrt() * eps
t_emb = self.time_mlp((t.float() / max(self.steps - 1, 1)).unsqueeze(-1))
inp = torch.cat([x_t, hierarchy_code, t_emb], dim=-1)
eps_hat = self.eps_net(inp)
return F.mse_loss(eps_hat, eps)
|