File size: 8,936 Bytes
141bacd | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 | """panda encoder with two input variants (pca / marker), sub-center prototypes."""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import Optional, List
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function
# gradient reversal layer
class GradReverse(Function):
@staticmethod
def forward(ctx, x, lam):
ctx.lam = lam
return x.view_as(x)
@staticmethod
def backward(ctx, g):
return -ctx.lam * g, None
def grad_reverse(x, lam):
return GradReverse.apply(x, lam)
# encoder with variant + sub-centers
class PANDAEncoder(nn.Module):
"""trunk + projection head + sub-center prototypes.
args:
variant : "pca" or "marker"
n_pca : 50
n_markers : m >= 0. required > 0 if variant == "marker".
n_classes : K
n_sub : sub-centers per class (default 3)
d_hidden, d_repr, d_proj: trunk sizing
n_datasets : for the dataset adversary head
"""
def __init__(
self,
variant: str = "pca",
n_pca: int = 50,
n_markers: int = 0,
d_hidden: int = 512,
d_repr: int = 256,
d_proj: int = 128,
n_classes: int = 10,
n_sub: int = 3,
n_datasets: int = 1,
dropout: float = 0.2,
):
super().__init__()
assert variant in ("pca", "marker"), variant
if variant == "marker":
assert n_markers > 0, "PANDA-Marker requires n_markers>0"
self.variant = variant
self.n_pca = n_pca
self.n_markers = n_markers if variant == "marker" else 0
self.n_classes = n_classes
self.n_sub = n_sub
self.n_datasets = n_datasets
input_dim = n_pca + self.n_markers
self.input_dim = input_dim
self.trunk = nn.Sequential(
nn.Linear(input_dim, d_hidden), nn.LayerNorm(d_hidden), nn.GELU(), nn.Dropout(dropout),
nn.Linear(d_hidden, d_hidden), nn.LayerNorm(d_hidden), nn.GELU(), nn.Dropout(dropout),
nn.Linear(d_hidden, d_repr), nn.LayerNorm(d_repr), nn.GELU(),
)
self.projection = nn.Sequential(
nn.Linear(d_repr, d_repr), nn.GELU(),
nn.Linear(d_repr, d_proj),
)
self.classifier = nn.Sequential(nn.Linear(d_repr + 2, n_classes))
self.dom_adv = nn.Sequential(nn.Linear(d_repr, 128), nn.ReLU(), nn.Linear(128, n_datasets))
self.depth_adv = nn.Sequential(nn.Linear(d_repr, 64), nn.ReLU(), nn.Linear(64, 1))
# sub-center prototypes (K, n_sub, d_proj), L2-normalised per sub-center
self.register_buffer(
"prototypes",
F.normalize(torch.randn(n_classes, n_sub, d_proj), dim=-1),
)
# EMA momentum as a buffer so we can overwrite it in place
self.register_buffer("proto_ema", torch.tensor(0.99))
@torch.no_grad()
def update_prototypes(self, z_norm: torch.Tensor, y: torch.Tensor):
"""ema update: assign each in-class cell to nearest sub-center, take the mean."""
ema = float(self.proto_ema.item())
for c in torch.unique(y):
mask = y == c
if not mask.any():
continue
zc = z_norm[mask] # (n_c, d_proj)
protos_c = self.prototypes[c] # (n_sub, d_proj)
sims = zc @ protos_c.T # (n_c, n_sub)
assign = sims.argmax(dim=1) # each cell -> nearest sub-center
for k in range(self.n_sub):
m2 = assign == k
if not m2.any():
continue
new = F.normalize(zc[m2].mean(dim=0), dim=0)
self.prototypes[c, k] = F.normalize(
ema * self.prototypes[c, k] + (1 - ema) * new, dim=0
)
@torch.no_grad()
def max_sub_cos(self, z_norm: torch.Tensor) -> torch.Tensor:
"""(B, K) cos(z, best sub-center) per class."""
B = z_norm.size(0); K, n_sub, D = self.prototypes.shape
sims = torch.einsum("bd,ksd->bks", z_norm, self.prototypes) # (B, K, n_sub)
return sims.max(dim=2).values # (B, K)
def forward(
self,
x_pca: torch.Tensor,
aux: torch.Tensor,
x_markers: Optional[torch.Tensor] = None,
lam_dann: float = 0.0,
) -> dict:
if self.variant == "marker":
assert x_markers is not None and x_markers.size(1) == self.n_markers
x = torch.cat([x_pca, x_markers], dim=1)
else:
x = x_pca
h = self.trunk(x)
z_raw = self.projection(h)
z = F.normalize(z_raw, dim=1)
logits = self.classifier(torch.cat([h, aux], dim=1))
h_rev = grad_reverse(h, lam_dann)
return {
"repr": h,
"z": z,
"logits": logits,
"dom": self.dom_adv(h_rev),
"depth": self.depth_adv(h_rev),
}
# losses
def supcon_loss(z: torch.Tensor, y: torch.Tensor, temperature: float = 0.1) -> torch.Tensor:
if z.size(0) < 2:
return z.new_zeros(())
sim = z @ z.T / temperature
sim_max, _ = sim.max(dim=1, keepdim=True)
sim = sim - sim_max.detach()
logits_mask = torch.ones_like(sim) - torch.eye(z.size(0), device=z.device)
exp_sim = torch.exp(sim) * logits_mask
log_prob = sim - torch.log(exp_sim.sum(dim=1, keepdim=True) + 1e-12)
labels_eq = (y.unsqueeze(0) == y.unsqueeze(1)).float() * logits_mask
denom = labels_eq.sum(dim=1).clamp_min(1.0)
per = -(labels_eq * log_prob).sum(dim=1) / denom
per = per * (labels_eq.sum(dim=1) > 0).float()
counts = torch.bincount(y, minlength=int(y.max().item()) + 1).float().clamp_min(1.0)
w = 1.0 / counts.sqrt()
return (per * w[y]).sum() / w[y].sum().clamp_min(1e-6)
def vicreg_loss(z: torch.Tensor, sim_weight: float = 0.0, var_weight: float = 25.0,
cov_weight: float = 1.0) -> torch.Tensor:
zc = z - z.mean(dim=0, keepdim=True)
std = (zc.var(dim=0) + 1e-4).sqrt()
var_loss = F.relu(1.0 - std).mean()
N, D = zc.shape
cov = (zc.T @ zc) / (N - 1)
off = cov - torch.diag(torch.diagonal(cov))
cov_loss = off.pow(2).sum() / D
return var_weight * var_loss + cov_weight * cov_loss
def hsic_biased(x: torch.Tensor, y: torch.Tensor,
sigma_x: float = 1.0, sigma_y: float = 1.0) -> torch.Tensor:
Nx = x.size(0)
if Nx < 2:
return x.new_zeros(())
K = torch.exp(-torch.cdist(x, x) ** 2 / (2 * sigma_x ** 2))
L = torch.exp(-torch.cdist(y, y) ** 2 / (2 * sigma_y ** 2))
H = torch.eye(Nx, device=x.device) - torch.ones(Nx, Nx, device=x.device) / Nx
return (K @ H @ L @ H).trace() / (Nx - 1) ** 2
def subcenter_angular_infonce(
z: torch.Tensor, # (B, d_proj) L2-normalised
y: torch.Tensor, # (B,)
prototypes: torch.Tensor, # (K, n_sub, d_proj)
margin: float = 0.15, # angular margin in radians
temperature: float = 0.07,
) -> torch.Tensor:
"""arcface-style angular-margin loss over sub-center prototypes."""
B = z.size(0); K, n_sub, D = prototypes.shape
sims = torch.einsum("bd,ksd->bks", z, prototypes) # (B, K, n_sub)
max_over_sub = sims.max(dim=2).values # (B, K)
# target class cosine, bump by angular margin, put back
target_cos = max_over_sub.gather(1, y.unsqueeze(1)).squeeze(1) # (B,)
target_cos = target_cos.clamp(-1 + 1e-7, 1 - 1e-7)
theta = torch.acos(target_cos)
target_new_cos = torch.cos(theta + margin)
logits = max_over_sub.clone()
logits.scatter_(1, y.unsqueeze(1), target_new_cos.unsqueeze(1))
logits = logits / temperature
return F.cross_entropy(logits, y)
def prototype_repulsion(prototypes: torch.Tensor, weight: float = 1.0) -> torch.Tensor:
"""penalise inter-class prototype cosine so eff-dim doesn't collapse."""
K, n_sub, D = prototypes.shape
centroids = F.normalize(prototypes.mean(dim=1), dim=1) # (K, D)
sim = centroids @ centroids.T # (K, K)
off = sim - torch.diag(torch.diagonal(sim))
return weight * off.pow(2).sum() / (K * (K - 1) + 1e-6)
def prototype_infonce_legacy(z, y, prototypes, temperature=0.07):
"""legacy single-prototype InfoNCE. kept for debugging + old checkpoints."""
if prototypes.dim() == 3:
prototypes = F.normalize(prototypes.mean(dim=1), dim=1) # collapse sub-centers
logits = z @ prototypes.T / temperature
return F.cross_entropy(logits, y)
|