Learnable mask token + Kronecker marker covariance GP

.gitignore

@@ -213,3 +213,9 @@ checkpoints/
 *.out
 *.err
 *.log
+
+# local-only files
+.DS_Store
+CLAUDE.md
+docs/superpowers/
+immuvis.py

dump_marker_covariance.py

@@ -0,0 +1,194 @@
+"""Inspect the learned marker covariance K_C of a Kronecker-marker GP model.
+
+K_C is the C×C correlation across markers that distinguishes the marker-covariance
+model from the plain GP model (which implicitly assumes K_C = I, markers independent).
+It is image-independent: the Hyperkernel marker embeddings are an nn.Embedding lookup
+(immuvis.py), projected and row-normalised, so
+
+    K_C = normalize(embedding_projection(E)) @ normalize(...).T + marker_jitter·I
+
+(see gp_covariance.py:549-555). This script loads the checkpoint, rebuilds K_C over the
+full marker vocabulary, and reports whether it carries real off-diagonal structure or is
+essentially identity — the direct test of "did the model learn anything interesting?".
+
+Runs on CPU in seconds; no dataset needed. Run on szary where the checkpoint lives.
+"""
+
+import argparse
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from ruamel.yaml import YAML
+from scipy.cluster.hierarchy import leaves_list, linkage
+from scipy.spatial.distance import squareform
+
+from multiplex_model.modules.immuvis import MultiplexAutoencoder
+from multiplex_model.utils.configuration import DecoderConfig, EncoderConfig
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Dump and visualise the learned marker covariance K_C.")
+    parser.add_argument(
+        "--checkpoint",
+        default="/raid_encrypted/immucan/models/last_checkpoint-ImVs-34.pth",
+        help="Checkpoint with model_state_dict AND gp_covariance_state_dict.",
+    )
+    parser.add_argument(
+        "--model-config",
+        default="/raid_encrypted/immucan/models/config.last_checkpoint-ImVs-34.yaml",
+        help="Model config YAML (encoder/decoder + marker_jitter).",
+    )
+    parser.add_argument(
+        "--tokenizer-config",
+        default="/home/mzmyslowski/marcin_multiplex/configs/all_markers_tokenizer.yaml",
+    )
+    parser.add_argument(
+        "--panel-config",
+        default="/home/mzmyslowski/marcin_multiplex/configs/all_panels_config.yaml",
+    )
+    parser.add_argument(
+        "--panel",
+        default=None,
+        help="Restrict K_C to one dataset's markers (e.g. 'hn'). Markers only ever share a "
+        "K_C within their own panel during training, so the full vocabulary is not meaningful.",
+    )
+    parser.add_argument(
+        "--out",
+        default="/home/mzmyslowski/marcin_multiplex/logs/marker_covariance_ImVs-34.png",
+    )
+    parser.add_argument("--top-pairs", type=int, default=15, help="How many strongest marker pairs to print.")
+    return parser.parse_args()
+
+
+def build_marker_covariance(
+    hyperkernel_weights: torch.Tensor,
+    projection_weight: torch.Tensor,
+    projection_bias: torch.Tensor,
+) -> np.ndarray:
+    """Reproduce the row-normalised marker embeddings that feed K_C (gp_covariance.py:549-555)."""
+    e = hyperkernel_weights @ projection_weight.T + projection_bias  # [C, D]
+    e = torch.nn.functional.normalize(e, p=2, dim=1)
+    return e.numpy()
+
+
+def participation_ratio(eig: np.ndarray) -> float:
+    """(Σλ)²/Σλ² — an effective dimensionality; low when one component dominates."""
+    eig = eig[eig > 0]
+    return float(eig.sum() ** 2 / (eig**2).sum())
+
+
+def residual_correlation(e: np.ndarray) -> np.ndarray:
+    """Correlation of embeddings after removing the shared component.
+
+    K_C is dominated by a mean 'everything co-varies' direction; the marker-specific
+    structure lives in the residual. This is the biologically informative view.
+    """
+    r = e - e.mean(axis=0, keepdims=True)
+    r = r / np.linalg.norm(r, axis=1, keepdims=True)
+    corr: np.ndarray = r @ r.T
+    return corr
+
+
+def signed_pairs(k: np.ndarray, names: list[str], n: int) -> tuple[list, list]:
+    c = k.shape[0]
+    pairs = [(names[i], names[j], float(k[i, j])) for i in range(c) for j in range(i + 1, c)]
+    pairs.sort(key=lambda p: p[2])
+    return pairs[-n:][::-1], pairs[:n]
+
+
+def main() -> None:
+    args = parse_args()
+    yaml = YAML(typ="safe")
+
+    with open(args.tokenizer_config, "r") as f:
+        tokenizer = yaml.load(f)
+    inv_tokenizer = {v: k for k, v in tokenizer.items()}
+    model_num_channels = len(tokenizer)  # nn.Embedding row count — must match checkpoint
+
+    if args.panel:
+        with open(args.panel_config, "r") as f:
+            panel_markers = yaml.load(f)["markers"][args.panel]
+        names = [m for m in panel_markers if m in tokenizer]
+        channel_ids = [tokenizer[m] for m in names]
+        print(f"Panel '{args.panel}': {len(names)}/{len(panel_markers)} markers in tokenizer")
+    else:
+        channel_ids = sorted(tokenizer.values())
+        names = [inv_tokenizer[i] for i in channel_ids]
+    num_channels = len(names)
+
+    with open(args.model_config, "r") as f:
+        model_config = yaml.load(f)
+    marker_jitter = model_config.get("marker_jitter", 1e-2)
+
+    model = MultiplexAutoencoder(
+        num_channels=model_num_channels,
+        encoder_config=EncoderConfig(**model_config["encoder"]).model_dump(),
+        decoder_config=DecoderConfig(**model_config["decoder"]).model_dump(),
+    )
+    checkpoint = torch.load(args.checkpoint, map_location="cpu")
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+
+    if "gp_covariance_state_dict" not in checkpoint:
+        raise KeyError(
+            f"{args.checkpoint} has no 'gp_covariance_state_dict'. The embedding_projection "
+            "weights are not in this checkpoint, so K_C cannot be reconstructed — the analysis "
+            "would use random projection weights and be meaningless."
+        )
+    gp_state = checkpoint["gp_covariance_state_dict"]
+    projection_weight = gp_state["embedding_projection.weight"]
+    projection_bias = gp_state["embedding_projection.bias"]
+
+    with torch.no_grad():
+        ids = torch.tensor(channel_ids, dtype=torch.long)
+        embeddings = model.encoder.hyperkernel.hyperkernel_weights(ids)  # [C, model_dim]
+        e = build_marker_covariance(embeddings, projection_weight, projection_bias)
+
+    k = e @ e.T + marker_jitter * np.eye(num_channels)
+    k_resid = residual_correlation(e)
+    eig = np.linalg.eigvalsh(k)[::-1]
+
+    # K_C is dominated by a shared 'everything co-varies' component; the marker-specific
+    # structure is the residual. Report both so the shared component is not mistaken for collapse.
+    print(f"Markers (C):                     {num_channels}")
+    print(f"Leading eigenvector of K_C:      {eig[0] / eig.sum():.1%} of total  (shared component)")
+    print(f"||mean of embeddings||:          {np.linalg.norm(e.mean(0)):.3f}  (1.0 = all markers identical)")
+    print(f"Residual effective dimensions:   {participation_ratio(np.linalg.eigvalsh(k_resid)):.1f} / {num_channels}")
+    top_pos, top_neg = signed_pairs(k_resid, names, args.top_pairs)
+    print(f"\nTop {args.top_pairs} co-grouped marker pairs (residual, shared component removed):")
+    for a, b, v in top_pos:
+        print(f"  {v:+.3f}  {a} — {b}")
+    print(f"\nTop {args.top_pairs} anti-grouped marker pairs (residual):")
+    for a, b, v in top_neg:
+        print(f"  {v:+.3f}  {a} — {b}")
+
+    order = leaves_list(linkage(squareform(np.clip(1.0 - k_resid, 0.0, 2.0), checks=False), method="average"))
+    off = ~np.eye(num_channels, dtype=bool)
+    panel_tag = f" — {args.panel}" if args.panel else ""
+
+    fig, axes = plt.subplots(1, 2, figsize=(21, 9))
+    for ax, mat, title in ((axes[0], k, "K_C (full)"), (axes[1], k_resid, "residual (shared component removed)")):
+        mat = mat[np.ix_(order, order)]
+        labels = [names[i] for i in order]
+        vmax = float(np.abs(mat[off]).max())
+        im = ax.imshow(mat, cmap="RdBu_r", vmin=-vmax, vmax=vmax)
+        ax.set_xticks(range(num_channels), labels, rotation=90, fontsize=6)
+        ax.set_yticks(range(num_channels), labels, fontsize=6)
+        ax.set_title(f"{title}{panel_tag}\n(markers clustered by residual)", fontsize=12)
+        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+
+    fig.tight_layout()
+    fig.savefig(args.out, dpi=140, bbox_inches="tight")
+    npz_out = args.out.rsplit(".", 1)[0] + ".npz"
+    np.savez(npz_out, k_c=k, k_residual=k_resid, marker_names=np.array(names), channel_ids=np.array(channel_ids))
+    print(f"\nSaved figure: {args.out}")
+    print(f"Saved matrix: {npz_out}")
+
+
+if __name__ == "__main__":
+    main()

