Upload liqi files

0ab2566 verified 1 day ago

16.4 kB

# 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

	# 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