| 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) |
|
|
| |
| z_lat = self._project(latent_hidden) |
| z_cot = self._project(cot_hidden) |
|
|
| |
| with torch.no_grad(): |
| match_logits_c2l = torch.bmm( |
| z_cot.detach(), z_lat.detach().transpose(1, 2) |
| ) / self.match_temperature |
|
|
| |
| 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) |
|
|
| |
| 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) |
| A_l2c = F.softmax(match_logits_l2c, dim=-1) |
|
|
| |
| z_lat_mix = torch.bmm(A_c2l, z_lat) |
| z_cot_mix = torch.bmm(A_l2c, z_cot) |
|
|
| |
| cos_c2l = F.cosine_similarity(z_cot.detach(), z_lat_mix, dim=-1) |
| cos_l2c = F.cosine_similarity(z_lat.detach(), z_cot_mix, dim=-1) |
|
|
| 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) |
|
|
| |
| with torch.no_grad(): |
| |
| A_c2l_valid = A_c2l[cot_mask.bool()] |
| match_entropy = -( |
| A_c2l_valid * A_c2l_valid.clamp(min=1e-8).log() |
| ).sum(-1).mean() |
|
|
| |
| |
| 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) |
| 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 |
|
|