multiplex_model/losses.py

@@ -304,11 +304,12 @@ def forward(
         B, C, H, W = target.shape
         N = H * W
 
-        assert H == W == self.covariance_module.grid_size, (
-            f"Image must be square with H == W == grid_size, "
-            f"got {H}×{W} vs grid_size={self.covariance_module.grid_size}. "
-            f"Check downscale_factor or grid_size."
-        )
+        if H != W or H != self.covariance_module.grid_size:
+            raise ValueError(
+                f"Image must be square with H == W == grid_size, "
+                f"got {H}×{W} vs grid_size={self.covariance_module.grid_size}. "
+                f"Check downscale_factor or grid_size."
+            )
 
         # Reshape to [B, N, C] — loop over batch, batch over channels
         target_bnc = target.reshape(B, C, N).permute(0, 2, 1)   # [B, N, C]
@@ -388,4 +389,149 @@ def forward(
             "gp_nll":        gp_nll.item(),
             "total_loss":    total_loss.item(),
         }
+        return total_loss, loss_dict
+
+
+
+class KroneckerMarkerGPNLLLoss(nn.Module):
+    """
+    GP-based NLL loss with joint spatial + marker covariance.
+
+    Uses KroneckerMarkerCovariance for triple Kronecker (K_x ⊗ K_y) ⊗ K_C
+    plus Woodbury for per-pixel sigma. Processes one image at a time,
+    computing joint log-prob over all N*C dimensions.
+
+    Requires square images (H == W == grid_size after downscaling).
+    """
+
+    def __init__(
+        self,
+        covariance_module,
+        downscale_factor: int = 1,
+        device=None,
+    ):
+        super().__init__()
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.device = device
+        self.covariance_module = covariance_module
+        self.downscale_factor = downscale_factor
+
+    def _downscale(self, tensor: torch.Tensor) -> torch.Tensor:
+        if self.downscale_factor == 1:
+            return tensor
+        return torch.nn.functional.avg_pool2d(
+            tensor,
+            kernel_size=self.downscale_factor,
+            stride=self.downscale_factor,
+        )
+
+    def forward(
+        self,
+        target: torch.Tensor,
+        mu: torch.Tensor,
+        sigma: torch.Tensor,
+        marker_embeddings: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            target: [B, C, H, W] ground truth
+            mu:     [B, C, H, W] predicted means
+            sigma:  [B, C, H, W] per-pixel std dev (not log)
+            marker_embeddings: [B, C, model_dim] Hyperkernel embeddings
+
+        Returns:
+            Scalar mean NLL per pixel per channel.
+        """
+        target = target.float()
+        mu = mu.float()
+        sigma = sigma.float()
+        marker_embeddings = marker_embeddings.float()
+
+        if self.downscale_factor > 1:
+            target = self._downscale(target)
+            mu = self._downscale(mu)
+            sigma = self._downscale(sigma)
+
+        B, C, H, W = target.shape
+        N = H * W
+
+        assert H == W == self.covariance_module.grid_size, (
+            f"Image must be square with H == W == grid_size, "
+            f"got {H}x{W} vs grid_size={self.covariance_module.grid_size}."
+        )
+
+        target_bnc = target.reshape(B, C, N).permute(0, 2, 1)  # [B, N, C]
+        mu_bnc = mu.reshape(B, C, N).permute(0, 2, 1)
+        sigma_bnc = sigma.reshape(B, C, N).permute(0, 2, 1)
+
+        total_log_prob = torch.zeros((), device=self.device, dtype=torch.float32)
+        for b in range(B):
+            total_log_prob = total_log_prob + self.covariance_module.log_prob_joint(
+                mu_bnc[b],
+                sigma_bnc[b],
+                target_bnc[b],
+                marker_embeddings[b],
+            )
+
+        return -total_log_prob / (B * N * C)
+
+
+class HybridKroneckerMarkerGPNLLLoss(nn.Module):
+    """
+    Hybrid loss: standard pixel-wise NLL + Kronecker marker GP NLL.
+
+    L = (1 - lambda_gp) * L_standard + lambda_gp * L_kronecker_marker_gp
+
+    Drop-in replacement for HybridKroneckerGPNLLLoss with additional
+    marker_embeddings argument in forward().
+    """
+
+    def __init__(
+        self,
+        covariance_module,
+        lambda_gp: float = 0.1,
+        downscale_factor: int = 1,
+        device=None,
+    ):
+        super().__init__()
+        self.lambda_gp = lambda_gp
+        self.gp_loss = KroneckerMarkerGPNLLLoss(
+            covariance_module=covariance_module,
+            downscale_factor=downscale_factor,
+            device=device,
+        )
+
+    def forward(
+        self,
+        target: torch.Tensor,
+        mu: torch.Tensor,
+        logvar: torch.Tensor,
+        marker_embeddings: torch.Tensor,
+    ) -> tuple[torch.Tensor, dict]:
+        """
+        Args:
+            target: [B, C, H, W] ground truth
+            mu:     [B, C, H, W] predicted means
+            logvar: [B, C, H, W] predicted log-variances
+            marker_embeddings: [B, C, model_dim] Hyperkernel embeddings
+
+        Returns:
+            total_loss: Combined scalar loss.
+            loss_dict:  {"standard_nll", "gp_nll", "total_loss"}.
+        """
+        var = torch.exp(logvar)
+        standard_nll = torch.mean((target - mu) ** 2 / (var + 1e-8) + logvar)
+
+        sigma = torch.sqrt(var)
+        gp_nll = self.gp_loss(target, mu, sigma, marker_embeddings)
+
+        total_loss = (1 - self.lambda_gp) * standard_nll + self.lambda_gp * gp_nll
+
+        loss_dict = {
+            "standard_nll": standard_nll.item(),
+            "gp_nll": gp_nll.item(),
+            "total_loss": total_loss.item(),
+        }
         return total_loss, loss_dict

multiplex_model/modules/__init__.py

@@ -97,6 +97,7 @@
 
 # Gaussian Process components
 from .gp_covariance import (
+    KroneckerMarkerCovariance,
     LowRankPlusSpatialCovariance,
 )
 
@@ -133,5 +134,6 @@
     "MultiplexImageDecoder",
     "MultiplexAutoencoder",
     # Gaussian Process components
+    "KroneckerMarkerCovariance",
     "LowRankPlusSpatialCovariance",
 ]