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