File size: 9,135 Bytes
0ab2566 | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
class DirectMatchSupervisor(nn.Module):
"""Direct soft matching supervisor without prototypes.
Ablation variant of PMPSLatentSupervisor: removes the prototype layer
and Sinkhorn-Knopp, directly matching latent and CoT embeddings
in the projected space via cosine similarity.
The matching structure (soft many-to-many, bidirectional, PSA-compatible)
is preserved. The loss changes from cross-entropy on prototype codes to
cosine similarity cross-reconstruction in embedding space.
Interface-compatible with PMPSLatentSupervisor for drop-in replacement.
"""
def __init__(
self,
hidden_dim,
feat_dim=128,
hidden_mlp=512,
match_temperature=0.3,
use_projector=True,
position_bias_strength=0.0,
):
super().__init__()
self.match_temperature = match_temperature
self.position_bias_strength = position_bias_strength
self.use_projector = use_projector
if use_projector:
self.projector = self._build_projector(hidden_dim, hidden_mlp, feat_dim)
self.feat_dim = feat_dim
else:
self.projector = nn.Identity()
self.feat_dim = hidden_dim
@staticmethod
def _build_projector(input_dim, hidden_dim, output_dim):
return nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(inplace=True),
nn.Linear(hidden_dim, output_dim),
)
def normalize_prototypes(self):
"""No-op: no prototypes in direct matching."""
pass
def _project(self, hidden_states):
"""Project (B, N, D) -> L2-normalised (B, N, feat_dim)."""
b, n, d = hidden_states.shape
flat = hidden_states.reshape(b * n, d)
emb = self.projector(flat)
emb = F.normalize(emb, dim=1, p=2)
return emb.reshape(b, n, -1)
@staticmethod
def _compute_position_bias(N_c, N_l, strength, cot_mask=None, device=None):
"""Gaussian positional bias (identical to PMPSLatentSupervisor)."""
if strength <= 0:
return 0.0
lat_pos = torch.arange(N_l, device=device, dtype=torch.float32)
lat_pos_norm = lat_pos / max(N_l - 1, 1)
if cot_mask is not None and cot_mask.dim() == 2:
B = cot_mask.size(0)
bias = torch.zeros(B, N_c, N_l, device=device)
for b in range(B):
n_valid = cot_mask[b].sum().long().item()
if n_valid <= 1:
continue
cot_pos = torch.arange(n_valid, device=device, dtype=torch.float32)
cot_pos_norm = cot_pos / max(n_valid - 1, 1)
delta = cot_pos_norm.unsqueeze(1) - lat_pos_norm.unsqueeze(0)
bias[b, :n_valid, :] = -strength * delta.pow(2)
return bias
else:
cot_pos = torch.arange(N_c, device=device, dtype=torch.float32)
cot_pos_norm = cot_pos / max(N_c - 1, 1)
delta = cot_pos_norm.unsqueeze(1) - lat_pos_norm.unsqueeze(0)
bias = -strength * delta.pow(2)
return bias.unsqueeze(0)
def forward(self, latent_hidden, cot_hidden, latent_mask=None, cot_mask=None,
position_bias_strength_override=None):
"""Compute direct soft matching loss between latent and CoT embeddings.
Same interface and return format as PMPSLatentSupervisor.forward().
Matching pipeline:
1. Project both sides through shared g(·) → L2-normalised z
2. Compute cosine similarity matrix z_cot · z_lat^T (stop-grad)
3. Apply PSA positional bias
4. Softmax → bidirectional matching weights A (no gradient)
5. Cross-reconstruct: z_hat = A · z (gradient through z)
6. Loss: 1 - cosine_similarity(z_target, z_hat)
Returns
-------
loss_c2l : scalar tensor (used for backprop, consistent with PMPS)
log_dict : diagnostic stats
"""
assert latent_hidden.dim() == 3 and cot_hidden.dim() == 3
assert latent_hidden.size(0) == cot_hidden.size(0)
assert latent_hidden.size(-1) == cot_hidden.size(-1)
B, N_l, _ = latent_hidden.shape
_, N_c, _ = cot_hidden.shape
device = latent_hidden.device
if latent_mask is None:
latent_mask = torch.ones(B, N_l, device=device, dtype=latent_hidden.dtype)
if cot_mask is None:
cot_mask = torch.ones(B, N_c, device=device, dtype=cot_hidden.dtype)
# 1. Project to shared space (gradient flows through projector)
z_lat = self._project(latent_hidden) # (B, N_l, d)
z_cot = self._project(cot_hidden) # (B, N_c, d)
# 2. Compute matching weights (stop-gradient, analogous to Sinkhorn codes)
with torch.no_grad():
match_logits_c2l = torch.bmm(
z_cot.detach(), z_lat.detach().transpose(1, 2)
) / self.match_temperature # (B, N_c, N_l)
# PSA positional bias
effective_strength = (
position_bias_strength_override
if position_bias_strength_override is not None
else self.position_bias_strength
)
if isinstance(effective_strength, (int, float)) and effective_strength > 0:
pos_bias = self._compute_position_bias(
N_c, N_l, effective_strength,
cot_mask=cot_mask, device=device,
)
match_logits_c2l = match_logits_c2l + pos_bias
match_logits_l2c = match_logits_c2l.transpose(1, 2) # (B, N_l, N_c)
# Mask out invalid positions
very_neg = torch.finfo(match_logits_c2l.dtype).min / 4
match_logits_c2l = match_logits_c2l.masked_fill(
~latent_mask.bool().unsqueeze(1), very_neg
)
match_logits_l2c = match_logits_l2c.masked_fill(
~cot_mask.bool().unsqueeze(1), very_neg
)
A_c2l = F.softmax(match_logits_c2l, dim=-1) # (B, N_c, N_l)
A_l2c = F.softmax(match_logits_l2c, dim=-1) # (B, N_l, N_c)
# 3. Cross-reconstruction (gradient flows through z_lat / z_cot → projector)
z_lat_mix = torch.bmm(A_c2l, z_lat) # (B, N_c, d)
z_cot_mix = torch.bmm(A_l2c, z_cot) # (B, N_l, d)
# 4. Loss: 1 - cosine_similarity (stop-gradient on targets)
cos_c2l = F.cosine_similarity(z_cot.detach(), z_lat_mix, dim=-1) # (B, N_c)
cos_l2c = F.cosine_similarity(z_lat.detach(), z_cot_mix, dim=-1) # (B, N_l)
loss_c2l = ((1 - cos_c2l) * cot_mask).sum() / cot_mask.sum().clamp(min=1)
loss_l2c = ((1 - cos_l2c) * latent_mask).sum() / latent_mask.sum().clamp(min=1)
loss = 0.5 * (loss_c2l + loss_l2c)
# 5. Diagnostics (same format as PMPSLatentSupervisor)
with torch.no_grad():
# Match entropy
A_c2l_valid = A_c2l[cot_mask.bool()]
match_entropy = -(
A_c2l_valid * A_c2l_valid.clamp(min=1e-8).log()
).sum(-1).mean()
# Match order correlation (Pearson between CoT position and
# attended latent position)
lat_indices = torch.arange(N_l, device=device, dtype=torch.float32)
attended_lat_pos = (
A_c2l * lat_indices.view(1, 1, -1)
).sum(dim=-1) # (B, N_c)
cot_indices = torch.arange(N_c, device=device, dtype=torch.float32)
corr_list = []
for b in range(B):
valid = cot_mask[b].bool()
x = cot_indices[valid]
y = attended_lat_pos[b, valid]
if len(x) > 2:
vx = x - x.mean()
vy = y - y.mean()
denom = (
vx.pow(2).sum().sqrt() * vy.pow(2).sum().sqrt()
).clamp(min=1e-8)
corr = (vx * vy).sum() / denom
if not torch.isnan(corr):
corr_list.append(corr)
match_order_corr = (
torch.stack(corr_list).mean()
if corr_list
else torch.tensor(0.0, device=device)
)
avg_cos_c2l = (cos_c2l * cot_mask).sum() / cot_mask.sum().clamp(min=1)
avg_cos_l2c = (cos_l2c * latent_mask).sum() / latent_mask.sum().clamp(min=1)
log_dict = {
"pmps_loss": loss.detach(),
"pmps_loss_c2l": loss_c2l.detach(),
"pmps_loss_l2c": loss_l2c.detach(),
"code_entropy_latent": torch.tensor(0.0, device=device),
"code_entropy_cot": torch.tensor(0.0, device=device),
"match_entropy_c2l": match_entropy,
"match_order_correlation": match_order_corr,
"avg_cosine_c2l": avg_cos_c2l,
"avg_cosine_l2c": avg_cos_l2c,
}
return loss, log_dict
|