# Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # # --------------------------------------------------------------- import torch import torch.nn as nn import torch.nn.functional as F class PMPSLatentSupervisor(nn.Module): """PMPS-style supervisor for latent reasoning hidden states. Setting ------- Each example yields two token sequences from the **same** language model: * latent hidden states ``H_lat ∈ R^{B × N_l × D}`` (few tokens) * CoT hidden states ``H_cot ∈ R^{B × N_c × D}`` (many tokens) Goal: supervise the latent tokens so that (a) one latent embedding may carry information from *multiple* CoT tokens (compression / coverage), and (b) different latent embeddings cover *different* CoT tokens (diversity / no collapse). """ def __init__( self, hidden_dim, feat_dim=128, hidden_mlp=512, nmb_prototypes=24, # ~4 × N_l (N_l=6 in CODI default) temperature=0.1, match_temperature=0.3, # softer cluster-space matching epsilon=0.05, sinkhorn_iterations=3, use_projector=True, position_bias_strength=0.0, # NEW: Gaussian positional bias strength; 0 = disabled ): super().__init__() self.nmb_prototypes = nmb_prototypes self.temperature = temperature self.match_temperature = match_temperature self.epsilon = epsilon self.sinkhorn_iterations = sinkhorn_iterations self.use_projector = use_projector self.position_bias_strength = position_bias_strength if use_projector: self.projector = self._build_projector(hidden_dim, hidden_mlp, feat_dim) proto_dim = feat_dim else: self.projector = nn.Identity() proto_dim = hidden_dim self.feat_dim = proto_dim self.prototypes = nn.Linear(proto_dim, nmb_prototypes, bias=False) # ------------------------------------------------------------------ # building blocks # ------------------------------------------------------------------ @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), ) @torch.no_grad() def normalize_prototypes(self): """L2-normalise prototype rows. Call once per step before forward, mirroring ``main_pmps.py``.""" w = self.prototypes.weight.data.clone() w = F.normalize(w, dim=1, p=2) self.prototypes.weight.copy_(w) def _project_and_score(self, hidden_states): """Project (B, N, D) -> normalised (B*N, feat_dim) and score against prototypes -> (B*N, K).""" 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) scores = self.prototypes(emb) # (B*N, K) return scores @torch.no_grad() def sinkhorn(self, scores): """Sinkhorn-Knopp on un-normalised prototype scores. scores : (M, K) cosine-similarity scores. returns Q : (M, K) rows sum to 1 (per-sample distribution over prototypes), columns sum to M / K (equipartition across prototypes). """ Q = torch.exp(scores / self.epsilon).t() # (K, M) Q = Q / (Q.sum() + 1e-8) K, M = Q.shape for _ in range(self.sinkhorn_iterations): Q = Q / (Q.sum(dim=1, keepdim=True) + 1e-8) Q = Q / K Q = Q / (Q.sum(dim=0, keepdim=True) + 1e-8) Q = Q / M Q = Q * M # rows now sum to 1 return Q.t() # (M, K) @torch.no_grad() def sinkhorn_per_sample(self, scores, batch_size, seq_len): """Per-sample Sinkhorn-Knopp. Instead of applying Sinkhorn on the flattened (B*N, K) matrix, apply it independently on each sample's (N, K) sub-matrix. This enforces that within EACH sample, the N tokens are distributed uniformly across K prototypes — directly preventing intra-sample collapse. Args: scores: (B*N, K) prototype scores (flattened). batch_size: B seq_len: N (number of tokens per sample, e.g. N_l or N_c) Returns: Q: (B*N, K) per-sample normalised codes. """ # Reshape to per-sample: (B, N, K) scores_3d = scores.reshape(batch_size, seq_len, K := scores.size(-1)) Q_list = [] for b in range(batch_size): Q_b = torch.exp(scores_3d[b] / self.epsilon).t() # (K, N) Q_b = Q_b / (Q_b.sum() + 1e-8) N = Q_b.size(1) for _ in range(self.sinkhorn_iterations): # Row normalisation: each prototype gets equal total mass Q_b = Q_b / (Q_b.sum(dim=1, keepdim=True) + 1e-8) Q_b = Q_b / K # Column normalisation: each token gets equal total mass Q_b = Q_b / (Q_b.sum(dim=0, keepdim=True) + 1e-8) Q_b = Q_b / N Q_b = Q_b * N # rows sum to 1 Q_list.append(Q_b.t()) # (N, K) return torch.cat(Q_list, dim=0) # (B*N, K) @torch.no_grad() def sinkhorn_per_sample_masked(self, scores_3d, mask): """Per-sample Sinkhorn with mask support. Args: scores_3d: (B, N, K) prototype scores. mask: (B, N) binary mask, 1 = valid token. Returns: Q: (B, N, K) per-sample normalised codes (masked positions = 0). """ B, N, K = scores_3d.shape Q_out = scores_3d.new_zeros(B, N, K) for b in range(B): valid = mask[b].bool() # (N,) n_valid = valid.sum().item() if n_valid == 0: continue s = scores_3d[b, valid, :] # (n_valid, K) Q_b = torch.exp(s / self.epsilon).t() # (K, n_valid) Q_b = Q_b / (Q_b.sum() + 1e-8) for _ in range(self.sinkhorn_iterations): Q_b = Q_b / (Q_b.sum(dim=1, keepdim=True) + 1e-8) Q_b = Q_b / K Q_b = Q_b / (Q_b.sum(dim=0, keepdim=True) + 1e-8) Q_b = Q_b / n_valid Q_b = Q_b * n_valid # each token's row sums to 1 Q_out[b, valid, :] = Q_b.t().to(dtype=Q_out.dtype) return Q_out # (B, N, K) # ------------------------------------------------------------------ # positional bias for progressive sequential matching # ------------------------------------------------------------------ @staticmethod def _compute_position_bias(N_c, N_l, strength, cot_mask=None, device=None): """Compute a Gaussian positional bias matrix. The bias encourages CoT token j to match latent token i when their normalised sequential positions are close, implementing the "progressive sequential PMPS" prior. Args: N_c: number of CoT tokens (padded length). N_l: number of latent tokens. strength: bias strength (larger = stronger sequential prior). 0 means no bias (original free matching). cot_mask: optional (B, N_c) binary mask. When provided the effective CoT length is computed per sample so that padding positions do not distort normalisation. device: target device. Returns: bias: (B, N_c, N_l) if cot_mask is given, else (1, N_c, N_l). """ 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) # [0, 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) # [0, 1] # (n_valid, 1) - (1, N_l) -> (n_valid, N_l) delta = cot_pos_norm.unsqueeze(1) - lat_pos_norm.unsqueeze(0) bias[b, :n_valid, :] = -strength * delta.pow(2) return bias # (B, N_c, N_l) 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) # (N_c, N_l) bias = -strength * delta.pow(2) return bias.unsqueeze(0) # (1, N_c, N_l) # ------------------------------------------------------------------ # forward # ------------------------------------------------------------------ def forward(self, latent_hidden, cot_hidden, latent_mask=None, cot_mask=None, position_bias_strength_override=None): """Compute the PMPS-style swapped loss between latent and CoT token sequences. Parameters ---------- latent_hidden : (B, N_l, D) cot_hidden : (B, N_c, D) latent_mask : optional (B, N_l) 0/1 cot_mask : optional (B, N_c) 0/1 Returns ------- loss : scalar 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), ( "latent and cot must share the hidden dim (same model)" ) B, N_l, _ = latent_hidden.shape _, N_c, _ = cot_hidden.shape K = self.nmb_prototypes device = latent_hidden.device # 1. score each token against the shared prototypes. latent_scores = self._project_and_score(latent_hidden) # (B*N_l, K) cot_scores = self._project_and_score(cot_hidden) # (B*N_c, K) # masks (flat) 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) latent_keep = latent_mask.reshape(-1).bool() cot_keep = cot_mask.reshape(-1).bool() # 2. compute per-token codes via Sinkhorn (no grad). with torch.no_grad(): q_latent_valid = self.sinkhorn(latent_scores[latent_keep]) # (M_l, K) q_cot_valid = self.sinkhorn(cot_scores[cot_keep]) # (M_c, K) Q_lat = latent_scores.new_zeros(B * N_l, K) Q_lat[latent_keep] = q_latent_valid.to(dtype=Q_lat.dtype) Q_lat = Q_lat.reshape(B, N_l, K) Q_cot = cot_scores.new_zeros(B * N_c, K) Q_cot[cot_keep] = q_cot_valid.to(dtype=Q_lat.dtype) Q_cot = Q_cot.reshape(B, N_c, K) # q_latent_valid = self.sinkhorn_per_sample(latent_scores[latent_keep], batch_size=B, seq_len=N_l) # cot_scores_3d = cot_scores.reshape(B, N_c, K) # Q_cot = self.sinkhorn_per_sample_masked(cot_scores_3d, mask=cot_mask) # (B, N_c, K) # Q_lat = latent_scores.new_zeros(B * N_l, K) # Q_lat[latent_keep] = q_latent_valid.to(dtype=Q_lat.dtype) # Q_lat = Q_lat.reshape(B, N_l, K) # Q_cot is already (B, N_c, K) with masked positions = 0, no need to reshape # 3. per-token softmax predictions over prototypes. log_p_lat = F.log_softmax(latent_scores.reshape(B, N_l, K) / self.temperature, dim=-1) log_p_cot = F.log_softmax(cot_scores.reshape(B, N_c, K) / self.temperature, dim=-1) p_lat = log_p_lat.exp() # (B, N_l, K) p_cot = log_p_cot.exp() # (B, N_c, K) # 4. cluster-based cross-view matching (codes-only) + positional bias. match_logits_c2l = torch.bmm(Q_cot, Q_lat.transpose(1, 2)) / self.match_temperature # (B, N_c, N_l) # NEW: apply Gaussian positional bias for progressive sequential matching 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) 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) # 5. mixture predictions. p_lat_mix_for_cot = torch.bmm(A_c2l, p_lat) # (B, N_c, K) p_cot_mix_for_lat = torch.bmm(A_l2c, p_cot) # (B, N_l, K) log_p_lat_mix_for_cot = (p_lat_mix_for_cot.clamp(min=1e-12)).log() log_p_cot_mix_for_lat = (p_cot_mix_for_lat.clamp(min=1e-12)).log() # 6. PMPS swapped CE loss (no pooling on the targets). ce_c2l = -(Q_cot * log_p_lat_mix_for_cot).sum(dim=-1) # (B, N_c) ce_l2c = -(Q_lat * log_p_cot_mix_for_lat).sum(dim=-1) # (B, N_l) loss_c2l = (ce_c2l * cot_mask).sum() / cot_mask.sum().clamp(min=1) loss_l2c = (ce_l2c * latent_mask).sum() / latent_mask.sum().clamp(min=1) loss = 0.5 * (loss_c2l + loss_l2c) # 7. diagnostics with torch.no_grad(): ent_lat = -(q_latent_valid * q_latent_valid.clamp(min=1e-8).log()).sum(-1).mean() ent_cot = -(q_cot_valid * q_cot_valid.clamp(min=1e-8).log()).sum(-1).mean() A_c2l_valid = A_c2l[cot_mask.bool()] match_entropy = -(A_c2l_valid * A_c2l_valid.clamp(min=1e-8).log()).sum(-1).mean() # NEW: match order correlation — measures whether the matching # preserves sequential order (Pearson correlation between CoT # token position and attended latent position). lat_indices = torch.arange(N_l, device=device, dtype=torch.float32) # weighted average latent position attended by each CoT token 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) ) log_dict = { "pmps_loss": loss.detach(), "pmps_loss_c2l": loss_c2l.detach(), "pmps_loss_l2c": loss_l2c.detach(), "code_entropy_latent": ent_lat, "code_entropy_cot": ent_cot, "match_entropy_c2l": match_entropy, "match_order_correlation": match_order_corr, } return loss_c2l, log_dict