update with orthobridge

v4/classes/bridges/ortho_bridge.py

@@ -0,0 +1,141 @@
+"""ortho_bridge — OrthoBridge: wraps any underlying bridge with cross-tower
+orthogonality regularization.
+
+Forces each tower to encode information that the other towers DON'T encode by
+penalising cross-tower representational similarity.  The penalty is added to
+the main training loss via the `modify_loss` hook, so the fusion stage runner
+sees it transparently — no other code changes required.
+
+How it works
+------------
+1. Inner bridge fuses the embeddings into the usual single tensor (forward).
+2. While forward is running, we compute pairwise linear-CKA between every
+   pair of input embeddings and stash the average value as `self._stashed`.
+3. `modify_loss(loss)` returns `loss + ortho_weight * stashed`.  Gradients
+   from the penalty flow back through the embeddings into the tower weights,
+   pushing each tower's representations apart.
+
+CKA reference: Kornblith et al., "Similarity of Neural Network Representations
+Revisited" (ICML 2019).  Linear CKA on centred embeddings is in [0, 1]:
+    0  = orthogonal (uncorrelated) representations
+    1  = identical (up to linear transform)
+
+Config example:
+    {
+      "name":   "nt",
+      "type":   "fusion",
+      "module": "v4.classes.bridges.ortho_bridge",
+      "class":  "OrthoBridge",
+      "args": {
+        "ortho_weight":  0.1,
+        "inner_module": "v4.classes.bridges.fusion_bridge",
+        "inner_class":  "FusionBridge",
+        "inner_args":   { "fusion_dim": 256, "use_se": true }
+      }
+    }
+"""
+from __future__ import annotations
+
+import importlib
+from itertools import combinations
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+def _linear_cka(z_a: torch.Tensor, z_b: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+    """Linear CKA between two (B, D_a) and (B, D_b) embedding batches.
+
+    Centred Gram-matrix similarity, normalised to [0, 1].
+    """
+    a = z_a - z_a.mean(dim=0, keepdim=True)
+    b = z_b - z_b.mean(dim=0, keepdim=True)
+    Ga = a @ a.T
+    Gb = b @ b.T
+    num = (Ga * Gb).sum()
+    den = torch.sqrt((Ga * Ga).sum() * (Gb * Gb).sum() + eps)
+    return num / den
+
+
+class OrthoBridge(nn.Module):
+    """Wrap any underlying bridge and add cross-tower orthogonality.
+
+    Parameters
+    ----------
+    input_dims    : ordered list of input embedding dims  (passed to inner bridge)
+    fusion_dim    : convenience alias passed to inner bridge if it accepts it
+    ortho_weight  : λ on the orthogonality penalty (default 0.1)
+    inner_module  : import path of the inner bridge class
+    inner_class   : class name within `inner_module`
+    inner_args    : kwargs forwarded to the inner bridge constructor
+
+    The OrthoBridge does NOT touch the inner bridge's forward output — it only
+    computes and stashes a penalty during forward, then exposes it via the
+    `modify_loss` hook.
+    """
+
+    def __init__(
+        self,
+        input_dims:   list[int],
+        fusion_dim:   int   = 256,
+        ortho_weight: float = 0.1,
+        inner_module: str   = "v4.classes.bridges.fusion_bridge",
+        inner_class:  str   = "FusionBridge",
+        inner_args:   dict | None = None,
+    ):
+        super().__init__()
+        self.ortho_weight = float(ortho_weight)
+
+        # Build the inner bridge.  Pass fusion_dim through unless the caller's
+        # inner_args overrides it.
+        inner_kwargs: dict[str, Any] = dict(inner_args or {})
+        inner_kwargs.setdefault("fusion_dim", fusion_dim)
+        mod = importlib.import_module(inner_module)
+        cls = getattr(mod, inner_class)
+        self.inner = cls(input_dims, **inner_kwargs)
+        self.out_dim = self.inner.out_dim
+
+        # Stash penalty here on every forward; modify_loss reads from it.
+        self.register_buffer("_stashed", torch.zeros(()), persistent=False)
+        self._last_pairs_cka: list[float] = []
+
+    # ── core ────────────────────────────────────────────────────────────────
+
+    def forward(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
+        # Compute pairwise linear-CKA across the raw tower embeddings.
+        # We use the RAW per-tower embeddings (not the inner bridge's projected
+        # versions) because we want to push the TOWERS apart, not the bridge's
+        # internal projections.
+        if len(embeddings) >= 2:
+            ckas = []
+            for i, j in combinations(range(len(embeddings)), 2):
+                ckas.append(_linear_cka(embeddings[i], embeddings[j]))
+            penalty = torch.stack(ckas).mean()
+            self._stashed = penalty
+            self._last_pairs_cka = [float(c.detach().cpu()) for c in ckas]
+        else:
+            self._stashed = torch.zeros((), device=embeddings[0].device)
+            self._last_pairs_cka = []
+
+        return self.inner(embeddings)
+
+    def modify_loss(self, loss: torch.Tensor) -> torch.Tensor:
+        return loss + self.ortho_weight * self._stashed
+
+    # ── delegate to inner ───────────────────────────────────────────────────
+
+    def set_phase(self, phase: str) -> None:
+        if hasattr(self.inner, "set_phase"):
+            self.inner.set_phase(phase)
+
+    # ── introspection ───────────────────────────────────────────────────────
+
+    @property
+    def last_cka(self) -> float:
+        """Mean cross-tower linear CKA from the most recent forward pass.
+
+        Useful for logging — should DECREASE during training if the penalty
+        is doing its job.
+        """
+        return float(self._stashed.detach().cpu()) if self._stashed.numel() else 0.0