diff --git a/.gitignore b/.gitignore
index ad539b8..f8de0e9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -213,3 +213,9 @@ checkpoints/
 *.out
 *.err
 *.log
+
+# local-only files
+.DS_Store
+CLAUDE.md
+docs/superpowers/
+immuvis.py
diff --git a/dump_marker_covariance.py b/dump_marker_covariance.py
new file mode 100644
index 0000000..68dfd3c
--- /dev/null
+++ b/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()
diff --git a/kronecker_marker_summary_pl.pdf b/kronecker_marker_summary_pl.pdf
new file mode 100644
index 0000000..2e395b6
Binary files /dev/null and b/kronecker_marker_summary_pl.pdf differ
diff --git a/logs/marker_covariance_ImVs-34_danenberg.png b/logs/marker_covariance_ImVs-34_danenberg.png
new file mode 100644
index 0000000..d9727cc
Binary files /dev/null and b/logs/marker_covariance_ImVs-34_danenberg.png differ
diff --git a/logs/marker_covariance_ImVs-34_hn.png b/logs/marker_covariance_ImVs-34_hn.png
new file mode 100644
index 0000000..bd77c91
Binary files /dev/null and b/logs/marker_covariance_ImVs-34_hn.png differ
diff --git a/logs/marker_covariance_ImVs-34_hoch-rna.png b/logs/marker_covariance_ImVs-34_hoch-rna.png
new file mode 100644
index 0000000..58eb9cb
Binary files /dev/null and b/logs/marker_covariance_ImVs-34_hoch-rna.png differ
diff --git a/multiplex_model/losses.py b/multiplex_model/losses.py
index 34cc1e2..23443be 100644
--- a/multiplex_model/losses.py
+++ b/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
\ No newline at end of file
diff --git a/multiplex_model/modules/__init__.py b/multiplex_model/modules/__init__.py
index abe1c6b..81e77f6 100644
--- a/multiplex_model/modules/__init__.py
+++ b/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",
 ]
diff --git a/multiplex_model/modules/gp_covariance.py b/multiplex_model/modules/gp_covariance.py
index b83875a..0b768f9 100644
--- a/multiplex_model/modules/gp_covariance.py
+++ b/multiplex_model/modules/gp_covariance.py
@@ -421,3 +421,225 @@ def log_prob_all_markers(
         mahal      = (E * K_inv_E).sum()                         # scalar
 
         return -0.5 * (mahal + log_det_K_total + N * C * math.log(2 * math.pi))
+
+
+class KroneckerMarkerCovariance(nn.Module):
+    """
+    GP covariance with triple Kronecker structure + marker covariance + Woodbury.
+
+    Models K = (K_x ⊗ K_y) ⊗ K_C + U_block·U_blockᵀ + jitter·I
+
+    K_C is computed from Hyperkernel marker embeddings projected to a lower
+    dimension: K_C = E·Eᵀ + marker_jitter·I. Eigendecomposed every forward
+    pass (O(C³), cheap for C ≤ 40).
+
+    Spatial K_x, K_y are 1D Matérn kernels eigendecomposed once at init
+    (same as KroneckerPlusSpatialCovariance).
+
+    The full NC×NC covariance is never materialised. A⁻¹v is computed via
+    three einsum contractions (spatial x, spatial y, marker).
+    """
+
+    def __init__(
+        self,
+        grid_size: int,
+        marker_embed_dim: int,
+        hyperkernel_model_dim: int,
+        kernel_jitter: float = 1e-2,
+        marker_jitter: float = 1e-2,
+        spatial_matern_kernel_nu: float = 1.5,
+        spatial_matern_kernel_length_scale: float = 5.0,
+        device=None,
+    ):
+        super().__init__()
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.kernel_jitter = kernel_jitter
+        self.marker_jitter = marker_jitter
+        self.grid_size = grid_size
+        self.N = grid_size * grid_size
+
+        # --- Spatial eigendecomposition (identical to KroneckerPlusSpatialCovariance) ---
+        x1d = torch.linspace(0, 1, grid_size, device=device).unsqueeze(-1)
+
+        k1d = gpytorch.kernels.MaternKernel(nu=spatial_matern_kernel_nu).to(device)
+        k1d.lengthscale = spatial_matern_kernel_length_scale
+        k1d.raw_lengthscale.requires_grad = False
+
+        with torch.no_grad():
+            K1d = k1d(x1d).evaluate()
+            lam, V = torch.linalg.eigh(K1d)
+
+        self.register_buffer("lam", lam)
+        self.register_buffer("V", V)
+
+        # Spatial-only Kronecker eigenvalues (without jitter — jitter added in triple_eigs)
+        kron_eigs = torch.outer(lam, lam)  # [n, n]
+        self.register_buffer("kron_eigs", kron_eigs)
+
+        # --- Marker embedding projection ---
+        self.embedding_projection = nn.Linear(hyperkernel_model_dim, marker_embed_dim)
+
+    def _A_solve_triple(
+        self,
+        v: torch.Tensor,
+        V_C: torch.Tensor,
+        triple_eigs: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Solve A⁻¹v where A = (K_x ⊗ K_y) ⊗ K_C + jitter·I, analytically.
+
+        A = (V_x ⊗ V_y ⊗ V_C) diag(triple_eigs) (V_x ⊗ V_y ⊗ V_C)ᵀ
+
+        Applied via six einsum contractions (3 forward + divide + 3 reverse).
+
+        Args:
+            v: [NC] or [NC, m]
+            V_C: [C, C] eigenvectors of K_C
+            triple_eigs: [n, n, C] = kron_eigs[i,j] * lam_C[k] + jitter
+
+        Returns:
+            A⁻¹v, same shape as v.
+        """
+        n = self.grid_size
+        C = V_C.shape[0]
+        squeeze = v.dim() == 1
+        if squeeze:
+            v = v.unsqueeze(-1)
+        m = v.shape[-1]
+
+        # Reshape [NC, m] -> [n, n, C, m] (spatial_x, spatial_y, marker, rhs)
+        V3 = v.reshape(n, n, C, m)
+
+        # Forward transform: (V_x ⊗ V_y ⊗ V_C)ᵀ v
+        # Contract marker axis with V_C
+        tmp = torch.einsum("ijcm, ck -> ijkm", V3, V_C)
+        # Contract spatial_y axis with V
+        tmp = torch.einsum("ijkm, jb -> ibkm", tmp, self.V)
+        # Contract spatial_x axis with V
+        tmp = torch.einsum("ibkm, ia -> abkm", tmp, self.V)
+
+        # Divide by eigenvalues
+        tmp = tmp / triple_eigs.unsqueeze(-1)
+
+        # Reverse transform: (V_x ⊗ V_y ⊗ V_C) tmp
+        tmp = torch.einsum("abkm, jb -> ajkm", tmp, self.V)
+        tmp = torch.einsum("ajkm, ia -> ijkm", tmp, self.V)
+        tmp = torch.einsum("ijkm, ck -> ijcm", tmp, V_C)
+
+        result = tmp.reshape(n * n * C, m)
+        return result.squeeze(-1) if squeeze else result
+
+    def _compute_marker_eigen(
+        self, marker_embeddings: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Project embeddings, build K_C, eigendecompose, compute triple eigenvalues.
+
+        Args:
+            marker_embeddings: [C, hyperkernel_model_dim]
+
+        Returns:
+            (V_C, triple_eigs, K_C):
+                V_C: [C, C] eigenvectors
+                triple_eigs: [n, n, C] eigenvalues of A
+                K_C: [C, C] marker covariance
+        """
+        E = self.embedding_projection(marker_embeddings)  # [C, D]
+        # Normalize rows to unit norm so K_C is a correlation matrix (diagonal = 1 + jitter).
+        # This decouples K_C conditioning from embedding_projection weight scale, keeping
+        # condition numbers bounded by C rather than growing with embedding magnitude.
+        E = nn.functional.normalize(E, p=2, dim=1)
+        C = E.shape[0]
+        K_C = E @ E.T + self.marker_jitter * torch.eye(C, device=E.device, dtype=E.dtype)
+        # Use float64 for eigh: when C > marker_embed_dim, K_C has C-D repeated eigenvalues
+        # at exactly marker_jitter. LAPACK's divide-and-conquer fails on near-repeated
+        # eigenvalues in float32; float64 precision resolves convergence reliably.
+        lam_C, V_C = torch.linalg.eigh(K_C.double())
+        lam_C = lam_C.to(E.dtype)
+        V_C = V_C.to(E.dtype)
+
+        # triple_eigs[i, j, k] = kron_eigs[i,j] * lam_C[k] + kernel_jitter
+        triple_eigs = self.kron_eigs.unsqueeze(-1) * lam_C.unsqueeze(0).unsqueeze(0) + self.kernel_jitter
+
+        return V_C, triple_eigs, K_C.double().to(E.dtype)
+
+    def log_prob_joint(
+        self,
+        mu_all: torch.Tensor,
+        U_all: torch.Tensor,
+        targets: torch.Tensor,
+        marker_embeddings: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Joint log p(targets | mu, K) over all N pixels and C markers.
+
+        K = (K_x ⊗ K_y) ⊗ K_C + U_block·U_blockᵀ + jitter·I
+
+        Uses Woodbury identity with rank-C U_block.
+
+        Args:
+            mu_all:  [N, C] predicted means
+            U_all:   [N, C] per-pixel std dev per channel
+            targets: [N, C] ground truth
+            marker_embeddings: [C, hyperkernel_model_dim] raw Hyperkernel embeddings
+
+        Returns:
+            Scalar log probability.
+        """
+        N, C = targets.shape
+        NC = N * C
+
+        V_C, triple_eigs, _ = self._compute_marker_eigen(marker_embeddings)
+
+        # Error vector in spatial-major order: [pix0_ch0, pix0_ch1, ..., pixN_chC]
+        e = (targets - mu_all).reshape(-1)  # [NC]
+
+        # Build U_block [NC, C] in spatial-major order: row (i*C + c) = pixel i, marker c
+        U_block = torch.diag_embed(U_all).reshape(NC, C)
+
+        # log det(A)
+        if (triple_eigs <= 0).any():
+            raise RuntimeError(
+                f"Non-positive triple eigenvalues (min={triple_eigs.min().item():.3e}). "
+                "Increase kernel_jitter or marker_jitter."
+            )
+        log_det_A = triple_eigs.log().sum()
+
+        # A⁻¹ applied to error and U_block columns (C+1 RHS, batched)
+        rhs = torch.cat([e.unsqueeze(-1), U_block], dim=-1)  # [NC, C+1]
+        A_inv_rhs = self._A_solve_triple(rhs, V_C, triple_eigs)  # [NC, C+1]
+        A_inv_e = A_inv_rhs[:, 0]       # [NC]
+        A_inv_U = A_inv_rhs[:, 1:]      # [NC, C]
+
+        # Woodbury inner matrix: M = I_C + U_blockᵀ A⁻¹ U_block  [C, C]
+        M = torch.eye(C, device=e.device, dtype=e.dtype) + U_block.T @ A_inv_U
+
+        # log det(K) = log det(A) + log det(M)
+        log_det_K = log_det_A + torch.linalg.slogdet(M)[1]
+
+        # K⁻¹ e = A⁻¹e - A⁻¹U M⁻¹ Uᵀ A⁻¹e
+        Ut_Ainv_e = U_block.T @ A_inv_e  # [C]
+        correction = A_inv_U @ torch.linalg.solve(M, Ut_Ainv_e)  # [NC]
+        K_inv_e = A_inv_e - correction
+
+        mahal = e @ K_inv_e
+
+        return -0.5 * (mahal + log_det_K + NC * math.log(2 * math.pi))
+
+    def compute_marker_correlation(self, marker_embeddings: torch.Tensor) -> torch.Tensor:
+        """
+        Compute C×C correlation matrix from projected marker embeddings.
+
+        Args:
+            marker_embeddings: [C, hyperkernel_model_dim]
+
+        Returns:
+            [C, C] correlation matrix (ones on diagonal).
+        """
+        E = nn.functional.normalize(self.embedding_projection(marker_embeddings), p=2, dim=1)
+        K_C = E @ E.T + self.marker_jitter * torch.eye(E.shape[0], device=E.device, dtype=E.dtype)
+        # Normalize to correlation: corr[i,j] = K_C[i,j] / sqrt(K_C[i,i] * K_C[j,j])
+        diag_sqrt = torch.sqrt(torch.diag(K_C))
+        return K_C / (diag_sqrt.unsqueeze(0) * diag_sqrt.unsqueeze(1))
diff --git a/multiplex_model/modules/immuvis.py b/multiplex_model/modules/immuvis.py
index fe4dc3f..5ed9fca 100644
--- a/multiplex_model/modules/immuvis.py
+++ b/multiplex_model/modules/immuvis.py
@@ -1,3 +1,4 @@
+import copy
 from typing import Literal
 
 import torch
@@ -151,6 +152,8 @@ def __init__(
         pm_layers_blocks: list[int],
         pm_embedding_dims: list[int],
         use_latent_norm: bool = False,
+        use_mask_token: bool = False,
+        mask_token_init: float = 0.0,
         encoder_type: str | type[Encoder] | dict = "convnext",
     ):
         """Initialize the Multiplex Image Encoder.
@@ -163,6 +166,8 @@ def __init__(
             pm_layers_blocks (List[int]): Number of blocks in each pan-marker layer.
             pm_embedding_dims (List[int]): Embedding dimensions for each pan-marker layer.
             use_latent_norm (bool, optional): Whether to apply LayerNorm to the latent representation. Defaults to False.
+            use_mask_token (bool, optional): Whether to use a learnable mask token for spatially masked pixels. Defaults to False.
+            mask_token_init (float, optional): Initial value for the learnable mask token. Defaults to 0.0.
             encoder_type (Union[str, Type[Encoder], Dict], optional): Type of encoder to use.
                 Can be a string (registry name), Encoder class, or config dict with 'type' and 'module_parameters'.
                 For ConvNeXtEncoder, module_parameters can include 'block_parameters' dict with ConvNextBlock parameters
@@ -171,6 +176,11 @@ def __init__(
         """
         super().__init__()
 
+        self.use_mask_token = use_mask_token
+        self.mask_token = (
+            nn.Parameter(torch.tensor(mask_token_init)) if use_mask_token else None
+        )
+
         # Resolve encoder class
         encoder_cls = resolve_encoder_class(encoder_type)
 
@@ -223,6 +233,7 @@ def forward(
         self,
         x: torch.Tensor,
         encoded_indices: torch.Tensor,
+        spatial_mask: torch.Tensor | None = None,
         return_features: bool = False,
     ) -> dict:
         """Forward pass of the encoder.
@@ -230,6 +241,8 @@ def forward(
         Args:
             x (torch.Tensor): Multiplex images batch tensor with shape [B, C, H, W]
             encoded_indices (torch.Tensor): Indices of the markers in channels tensor with shape [B, C].
+            spatial_mask (torch.Tensor | None, optional): Binary mask indicating spatially masked pixels [B, C, H, W].
+                When provided and use_mask_token is True, masked pixels are replaced with the mask token. Defaults to None.
             return_features (bool, optional): If True, returns the features after each block. Defaults to False.
 
         Returns:
@@ -239,6 +252,9 @@ def forward(
         features = []
 
         B, C, H, W = x.shape
+        if self.use_mask_token and spatial_mask is not None:
+            mask_token = self.mask_token.to(dtype=x.dtype)
+            x = torch.where(spatial_mask, mask_token, x)
         x = x.reshape(B * C, 1, H, W)
         x = self.marker_agnostic_encoder(x, return_features=return_features)
         if return_features:
@@ -369,6 +385,11 @@ def __init__(
             decoder_config (dict): Configuration for the decoder.
         """
         super().__init__()
+        self._architecture_config = {
+            "num_channels": num_channels,
+            "encoder_config": copy.deepcopy(encoder_config),
+            "decoder_config": copy.deepcopy(decoder_config),
+        }
         self.latent_dim = encoder_config["pm_embedding_dims"][-1]
         self.num_channels = num_channels
 
@@ -390,10 +411,70 @@ def __init__(
             **decoder_config,
         )
 
+    def get_architecture_config(self, by_alias: bool = False) -> dict:
+        """Returns the stored architecture configuration.
+
+        Args:
+            by_alias (bool, optional): If True, renames keys to use aliases for compatibility.
+                Defaults to False.
+
+        Returns:
+            dict: A deep copy of the architecture configuration.
+        """
+        config = copy.deepcopy(self._architecture_config)
+        if by_alias:
+            config = config.copy()
+            config["encoder"] = config.pop("encoder_config")
+            config["decoder"] = config.pop("decoder_config")
+            config["encoder"]["hyperkernel"] = config["encoder"].pop("hyperkernel_config")
+            config["decoder"]["hyperkernel"] = config["decoder"].pop("hyperkernel_config")
+        return config
+
+    @classmethod
+    def load_from_checkpoint(
+        cls,
+        checkpoint: str | dict,
+        map_location: str | torch.device | None = None,
+        model_config: dict | None = None,
+        strict: bool = True,
+    ) -> "MultiplexAutoencoder":
+        """Load a MultiplexAutoencoder from a checkpoint.
+
+        Args:
+            checkpoint (str | dict): Path to checkpoint file or checkpoint dict.
+            map_location (str | torch.device | None, optional): Location to map checkpoint to.
+                Defaults to None.
+            model_config (dict | None, optional): Model configuration to use if checkpoint
+                does not contain 'model_config'. Defaults to None.
+            strict (bool, optional): Whether to strictly enforce that all state_dict keys
+                match the model. Defaults to True.
+
+        Returns:
+            MultiplexAutoencoder: Loaded model with state_dict applied.
+
+        Raises:
+            ValueError: If checkpoint and model_config do not provide configuration.
+        """
+        if isinstance(checkpoint, dict):
+            checkpoint_data = checkpoint
+        else:
+            checkpoint_data = torch.load(checkpoint, map_location=map_location)
+
+        resolved_config = checkpoint_data.get("model_config", model_config)
+        if resolved_config is None:
+            raise ValueError(
+                "Checkpoint is missing 'model_config'; provide model_config to load the model."
+            )
+
+        model = cls(**resolved_config)
+        model.load_state_dict(checkpoint_data["model_state_dict"], strict=strict)
+        return model
+
     def encode(
         self,
         x: torch.Tensor,
         encoded_indices: torch.Tensor,
+        spatial_mask: torch.Tensor | None = None,
         return_features: bool = False,
     ) -> dict:
         """Encode the input images using the encoder.
@@ -401,13 +482,19 @@ def encode(
         Args:
             x (torch.Tensor): Input images tensor with shape (B, C, H, W).
             encoded_indices (torch.Tensor): Indices of the markers in channels.
+            spatial_mask (torch.Tensor | None, optional): Binary mask indicating spatially masked pixels [B, C, H, W].
+                When provided and the encoder has use_mask_token enabled, masked pixels are replaced with the mask token.
+                Defaults to None.
             return_features (bool, optional): If True, returns the features after encoding. Defaults to False.
 
         Returns:
             dict: A dictionary containing the encoded images tensor (under 'output') and optionally the features.
         """
         encoding_output = self.encoder(
-            x, encoded_indices, return_features=return_features
+            x,
+            encoded_indices,
+            spatial_mask=spatial_mask,
+            return_features=return_features,
         )
         outputs = {"output": encoding_output["output"]}
 
@@ -437,6 +524,7 @@ def forward(
         x: torch.Tensor,
         encoded_indices: torch.Tensor,
         decoded_indices: torch.Tensor,
+        spatial_mask: torch.Tensor | None = None,
         return_features: bool = False,
     ) -> dict:
         """Forward pass of the Multiplex Autoencoder.
@@ -447,12 +535,16 @@ def forward(
                 for encoding.
             decoded_indices (torch.Tensor): Indices of the markers in channels
                 for decoding.
+            spatial_mask (torch.Tensor | None, optional): Binary mask indicating spatially masked pixels [B, C, H, W].
+                When provided and the encoder has use_mask_token enabled, masked pixels are replaced with the mask token.
+                Defaults to None.
+            return_features (bool, optional): If True, returns the features after encoding. Defaults to False.
 
         Returns:
             dict: A dictionary containing the reconstructed images tensor (under 'output') and optionally the features.
         """
         encoding_output = self.encode(
-            x, encoded_indices, return_features=return_features
+            x, encoded_indices, spatial_mask=spatial_mask, return_features=return_features
         )
         x = encoding_output["output"]
         x = self.decode(x, decoded_indices)
diff --git a/multiplex_model/utils/__init__.py b/multiplex_model/utils/__init__.py
index 9aeb72d..c5a1d6a 100644
--- a/multiplex_model/utils/__init__.py
+++ b/multiplex_model/utils/__init__.py
@@ -13,6 +13,7 @@
 from .masking import (
     apply_channel_masking,
     apply_spatial_masking,
+    get_pixel_mask,
 )
 from .optim import (
     ClampWithGrad,
@@ -25,6 +26,7 @@
     get_run_name,
     init_experiment,
     log_training_metrics,
+    log_validation_batch_metrics,
     log_validation_images,
     log_validation_metrics,
     plot_reconstructs_with_masks,
@@ -41,11 +43,13 @@
     # Masking
     "apply_channel_masking",
     "apply_spatial_masking",
+    "get_pixel_mask",
     # Logging
     "plot_reconstructs_with_uncertainty",
     "plot_reconstructs_with_masks",
     "init_experiment",
     "log_training_metrics",
+    "log_validation_batch_metrics",
     "log_validation_metrics",
     "log_validation_images",
     "get_run_name",
diff --git a/multiplex_model/utils/configuration.py b/multiplex_model/utils/configuration.py
index 736ff46..1dcd88f 100644
--- a/multiplex_model/utils/configuration.py
+++ b/multiplex_model/utils/configuration.py
@@ -3,7 +3,7 @@
 import os
 from typing import Any
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, Field, ValidationInfo, field_validator, model_validator
 
 from .train_logging import get_run_name
 
@@ -84,6 +84,14 @@ class EncoderConfig(BaseModel):
             "Can be a string (e.g., 'convnext') or a dict with 'type' and 'module_parameters'."
         ),
     )
+    use_mask_token: bool = Field(
+        default=False,
+        description="Whether to replace spatially-masked pixels with a learnable scalar token",
+    )
+    mask_token_init: float = Field(
+        default=0.0,
+        description="Initial value for the learnable mask token",
+    )
 
     @field_validator("ma_layers_blocks", "pm_layers_blocks")
     @classmethod
@@ -101,7 +109,7 @@ def validate_embedding_dims(cls, v: list[int]) -> list[int]:
 
     @field_validator("ma_embedding_dims")
     @classmethod
-    def validate_ma_lengths(cls, v: list[int], info) -> list[int]:
+    def validate_ma_lengths(cls, v: list[int], info: ValidationInfo) -> list[int]:
         if "ma_layers_blocks" in info.data:
             blocks = info.data["ma_layers_blocks"]
             if len(v) != len(blocks):
@@ -121,7 +129,7 @@ def validate_pm_not_empty(cls, v: list[int]) -> list[int]:
 
     @field_validator("pm_embedding_dims")
     @classmethod
-    def validate_pm_lengths(cls, v: list[int], info) -> list[int]:
+    def validate_pm_lengths(cls, v: list[int], info: ValidationInfo) -> list[int]:
         if len(v) == 0:
             raise ValueError(
                 "pm_embedding_dims cannot be empty - at least one pan-marker layer is required"
@@ -160,7 +168,7 @@ class DecoderConfig(BaseModel):
 
     @field_validator("block_type", mode="before")
     @classmethod
-    def validate_block_type(cls, v) -> ModuleConfig:
+    def validate_block_type(cls, v: Any) -> ModuleConfig:
         if v is None:
             return ModuleConfig(type="convnext")
         if isinstance(v, ModuleConfig):
@@ -242,6 +250,15 @@ class TrainingConfig(BaseModel):
     use_kronecker_gp: bool = Field(
         False, description="Whether to use Kronecker GP loss instead of CG-based GP loss"
     )
+    use_marker_covariance: bool = Field(
+        False, description="Whether to use marker covariance in Kronecker GP loss (requires use_kronecker_gp=True)"
+    )
+    marker_embed_dim: int = Field(
+        32, gt=0, description="Projection dimension for marker embeddings in K_C computation"
+    )
+    marker_jitter: float = Field(
+        1e-2, gt=0, description="Jitter added to marker covariance K_C for numerical stability"
+    )
 
     # Model architecture
     encoder_config: EncoderConfig = Field(
@@ -256,6 +273,10 @@ class TrainingConfig(BaseModel):
         None,
         description="Path to checkpoint to resume from. Use 'last' to load last checkpoint if available",
     )
+    reset_lr_schedule: bool = Field(
+        False,
+        description="When resuming, ignore checkpoint's scheduler/optimizer state and start a fresh LR schedule",
+    )
     checkpoints_dir: str = Field(
         "checkpoints", description="Directory to save checkpoints"
     )
@@ -274,6 +295,12 @@ class TrainingConfig(BaseModel):
     )
     run_name: str | None = Field(None, description="Name for Comet.ml experiment")
 
+    @model_validator(mode='after')
+    def _validate_marker_covariance_requires_kronecker(self) -> 'TrainingConfig':
+        if self.use_marker_covariance and not self.use_kronecker_gp:
+            raise ValueError("use_marker_covariance=True requires use_kronecker_gp=True")
+        return self
+
     def resolve_checkpoint(self) -> bool:
         """Resolve checkpoint path and determine if checkpoint should be loaded.
 
diff --git a/multiplex_model/utils/masking.py b/multiplex_model/utils/masking.py
index 5501350..7b7d0cc 100644
--- a/multiplex_model/utils/masking.py
+++ b/multiplex_model/utils/masking.py
@@ -58,15 +58,15 @@ def apply_channel_masking(
     )
     num_channels_to_mask = np.random.randint(1, max_channels_to_mask + 1)
 
-    masked_img = []
-    active_channel_ids = []
+    masked_img_list: list[torch.Tensor] = []
+    active_channel_ids_list: list[torch.Tensor] = []
     for b_i in range(batch_size):
         channels_to_keep = torch.randperm(num_sampled_channels)[num_channels_to_mask:]
-        masked_img.append(img[b_i : b_i + 1, channels_to_keep, :, :])
-        active_channel_ids.append(channel_ids[b_i : b_i + 1, channels_to_keep])
+        masked_img_list.append(img[b_i : b_i + 1, channels_to_keep, :, :])
+        active_channel_ids_list.append(channel_ids[b_i : b_i + 1, channels_to_keep])
 
-    masked_img = torch.cat(masked_img, dim=0)  # [B, C_active, H, W]
-    active_channel_ids = torch.cat(active_channel_ids, dim=0)  # [B, C_active]
+    masked_img = torch.cat(masked_img_list, dim=0)  # [B, C_active, H, W]
+    active_channel_ids = torch.cat(active_channel_ids_list, dim=0)  # [B, C_active]
 
     return img, channel_ids, masked_img, active_channel_ids
 
@@ -107,3 +107,25 @@ def apply_spatial_masking(
     masked_img[pixel_mask] = mask_fill_value
 
     return masked_img, pixel_mask
+
+
+def get_pixel_mask(
+    img: torch.Tensor,
+    spatial_masking_ratio: float = 0.6,
+    mask_patch_size: int = 8,
+) -> torch.Tensor:
+    """Create a boolean spatial patch mask without modifying the image.
+
+    Unlike apply_spatial_masking, this only returns the boolean mask — the
+    caller decides how masked pixels are filled (e.g., via a learnable token).
+
+    Args:
+        img (torch.Tensor): Input images [B, C, H, W] — used only for shape/device.
+        spatial_masking_ratio (float): Fraction of patches to mask.
+        mask_patch_size (int): Size of each square patch to mask.
+
+    Returns:
+        torch.Tensor: Boolean mask [B, C, H, W], True where pixels are masked.
+    """
+    _, mask = apply_spatial_masking(img, spatial_masking_ratio, mask_patch_size)
+    return mask
diff --git a/multiplex_model/utils/optim.py b/multiplex_model/utils/optim.py
index b34b3cf..2882b4d 100644
--- a/multiplex_model/utils/optim.py
+++ b/multiplex_model/utils/optim.py
@@ -10,15 +10,15 @@ class ClampWithGrad(torch.autograd.Function):
     """Custom autograd function for clamping with smooth gradients."""
 
     @staticmethod
-    def forward(ctx, x, min_val=-15.0, max_val=15.0):
+    def forward(ctx: torch.autograd.function.FunctionCtx, x: torch.Tensor, min_val: float = -15.0, max_val: float = 15.0) -> torch.Tensor:  # type: ignore[override]
         ctx.save_for_backward(x)
-        ctx.min_val, ctx.max_val = min_val, max_val
+        ctx.min_val, ctx.max_val = min_val, max_val  # type: ignore[attr-defined,misc]
         return x.clamp(min_val, max_val)
 
     @staticmethod
-    def backward(ctx, grad_output):
-        (x,) = ctx.saved_tensors
-        min_val, max_val = ctx.min_val, ctx.max_val
+    def backward(ctx: torch.autograd.function.FunctionCtx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:  # type: ignore[override]
+        (x,) = ctx.saved_tensors  # type: ignore[attr-defined,misc]
+        min_val, max_val = ctx.min_val, ctx.max_val  # type: ignore[attr-defined]
 
         grad_input = grad_output.clone()
         tanh_x = torch.tanh(x)
@@ -58,7 +58,7 @@ def get_scheduler_with_warmup(
     """
     final_lr_mult = final_lr / peak_lr
 
-    def lr_lambda(current_step, type: Literal["cosine", "linear"] = "cosine"):
+    def lr_lambda(current_step: int, type: Literal["cosine", "linear"] = "cosine") -> float:
         if current_step < num_warmup_steps:
             return float(max(1, current_step)) / float(max(1, num_warmup_steps))
         elif current_step >= num_annealing_steps + num_warmup_steps:
diff --git a/multiplex_model/utils/train_logging.py b/multiplex_model/utils/train_logging.py
index 83e7ddd..8c4043c 100644
--- a/multiplex_model/utils/train_logging.py
+++ b/multiplex_model/utils/train_logging.py
@@ -1,5 +1,6 @@
 """Logging and visualization utilities for training and validation."""
 
+import os
 import re
 from datetime import datetime
 from io import BytesIO
@@ -29,7 +30,7 @@ def plot_reconstructs_with_uncertainty(
     ncols: int = 9,
     scale_by_max: bool = True,
     partially_masked_ids: list[int] = [],
-):
+) -> plt.Figure:
     """Plot the original image and the reconstructed image with uncertainty.
 
     Args:
@@ -66,7 +67,7 @@ def plot_reconstructs_with_uncertainty(
         ax_uncertainty.axis("off")
 
         if j < num_channels:
-            marker_name = markers_names_map[channel_ids[0, j].item()]
+            marker_name = markers_names_map[int(channel_ids[0, j].item())]
             ax_img.imshow(orig_img[0, j].cpu().numpy(), cmap="CMRmap", vmin=0, vmax=1)
             ax_img.set_title(f"Original\n{marker_name}")
 
@@ -111,7 +112,7 @@ def plot_reconstructs_with_masks(
     fully_masked_ids: list[int],
     markers_names_map: dict[int, str],
     ncols: int = 9,
-):
+) -> plt.Figure:
     """Plot the original image, masked image (with white pixels where masked), and reconstruction.
 
     Args:
@@ -148,7 +149,7 @@ def plot_reconstructs_with_masks(
         ax_reconstructed = ax_flat[i + 2]
 
         if j < num_channels:
-            channel_id = channel_ids[0, j].item()
+            channel_id = int(channel_ids[0, j].item())
             marker_name = markers_names_map[channel_id]
 
             # Show original
@@ -182,7 +183,7 @@ def plot_reconstructs_with_masks(
                 masked_idx = channel_to_masked_idx[channel_id]
 
                 # Convert grayscale to RGBA using colormap (image already normalized to 0-1)
-                cmap = plt.cm.CMRmap
+                cmap = plt.cm.CMRmap  # type: ignore[attr-defined]
                 img_data = orig_img[0, j].cpu().numpy()
                 rgba_img = cmap(img_data)  # Apply colormap directly
 
@@ -259,8 +260,8 @@ def get_next_version_number(
 
         latest_experiment = experiments[0]
 
-        version = re.match(version_pattern, latest_experiment.name)
-        version = int(version.group(1))
+        m = re.match(version_pattern, latest_experiment.name)
+        version = int(m.group(1)) if m else 0
 
         # Return next version (1 if no versions exist)
         return version + 1
@@ -295,13 +296,20 @@ def init_experiment(config: dict[str, Any]) -> None:
                 api_key=config.get("comet_api_key"),
             )
             run_name = f"ImVs-{version}"
+            # Parallel jobs race on the version query and can get the same number;
+            # the SLURM job id disambiguates so their checkpoints don't overwrite.
+            slurm_job_id = os.environ.get("SLURM_JOB_ID")
+            if slurm_job_id:
+                run_name = f"{run_name}-{slurm_job_id}"
         else:
             # Fallback to date-time as default run name
             run_name = datetime.now().strftime("%m%d_%H:%M:%S")
 
     print(f"Run name: {run_name}")
     _experiment.set_name(run_name)
-    _experiment.add_tags(config.get("tags", []))
+    tags = config.get("tags", [])
+    if tags:
+        _experiment.add_tags(tags)
     _experiment.log_parameters(config)
 
 
@@ -315,6 +323,7 @@ def log_training_metrics(
     step: int | None = None,
     standard_nll: float | None = None,
     gp_nll: float | None = None,
+    mask_token: float | None = None,
 ) -> None:
     """Log training metrics to Comet.ml.
 
@@ -328,6 +337,7 @@ def log_training_metrics(
         step (int | None): Step number for logging
         standard_nll (float | None): Standard NLL loss component (GP training)
         gp_nll (float | None): GP-based NLL loss component (GP training)
+        mask_token (float | None): Learnable mask token value
     """
     if _experiment is None:
         return
@@ -344,6 +354,8 @@ def log_training_metrics(
         metrics["train/standard_nll"] = standard_nll
     if gp_nll is not None:
         metrics["train/gp_nll"] = gp_nll
+    if mask_token is not None:
+        metrics["train/mask_token"] = mask_token
     _experiment.log_metrics(metrics, step=step)
 
 
@@ -354,6 +366,9 @@ def log_validation_metrics(
     latent_rankme: float,
     epoch: int,
     variance_mae_correlation: float | None = None,
+    variance_mse_correlation: float | None = None,
+    val_standard_nll: float | None = None,
+    val_gp_nll: float | None = None,
 ) -> None:
     """Log validation metrics to Comet.ml.
 
@@ -364,6 +379,7 @@ def log_validation_metrics(
         latent_rankme (float): RankMe metric for latent representations
         epoch (int): Current epoch number
         variance_mae_correlation (Optional[float]): Pearson correlation between predicted variances and MAEs per channel
+        variance_mse_correlation (Optional[float]): Pearson correlation between predicted variances and MSEs per channel
     """
     if _experiment is None:
         return
@@ -376,9 +392,33 @@ def log_validation_metrics(
     }
     if variance_mae_correlation is not None:
         metrics["val/variance_mae_correlation"] = variance_mae_correlation
+    if variance_mse_correlation is not None:
+        metrics["val/variance_mse_correlation"] = variance_mse_correlation
+    if val_standard_nll is not None:
+        metrics["val/standard_nll"] = val_standard_nll
+    if val_gp_nll is not None:
+        metrics["val/gp_nll"] = val_gp_nll
     _experiment.log_metrics(metrics, epoch=epoch)
 
 
+def log_validation_batch_metrics(
+    variance_mse_correlation_per_batch: float,
+    step: int,
+) -> None:
+    """Log per-batch validation metrics to Comet.ml.
+
+    Args:
+        variance_mse_correlation_per_batch (float): Pearson correlation between predicted variances and MSEs per channel for a single batch
+        step (int): Global step number
+    """
+    if _experiment is None:
+        return
+    _experiment.log_metrics(
+        {"val/variance_mse_correlation_per_batch": variance_mse_correlation_per_batch},
+        step=step,
+    )
+
+
 def log_validation_images(
     fig: plt.Figure,
     panel_idx: int,
@@ -386,6 +426,7 @@ def log_validation_images(
     epoch: int,
     masked_channels_names: str,
     img_idx: int,
+    name_suffix: str = "",
 ) -> None:
     """Log validation reconstruction images to Comet.ml.
 
@@ -396,6 +437,7 @@ def log_validation_images(
         epoch (int): Current epoch number
         masked_channels_names (str): Names of masked channels
         img_idx (int): Index of the image in the batch
+        name_suffix (str): Optional suffix appended to the image name
     """
     if _experiment is None:
         return
@@ -408,7 +450,7 @@ def log_validation_images(
 
     _experiment.log_image(
         img,
-        name=f"val/reconstructions_panel-{panel_idx}_epoch-{epoch + 1}_img-{img_idx}",
+        name=f"val/reconstructions_panel-{panel_idx}_epoch-{epoch + 1}_img-{img_idx}{name_suffix}",
         step=epoch,
         metadata={
             "panel_idx": panel_idx,
diff --git a/notes/imvs34_marker_covariance.md b/notes/imvs34_marker_covariance.md
new file mode 100644
index 0000000..be90416
--- /dev/null
+++ b/notes/imvs34_marker_covariance.md
@@ -0,0 +1,160 @@
+# ImVs-34 (kronecker-learnmask) — analiza kowariancji markerów K_C
+
+Notatka / wiadomość do grupy. Panel hn (40 markerów); danenberg i hoch-rna analogicznie.
+Skrypt: `dump_marker_covariance.py --panel <dataset>`. Figury w `logs/marker_covariance_ImVs-34_*.png`.
+
+---
+
+Cześć,
+
+Podzielę się wynikiem małej analizy modelu ImVs-34 (wariant kronecker-learnmask). Chciałem sprawdzić, czy ten model nauczył się jakiejś ciekawej reprezentacji niepewności, której nie ma zwykły GP. Kalibracja marginalna (Pearson log-MSE vs log-sigma oraz log-sigma vs log-MAE) wychodziła praktycznie identyczna jak w bazowym GP, więc zajrzałem bezpośrednio w to, co ten wariant realnie dodaje.
+
+Krótko o co chodzi, dla tych co nie siedzieli w części GP: w losie GP, oprócz kowariancji przestrzennej po pikselach, jest dodatkowo kowariancja po markerach — macierz K_C. Powstaje ona z embeddingów markerów z hyperkernela (rzutowanych liniowo i znormalizowanych): K_C = E·Eᵀ. Bazowy model traktuje markery jako niezależne, czyli de facto K_C = I. Ważne: K_C nie zależy od obrazka — to czysty lookup po embeddingach markerów, więc można ją policzyć raz i po prostu obejrzeć.
+
+W załączniku K_C dla panelu hn (40 markerów), dwa panele:
+- lewy: pełna K_C,
+- prawy: K_C po odjęciu dominującej wspólnej składowej („residual").
+
+Kolor = korelacja między dwoma markerami w przestrzeni embeddingów (czerwony dodatni, niebieski ujemny). Markery na obu osiach są ułożone tak samo — uporządkowane przez klasteryzację panelu residualnego, żeby podobne markery leżały obok siebie i tworzyły bloki (to tylko kolejność osi, nie zmienia wartości).
+
+Co widać:
+- Lewy panel jest głównie czerwony — K_C jest zdominowana przez jedną wspólną składową (pierwszy wektor własny to ~55% macierzy). Wszystko koreluje dodatnio, co jest bardzo blisko tego, co dostalibyśmy w ogóle bez kowariancji markerów.
+- Dopiero prawy panel (po odjęciu tej składowej) pokazuje właściwą strukturę: duży blok limfoidalny (FOXP3, PD1, CD27, ICOS, LAG3, CD20 — wszystkie mocno dodatnio) przeciwstawiony blokowi mieloidalnemu (CD11c, CD16, MPO; plus cl.PARP), który jest niebieski względem limfoidalnego. Czyli model ustawił limfoidalne vs mieloidalne na jednej osi — sensowna biologicznie struktura, której bazowy model (K_C = I) w ogóle nie jest w stanie wyrazić.
+- Sanity check: DNA1 i DNA2 lądują jako osobna para (~+1), czyli dwa barwienia jądrowe rozpoznane jako praktycznie identyczne. Markery housekeeping/jądrowe (Histone H3, Ki67, SMA, B2M) są blade, niezależne od osi immunologicznej — co też ma sens.
+
+Sprawdziłem też, skąd ta struktura pochodzi, i to jest ciekawe: **nie jest odziedziczona z rekonstrukcji**. Surowe embeddingi hyperkernela (te same, które ma model bazowy) są nieustrukturyzowane — niemal pełnorzędowe i wzajemnie prawie ortogonalne, korelacje ~±0.03, żadnych bloków (rekonstrukcja pcha embeddingi ku odrębnym filtrom per marker, a nie ku grupowaniu). Cała struktura limfoidalna/mieloidalna powstaje dopiero w warstwie projekcji (`embedding_projection`), która (a) w modelu bazowym w ogóle nie istnieje i (b) jest trenowana wyłącznie przez loss K_C. Korelacja między strukturą surowych a rzutowanych embeddingów to ~0.08, czyli praktycznie zero. Innymi słowy: to grupowanie jest realną „zasługą" kowariancji markerów, a nie efektem ubocznym rekonstrukcji.
+
+Danenberg i hoch-rna wyglądają analogicznie: danenberg dokłada parę stromalną FSP1–Podoplanin, a hoch-rna to panel RNA, więc rzadkie chemokiny + kontrola DapB zlewają się w jedną grupę.
+
+Wniosek: K_C faktycznie nauczyła się nietrywialnej, biologicznie sensownej struktury po markerach — i to struktury specyficznej dla mechanizmu kowariancji markerów, nie czegoś, co model bazowy też by miał. ALE dwie rzeczy tłumaczą, czemu nie widać tego w naszej kalibracji:
+1. Ta struktura jest drugorzędna — dominuje wspólny „globalny" komponent, który działa niemal jak brak kowariancji markerów.
+2. Co ważniejsze: K_C wchodzi tylko do losa (łączny log-likelihood po pikselach × markerach), a niepewność, którą raportujemy i kalibrujemy, to marginalna wariancja z głowicy dekodera (logvar), a nie wariancja a posteriori GP. Czyli nasze wykresy kalibracyjne strukturalnie nie mogą tego „zobaczyć" — dlatego wychodzą identyczne jak baseline.
+
+Gdybyśmy chcieli pokazać efekt K_C na samej niepewności, trzeba by albo policzyć wariancję predykcyjną GP (która realnie używa K_C), albo zrobić test łączny — np. czy markery, które K_C grupuje razem, mają skorelowane błędy w leave-one-out.
+
+Skrypt (dump_marker_covariance.py) jest w repo, liczy się na CPU w kilka sekund, przyjmuje --panel <dataset>. Dajcie znać co myślicie.
+
+---
+
+_Uwaga do potwierdzenia: fragment o pochodzeniu struktury opiera się na dowodzie pośrednim (surowe embeddingi samego ImVs-34). Twarde potwierdzenie = ta sama analiza na checkpoincie modelu bazowego (GP bez marker covariance)._
+
+---
+
+## Test: czy K_C przewiduje skorelowane błędy LOO?
+
+Skrypt `test_kc_error_corr.py` (149 rekonstrukcji LOO, panel hn, 40 markerów). Liczy
+empiryczną korelację map residuów (recon − target) między markerami, uśrednioną po
+obrazach, i porównuje ją z K_C (pełnym i residualnym), z testem permutacyjnym.
+
+**Wynik: NIE — K_C nie przewiduje skorelowanych błędów LOO.**
+
+```
+corr(K_C pełne,    korelacja-błędów pełna)     = -0.016   (praktycznie zero)
+corr(K_C residual, korelacja-błędów residual)  = -0.194   (perm p = 0.025; null |r| max 0.083)
+```
+
+Pełne K_C: brak związku. Residualne K_C: słaby, istotny, ale UJEMNY — markery grupowane
+przez K_C mają odrobinę *mniej* skorelowane błędy, odwrotnie niż hipoteza „K_C łapie
+kowariancję błędów". Parami:
+
+| para | K_C_resid | błąd_resid |
+|---|---|---|
+| DNA1 – DNA2 | +0.999 | −0.118 |
+| CD163 – CD206 | +0.351 | −0.080 |
+| CD163 – CD68 | +0.273 | −0.073 |
+| CD14 – CD163 | +0.176 | −0.162 |
+
+**Mechanizm (DNA1/DNA2):** K_C = +0.999, bo to prawie identyczne barwienia. Ale w LOO,
+maskując DNA1, model ma DNA2 na wejściu → odtwarza DNA1 kopiując DNA2 → mały błąd (i
+odwrotnie). Czyli K_C mierzy **podobieństwo/redundancję** markerów, a nie kowariancję
+błędów; redundantne markery łatwo zaimputować z siebie → mały, zdekorelowany błąd. K_C
+wiąże się więc raczej z **wielkością** błędu (grupa → niski błąd) niż z jego korelacją.
+
+**Wniosek:** K_C to ciekawa wyuczona reprezentacja relacji markerów (biologicznie sensowna,
+CD45RA/CD45RO ≈ 0.30), ale **nie działa jako operacyjny predyktor łącznego zachowania
+błędów / skorelowanej niepewności**. „Ciekawa reprezentacja" ≠ „ciekawa niepewność"
+w sensie mierzalnym na wyjściu — spójne z tym, że raportowana niepewność nie płynie z K_C.
+
+Naturalny następny test: czy markery z wieloma sąsiadami w K_C rekonstruują się lepiej
+w LOO (redundancja → niższy MSE) — to wielkość, z którą K_C faktycznie się wiąże.
+
+---
+
+## Test: czy redundancja w K_C przewiduje niższe MSE w LOO?
+
+Skrypt `test_kc_redundancy_mse.py`. Per marker liczy „redundancję" z K_C i koreluje ją ze
+średnim MSE per marker z CSV LOO (149 obrazów, 40 markerów). Jeden punkt korelacji = jeden
+marker (n = 40). Pearson (recon–target, niezależny od skali) dodany jako kontrola confoundu.
+
+**Wynik: TAK na surowym MSE — ale ⚠ patrz sekcja niżej: na metryce niezależnej od skali
+efekt znika (był confoundem skali).** Markery z bliskim „bliźniakiem" w K_C mają niższe surowe MSE.
+
+```
+score                             vs MSE:  Spearman   Pearson
+kc_max  (najlepszy bliźniak)               -0.878     -0.771    <- najsilniejszy
+kc_nn05 (# sąsiadów > 0.5)                 -0.684     -0.535
+kc_mean (średnie podobieństwo)             -0.480     -0.245
+```
+
+Najsilniejszy predyktor to `kc_max` — do imputacji zamaskowanego markera wystarczy jeden
+bardzo podobny marker na wejściu, z którego można „skopiować".
+
+Ranking:
+- redundantne (kc_max ≈ 0.99): CD4, ICOS, PD1, CD27, FOXP3, CD3, LAG3, CD20 → MSE 0.0006–0.0036
+- unikalne (kc_max ≈ 0, 0 sąsiadów): Ki67 (0.032), CD15 (0.016), Ecad (0.019) → najtrudniejsze (wyjątek: SMA)
+- niuans: MPO/cl.PARP/CD16/CD11c mają ujemne kc_mean (anty-limfoidalne), ale kc_max ≈ 0.98
+  (bliźniacy we własnym klastrze mieloidalnym) → niskie MSE. Liczy się posiadanie *jakiegokolwiek*
+  bliźniaka (kc_max), nie ogólne podobieństwo do panelu (kc_mean).
+
+**Zastrzeżenie (confound intensywności):** związek silny z MSE, ale słaby z Pearsonem
+(kc_max vs pearson: Spearman −0.12). MSE zależy od skali, więc część efektu to fakt, że
+redundantne markery immunologiczne bywają niżej-sygnałowe. Ale mechanizm bliźniaka jest realny
+i niezależny od skali (DNA1/DNA2 i klaster mieloidalny są jasne, a mimo to mają niskie MSE).
+
+## Test: kontrola confoundu — metryka niezależna od skali
+
+Skrypt `test_kc_redundancy_nmse.py`. Zamiast surowego MSE używa NMSE = MSE/Var(target)
+(= 1 − R²) oraz Pearsona(recon, target), liczonych z NPZ — obie niezależne od dynamiki markera.
+
+**Wynik: efekt redundancji ZNIKA.** Poprzednie −0.88 (kc_max vs MSE) było prawie w całości
+confoundem skali.
+
+```
+score     vs        Spearman   (chcemy)
+kc_max    NMSE        +0.008       -      -> zero
+kc_max    R^2         -0.008       +      -> zero
+kc_max    pearson     -0.153       +      -> słabo, zły kierunek
+kc_nn05   pearson     -0.309       +      -> słabo, zły kierunek
+```
+
+| marker | kc_max | R² | pearson |
+|---|---|---|---|
+| DNA1 / DNA2 | 0.999 | 0.97 | 0.99 |
+| PD1 | 1.000 | 0.11 | 0.50 |
+| LAG3 | 0.999 | 0.08 | 0.37 |
+| cl.PARP | 0.999 | 0.01 | 0.21 |
+| CD14 (unikalny) | 0.214 | 0.60 | 0.83 |
+| HLADR (unikalny) | 0.136 | 0.54 | 0.81 |
+| Ecad (unikalny) | 0.097 | 0.49 | 0.79 |
+
+- Tylko prawdziwe duplikaty działają: DNA1/DNA2 (kc_max ≈ 1 **i** R² ≈ 0.97).
+- Reszta klastra limfoidalnego (PD1, LAG3, cl.PARP): kc_max ≈ 1, ale R² 0.01–0.11 — podobieństwo
+  embeddingów ≠ kopiowalność pikseli.
+- Unikalne markery (Ecad, HLADR, CD14) bywają lepiej odtwarzalne niż redundantne limfoidalne.
+
+Redundantne w K_C to po prostu rzadkie, niskosygnałowe markery immunologiczne → małe MSE
+mechanicznie (mała wariancja), a nie „łatwe do zaimputowania".
+
+## Spięcie wszystkich testów — operacyjne znaczenie K_C (wersja finalna)
+
+1. K_C uczy się sensownej biologicznie **struktury podobieństwa markerów** (limfoidalne/mieloidalne;
+   CD45RA/CD45RO ≈ 0.30, zgodne z paperem ImmuVis). ✓
+2. **Nie** przewiduje skorelowanych błędów LOO (r ≈ 0 / słabo ujemne). ✗
+3. **Nie** przewiduje jakości rekonstrukcji na metryce niezależnej od skali (NMSE/R²/Pearson ≈ 0);
+   pozorny efekt na surowym MSE był confoundem skali. ✗ (teza o „mapie imputowalności" — WYCOFANA)
+
+**Wniosek finalny:** K_C to interpretowalna wyuczona mapa podobieństwa markerów, ale **bez
+wykrywalnego operacyjnego śladu na wyjściu modelu** — ani skorelowanej niepewności, ani realnej
+jakości rekonstrukcji. Ciekawa reprezentacja, która (na razie) nie przekłada się na mierzalny
+efekt w predykcjach. Jedyny czysty przypadek „podobieństwo → kopiowalność" to dosłowne
+duplikaty (DNA1/DNA2).
diff --git a/pyproject.toml b/pyproject.toml
index b687dd8..e92e017 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -89,4 +89,12 @@ python_version = "3.12"
 warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = false
+check_untyped_defs = true
 ignore_missing_imports = true
+explicit_package_bases = true
+
+[[tool.mypy.overrides]]
+# multiplex_model.utils is fully typed — enforce it stays that way
+module = "multiplex_model.utils.*"
+disallow_untyped_defs = true
+check_untyped_defs = true
diff --git a/run_embed.py b/run_embed.py
new file mode 100644
index 0000000..90477ce
--- /dev/null
+++ b/run_embed.py
@@ -0,0 +1,256 @@
+import os
+import numpy as np
+import torch
+from ruamel.yaml import YAML
+from tqdm.auto import tqdm
+from torch.utils.data import DataLoader
+import pandas as pd
+from glob import glob
+
+
+from multiplex_model.data import DatasetFromTIFF, PanelBatchSampler
+from multiplex_model.modules.immuvis import MultiplexAutoencoder
+from multiplex_model.utils.configuration import EncoderConfig, DecoderConfig
+
+
+models_path = "/raid_encrypted/immucan/models"
+embeddings_path = "/raid_encrypted/immucan/embeddings"
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+panel_config = "/home/mzmyslowski/marcin_multiplex/configs/all_panels_config.yaml"
+tokenizer_config = "/home/mzmyslowski/marcin_multiplex/configs/all_markers_tokenizer.yaml"
+
+PATCH_SIZE = 128
+BATCH_SIZE = 1
+NUM_WORKERS = 8
+SAVE_EVERY = 200  # Save intermediate results every N images
+
+print(f"Using device: {DEVICE}")
+
+
+# Load configuration
+yaml = YAML(typ="safe")
+with open(panel_config, "r") as f:
+    panel_config_dict = yaml.load(f)
+
+panel_config_dict['datasets'] = ['hn']
+# Override paths to raw TIFFs (all_panels_config.yaml points to pre-patched .npy used for training)
+panel_config_dict['paths']['train'] = '/raid_encrypted/immucan/immuvis_split/train'
+panel_config_dict['paths']['test'] = '/raid_encrypted/immucan/immuvis_split/test'
+
+# Load tokenizer
+with open(tokenizer_config, "r") as f:
+    TOKENIZER = yaml.load(f)
+
+# Create inverse tokenizer for channel names
+INV_TOKENIZER = {v: k for k, v in TOKENIZER.items()}
+num_channels = len(TOKENIZER)
+
+print(f"Number of channels: {num_channels}")
+print(f"Sample markers: {list(TOKENIZER.keys())[:5]}")
+
+# Create train and test datasets
+train_dataset = DatasetFromTIFF(
+    panels_config=panel_config_dict,
+    split='train',
+    marker_tokenizer=TOKENIZER,
+    transform=None,
+    use_median_denoising=False,
+    use_butterworth_filter=True,
+    use_minmax_normalization=False,
+    use_global_clip_limits=False,
+    use_clip_normalization=True,
+)
+
+test_dataset = DatasetFromTIFF(
+    panels_config=panel_config_dict,
+    split='test',
+    marker_tokenizer=TOKENIZER,
+    transform=None,
+    use_median_denoising=False,
+    use_butterworth_filter=True,
+    use_minmax_normalization=False,
+    use_global_clip_limits=False,
+    use_clip_normalization=True,
+)
+
+train_batch_sampler = PanelBatchSampler(train_dataset, BATCH_SIZE, shuffle=False)
+test_batch_sampler = PanelBatchSampler(test_dataset, BATCH_SIZE, shuffle=False)
+
+train_dataloader = DataLoader(train_dataset, batch_sampler=train_batch_sampler, num_workers=NUM_WORKERS)
+test_dataloader = DataLoader(test_dataset, batch_sampler=test_batch_sampler, num_workers=NUM_WORKERS)
+
+print(f"Train dataset size: {len(train_dataset)} images")
+print(f"Test dataset size: {len(test_dataset)} images")
+
+
+def get_all_patches(img, patch_size: int = 128):
+    """Extract all non-overlapping patches from an image."""
+    H, W = img.shape[2:]
+    i0, j0 = 0, 0
+    i1, j1 = patch_size, patch_size
+    patches = []
+    coords = []
+    
+    while True:
+        while True:
+            patch = img[:, :, i0:i1, j0:j1]
+            patches.append(patch)
+            coords.append([(i0, j0), (i1, j1)])
+            
+            j1 += patch_size
+            if j1 > W:
+                break
+            j0 = j1 - patch_size
+        
+        i1 += patch_size
+        if i1 > H:
+            break
+        i0 = i1 - patch_size
+        j0 = 0
+        j1 = patch_size
+
+    return patches, coords
+
+
+def embed_images(
+    model,
+    dataloader, 
+    device, 
+    patch_size=128,
+    outpath=None,
+    split_name=None,
+    model_prefix=None,
+    save_every=200
+):
+    """Embed all images in the dataloader by extracting patches and encoding them."""
+    model.eval()
+
+    embeddings = []
+    metadata = []
+    batch_idx = 0
+    
+    for i, (img, channel_ids, panel_idx, img_path) in enumerate(tqdm(dataloader, desc=f"Embedding {split_name} images")):
+        B, C, H, W = img.shape
+        if H < patch_size or W < patch_size:
+            print(f'Image is smaller than patch size: {img.shape} at {img_path[0]}')
+            continue
+        
+        channel_ids = channel_ids.to(device)
+
+        for patch, (coords0, coords1) in zip(*get_all_patches(img, patch_size)):
+            patch = patch.to(torch.float32).to(device)
+            metadata.append((os.path.realpath(img_path[0]), panel_idx[0], coords0, coords1))
+
+            with torch.no_grad():
+                latent = model.encode(patch, channel_ids)['output']
+                embeddings.append(latent.cpu().numpy().squeeze(0))
+
+        if (i + 1) % save_every == 0:
+            print(f'Processed {i + 1} images, saving batch {batch_idx}...')
+            # Save intermediate results
+            if outpath:
+                embeddings_array = np.stack(embeddings)
+                np.save(
+                    os.path.join(outpath, f'{model_prefix}_{split_name}_image_patches_embeddings_batch_{batch_idx}.npy'), 
+                    embeddings_array
+                )
+                pd.DataFrame(
+                    metadata, 
+                    columns=['img_path', 'panel', 'coords0', 'coords1']
+                ).to_csv(
+                    os.path.join(outpath, f'{model_prefix}_{split_name}_image_patches_metadata_batch_{batch_idx}.csv'), 
+                    index=False
+                )
+
+                embeddings = []
+                metadata = []
+                batch_idx += 1
+
+    # Save remaining embeddings
+    if embeddings:
+        embeddings_array = np.stack(embeddings)
+        np.save(
+            os.path.join(outpath, f'{model_prefix}_{split_name}_image_patches_embeddings_batch_{batch_idx}.npy'), 
+            embeddings_array
+        )
+        pd.DataFrame(
+            metadata, 
+            columns=['img_path', 'panel', 'coords0', 'coords1']
+        ).to_csv(
+            os.path.join(outpath, f'{model_prefix}_{split_name}_image_patches_metadata_batch_{batch_idx}.csv'), 
+            index=False
+        )
+        print(f'Saved final batch {batch_idx}')
+    
+    print(f'Finished embedding {split_name} images!')
+
+MODEL_WEIGHTS_PATH = "/home/mzmyslowski/marcin_multiplex/checkpoints/last_checkpoint-ImVs-25.pth"
+MODEL_CONFIG_PATH = "/home/mzmyslowski/marcin_multiplex/train_masked_gp_marker_config_resume3.yaml"
+
+for model_name in [MODEL_WEIGHTS_PATH]:
+    model_checkpoint = os.path.basename(model_name)
+    model_prefix = model_checkpoint.replace('.pth', '')
+
+    print(f"\n{'='*80}")
+    print(f"Processing model: {model_checkpoint}")
+    print(f"{'='*80}")
+
+    with open(MODEL_CONFIG_PATH, "r") as f:
+        model_config_dict = yaml.load(f)
+
+    encoder_config = EncoderConfig(**model_config_dict["encoder"])
+    decoder_config = DecoderConfig(**model_config_dict["decoder"])
+
+    # Initialize model
+    model = MultiplexAutoencoder(
+        num_channels=num_channels,
+        encoder_config=encoder_config.model_dump(),
+        decoder_config=decoder_config.model_dump(),
+    ).to(DEVICE)
+
+    # Load model weights
+    print(f"Loading model weights from: {MODEL_WEIGHTS_PATH}")
+    checkpoint = torch.load(MODEL_WEIGHTS_PATH, map_location="cpu")
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+
+    print("Model loaded successfully!")
+
+    
+    # Embed test images
+    print("\nEmbedding test dataset...")
+    embed_images(
+        model,
+        test_dataloader, 
+        DEVICE, 
+        patch_size=PATCH_SIZE,
+        outpath=embeddings_path,
+        split_name='test',
+        model_prefix=model_prefix,
+        save_every=SAVE_EVERY
+    )
+    
+    # Embed train images
+    print("\nEmbedding train dataset...")
+    embed_images(
+        model,
+        train_dataloader, 
+        DEVICE, 
+        patch_size=PATCH_SIZE,
+        outpath=embeddings_path,
+        split_name='train',
+        model_prefix=model_prefix,
+        save_every=SAVE_EVERY
+    )
+    
+    print(f"\nCompleted embedding for {model_checkpoint}")
+    
+    # Clean up to free memory
+    del model
+    del checkpoint
+    torch.cuda.empty_cache()
+
+print(f"\n{'='*80}")
+print("All models processed successfully!")
+print(f"{'='*80}")
diff --git a/run_validation_leave_one_out.py b/run_validation_leave_one_out.py
new file mode 100644
index 0000000..77c58d4
--- /dev/null
+++ b/run_validation_leave_one_out.py
@@ -0,0 +1,293 @@
+import argparse
+import json
+import os
+from glob import glob
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from ruamel.yaml import YAML
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+
+from multiplex_model.data import DatasetFromTIFF, TestCrop
+from multiplex_model.modules.immuvis import MultiplexAutoencoder
+from multiplex_model.utils.configuration import DecoderConfig, EncoderConfig
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run leave-one-out validation (mask one channel at a time)."
+    )
+    parser.add_argument(
+        "--versions",
+        type=int,
+        nargs="+",
+        default=list(range(0, 19)),
+        help="Model versions to evaluate (e.g. --versions 14 15).",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default=None,
+        help="Direct path to a single model checkpoint (bypasses --versions glob).",
+    )
+    parser.add_argument(
+        "--model-config",
+        type=str,
+        default=None,
+        help="Direct path to model config YAML (required with --checkpoint).",
+    )
+    parser.add_argument(
+        "--max-images",
+        type=int,
+        default=None,
+        help="Maximum number of test images to evaluate per model (default: all).",
+    )
+    parser.add_argument(
+        "--models-path",
+        default="/raid_encrypted/immucan/models",
+        help="Path to model checkpoints and configs.",
+    )
+    parser.add_argument(
+        "--results-dir",
+        default="/raid_encrypted/immucan/results/with_reconstructs",
+        help="Where to save CSV outputs.",
+    )
+    parser.add_argument(
+        "--recon-dir",
+        default="/raid_encrypted/immucan/recons/immuvis-beta",
+        help="Where to save leave-one-out reconstructions (npz).",
+    )
+    parser.add_argument(
+        "--save-reconstructions",
+        action="store_true",
+        help="Save leave-one-out reconstructions to NPZ files.",
+    )
+    parser.add_argument(
+        "--panel-config",
+        default="/home/mzmyslowski/marcin_multiplex/configs/all_panels_config.yaml",
+    )
+    parser.add_argument(
+        "--tokenizer-config",
+        default="/home/mzmyslowski/marcin_multiplex/configs/all_markers_tokenizer.yaml",
+    )
+    return parser.parse_args()
+
+
+def create_leave_one_out_batch(
+    img: torch.Tensor, channel_ids: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Create leave-one-out batch for a single image.
+
+    Args:
+        img: [C, H, W]
+        channel_ids: [C]
+
+    Returns:
+        masked_img: [C, C-1, H, W]
+        active_channel_ids: [C, C-1]
+        output_channel_ids: [C, 1]
+        masked_indices: [C]
+    """
+    num_channels, height, width = img.shape
+    keep_mask = ~torch.eye(num_channels, dtype=torch.bool, device=img.device)
+
+    img_expand = img.unsqueeze(0).expand(num_channels, -1, -1, -1)
+    masked_img = img_expand[keep_mask].view(num_channels, num_channels - 1, height, width)
+
+    channel_ids_expand = channel_ids.unsqueeze(0).expand(num_channels, -1)
+    active_channel_ids = channel_ids_expand[keep_mask].view(num_channels, num_channels - 1)
+
+    output_channel_ids = channel_ids.view(num_channels, 1)
+    masked_indices = torch.arange(num_channels, device=img.device)
+
+    return masked_img, active_channel_ids, output_channel_ids, masked_indices
+
+
+def main() -> None:
+    args = parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+
+    yaml = YAML(typ="safe")
+    with open(args.panel_config, "r") as f:
+        panel_config_dict = yaml.load(f)
+
+    panel_config_dict["datasets"] = ["hn"]
+
+    with open(args.tokenizer_config, "r") as f:
+        tokenizer = yaml.load(f)
+
+    inv_tokenizer = {v: k for k, v in tokenizer.items()}
+    num_channels = len(tokenizer)
+
+    print(f"Number of channels: {num_channels}")
+    print(f"Sample markers: {list(tokenizer.keys())[:5]}")
+
+    test_transform = TestCrop(128)
+    test_dataset = DatasetFromTIFF(
+        panels_config=panel_config_dict,
+        split="test",
+        marker_tokenizer=tokenizer,
+        use_preprocessing=False,
+        use_median_denoising=False,
+        use_butterworth_filter=True,
+        use_minmax_normalization=False,
+        use_clip_normalization=True,
+        file_extension="npy",
+        transform=test_transform,
+    )
+
+    print(f"Test dataset size: {len(test_dataset)} images")
+    dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False)
+
+    if args.checkpoint:
+        if not args.model_config:
+            raise ValueError("--model-config is required when using --checkpoint")
+        model_entries = [(args.checkpoint, args.model_config)]
+    else:
+        patterns = [f"Immu*-6{v:02d}-beta-*.pth" for v in args.versions]
+        model_files: list[str] = []
+        for pattern in patterns:
+            model_files.extend(glob(f"{args.models_path}/{pattern}"))
+
+        if not model_files:
+            raise FileNotFoundError(
+                f"No model checkpoints found in {args.models_path} for versions {args.versions}."
+            )
+        model_entries = []
+        for mf in sorted(model_files):
+            cfg = f"{args.models_path}/config.{Path(mf).stem}.yaml"
+            model_entries.append((mf, cfg))
+
+    os.makedirs(args.results_dir, exist_ok=True)
+
+    for model_weights_path, model_config_path in model_entries:
+        model_checkpoint = os.path.basename(model_weights_path)
+        model_idx = Path(model_weights_path).stem
+
+        with open(model_config_path, "r") as f:
+            model_config_dict = yaml.load(f)
+
+        encoder_config = EncoderConfig(**model_config_dict["encoder"])
+        decoder_config = DecoderConfig(**model_config_dict["decoder"])
+
+        model = MultiplexAutoencoder(
+            num_channels=num_channels,
+            encoder_config=encoder_config.model_dump(),
+            decoder_config=decoder_config.model_dump(),
+        ).to(device)
+
+        print(f"Loading model weights from: {model_weights_path}")
+        checkpoint = torch.load(model_weights_path, map_location="cpu")
+        model.load_state_dict(checkpoint["model_state_dict"])
+        model.eval()
+
+        all_mse = []
+        all_uncertainties = []
+        all_pearson_r = []
+        all_channel_ids = []
+        all_dataset_names = []
+        all_image_paths = []
+
+        recon_dir = Path(args.recon_dir) / f"immuvis_{model_idx}_loo"
+        if args.save_reconstructions:
+            recon_dir.mkdir(parents=True, exist_ok=True)
+
+        with torch.no_grad():
+            for img_idx, (img, channel_ids, ds_name, img_path) in enumerate(
+                tqdm(dataloader, desc="Leave-one-out validation")
+            ):
+                if args.max_images is not None and img_idx >= args.max_images:
+                    break
+
+                img = img.squeeze(0).to(device, dtype=torch.float32)
+                channel_ids = channel_ids.squeeze(0).to(device, dtype=torch.long)
+
+                (
+                    masked_img,
+                    active_channel_ids,
+                    output_channel_ids,
+                    masked_indices,
+                ) = create_leave_one_out_batch(img=img, channel_ids=channel_ids)
+
+                output = model(masked_img, active_channel_ids, output_channel_ids)[
+                    "output"
+                ]
+
+                mi, logvar = output.unbind(dim=-1)
+                mi = torch.sigmoid(mi).squeeze(1)
+                logvar = logvar.squeeze(1)
+
+                target_channels = img[masked_indices]
+                mse = (mi - target_channels).pow(2).mean(dim=(1, 2))
+
+                mi_mean = mi.mean(dim=(1, 2), keepdim=True)
+                target_mean = target_channels.mean(dim=(1, 2), keepdim=True)
+                pearson_r = ((mi - mi_mean) * (target_channels - target_mean)).mean(
+                    dim=(1, 2)
+                ) / (mi.std(dim=(1, 2)) * target_channels.std(dim=(1, 2)) + 1e-8)
+
+                all_mse.append(mse.flatten().cpu().numpy())
+                all_uncertainties.append(logvar.mean(dim=(1, 2)).flatten().cpu().numpy())
+                all_pearson_r.append(pearson_r.flatten().cpu().numpy())
+                all_channel_ids.append(channel_ids[masked_indices].flatten().cpu().numpy())
+
+                num_observations = masked_indices.numel()
+                all_dataset_names.extend([ds_name[0]] * num_observations)
+                all_image_paths.extend([img_path[0]] * num_observations)
+
+                if args.save_reconstructions:
+                    masked_channel_ids = channel_ids.detach().cpu().numpy()
+                    masked_marker_names = [
+                        inv_tokenizer.get(int(cid), "Unknown")
+                        for cid in masked_channel_ids.tolist()
+                    ]
+                    metadata = {
+                        "image_index": int(img_idx),
+                        "image_path": str(img_path[0]),
+                        "dataset_name": str(ds_name[0]),
+                        "masked_strategy": "leave_one_out",
+                        "num_channels": int(masked_channel_ids.shape[0]),
+                    }
+                    out_path = recon_dir / f"recn-{img_idx:05d}.npz"
+                    np.savez_compressed(
+                        out_path,
+                        recon=mi.detach().cpu().numpy(),
+                        variance=torch.exp(logvar).detach().cpu().numpy(),
+                        target=img.detach().cpu().numpy(),
+                        channel_ids=masked_channel_ids,
+                        marker_names=np.array(masked_marker_names),
+                        masked_channel_ids=masked_channel_ids,
+                        masked_marker_names=np.array(masked_marker_names),
+                        metadata=np.array(json.dumps(metadata)),
+                    )
+
+        mses = np.concatenate(all_mse, axis=0)
+        uncertainties = np.concatenate(all_uncertainties, axis=0)
+        pearson_rs = np.concatenate(all_pearson_r, axis=0)
+        masked_channel_ids = np.concatenate(all_channel_ids, axis=0)
+
+        all_vals = np.stack(
+            [mses, uncertainties, pearson_rs, masked_channel_ids],
+            axis=1,
+        )
+
+        df = pd.DataFrame(all_vals, columns=["mse", "logsigma", "pearson", "Channel_ID"])
+        df["marker"] = df["Channel_ID"].map(lambda x: inv_tokenizer.get(x, "Unknown"))
+        df["masked"] = "leave_one_out"
+        df["masked_count"] = 1
+        df["model"] = f"ImmuViT-{model_idx}"
+        df["dataset_name"] = all_dataset_names
+        df["image_path"] = all_image_paths
+
+        output_file = os.path.join(args.results_dir, f"immuvis_{model_idx}_loo.csv")
+        print(f"Saving results to: {output_file}")
+        df.to_csv(output_file, index=False)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/setup_venv.sh b/setup_venv.sh
new file mode 100755
index 0000000..94a0ed5
--- /dev/null
+++ b/setup_venv.sh
@@ -0,0 +1,19 @@
+#!/bin/bash
+# Run this once on bury/szary to set up the virtual environment
+set -e
+
+# Install uv if not present
+if ! command -v uv &> /dev/null; then
+    curl -LsSf https://astral.sh/uv/install.sh | sh
+    source "$HOME/.local/bin/env"
+fi
+
+export PATH="$HOME/.local/bin:$PATH"
+
+cd "$(dirname "$0")"
+
+uv venv ~/venv
+source ~/venv/bin/activate
+uv pip install -e ".[dev]"
+
+echo "Venv ready at ~/venv"
diff --git a/summary_pl.py b/summary_pl.py
new file mode 100644
index 0000000..b78206a
--- /dev/null
+++ b/summary_pl.py
@@ -0,0 +1,205 @@
+"""Generate PDF summary of KroneckerMarkerCovariance method in Polish."""
+
+from reportlab.lib.pagesizes import A4
+from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
+from reportlab.lib.units import cm
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
+from reportlab.lib.enums import TA_LEFT, TA_CENTER
+from reportlab.lib import colors
+from reportlab.platypus import HRFlowable
+
+
+OUTPUT = "kronecker_marker_summary_pl.pdf"
+
+doc = SimpleDocTemplate(
+    OUTPUT,
+    pagesize=A4,
+    leftMargin=2.5 * cm,
+    rightMargin=2.5 * cm,
+    topMargin=2.5 * cm,
+    bottomMargin=2.5 * cm,
+)
+
+styles = getSampleStyleSheet()
+
+title_style = ParagraphStyle(
+    "Title",
+    parent=styles["Normal"],
+    fontSize=16,
+    fontName="Helvetica-Bold",
+    spaceAfter=6,
+    alignment=TA_CENTER,
+)
+subtitle_style = ParagraphStyle(
+    "Subtitle",
+    parent=styles["Normal"],
+    fontSize=11,
+    fontName="Helvetica",
+    textColor=colors.HexColor("#555555"),
+    spaceAfter=16,
+    alignment=TA_CENTER,
+)
+h2_style = ParagraphStyle(
+    "H2",
+    parent=styles["Normal"],
+    fontSize=12,
+    fontName="Helvetica-Bold",
+    spaceBefore=14,
+    spaceAfter=4,
+    textColor=colors.HexColor("#1a1a2e"),
+)
+body_style = ParagraphStyle(
+    "Body",
+    parent=styles["Normal"],
+    fontSize=10,
+    fontName="Helvetica",
+    leading=15,
+    spaceAfter=6,
+)
+eq_style = ParagraphStyle(
+    "Eq",
+    parent=styles["Normal"],
+    fontSize=10,
+    fontName="Courier",
+    leading=14,
+    leftIndent=20,
+    spaceBefore=4,
+    spaceAfter=4,
+    backColor=colors.HexColor("#f5f5f5"),
+)
+bullet_style = ParagraphStyle(
+    "Bullet",
+    parent=styles["Normal"],
+    fontSize=10,
+    fontName="Helvetica",
+    leading=14,
+    leftIndent=16,
+    spaceAfter=3,
+    bulletIndent=0,
+)
+
+
+def h2(text):
+    return Paragraph(text, h2_style)
+
+
+def body(text):
+    return Paragraph(text, body_style)
+
+
+def eq(text):
+    return Paragraph(text, eq_style)
+
+
+def bullet(text):
+    return Paragraph(f"&bull;  {text}", bullet_style)
+
+
+def gap(n=6):
+    return Spacer(1, n)
+
+
+def hr():
+    return HRFlowable(width="100%", thickness=0.5, color=colors.HexColor("#cccccc"), spaceAfter=6)
+
+
+story = [
+    Paragraph("Kowariancja z Kroneckerem i markerami", title_style),
+    Paragraph("Krótkie podsumowanie metody — Multiplex Image Model", subtitle_style),
+    hr(),
+
+    # --- Problem ---
+    h2("1. Problem"),
+    body(
+        "Rekonstruujemy obraz multiplex złożony z <b>C markerów</b> i <b>H×W pikseli</b>. "
+        "Model predykuje średnią μ i niepewność σ per piksel per marker. "
+        "Celem jest modelowanie <b>korelacji przestrzennych i między markerami</b> — "
+        "nie zakładamy niezależności pikseli."
+    ),
+
+    # --- Struktura kowariancji ---
+    h2("2. Struktura kowariancji"),
+    body(
+        "Definiujemy kowariancję na przestrzeni <b>NC wymiarów</b> "
+        "(N = H·W pikseli, C markerów):"
+    ),
+    gap(),
+    eq("K = (Kx ⊗ Ky) ⊗ K_C  +  U_block · U_block\u1d40  +  \u03b5I"),
+    gap(10),
+
+    body("<b>Kx, Ky ∈ ℝ^{n×n}</b> — jądro Matérna 1D na osiach przestrzennych:"),
+    bullet("Separowalna aproksymacja izotropowego Matérna."),
+    bullet("Kx = Ky — ta sama siatka n punktów równomiernie w [0, 1]."),
+    bullet("ν = 1.5 (raz różniczkowalne), lengthscale = 5.0 (szeroka korelacja)."),
+    gap(4),
+
+    body("<b>K_C ∈ ℝ^{C×C}</b> — kowariancja markerów:"),
+    eq("K_C = E · E\u1d40  +  \u03b4I"),
+    body(
+        "gdzie E ∈ ℝ^{C×D} to znormalizowane rzutowanie embedingów Hyperkernel "
+        "(nn.Linear → normalize po wierszach). "
+        "Normalizacja wierszy sprawia, że K_C jest macierzą korelacji "
+        "(jedynki na diagonali), co ogranicza liczbę uwarunkowania do max C."
+    ),
+    gap(4),
+
+    body("<b>U_block ∈ ℝ^{NC×C}</b> — niskorangowy składnik per piksel (z dekodera)."),
+    body("<b>εI</b> — jitter numeryczny dla stabilności."),
+
+    # --- Obliczenia ---
+    h2("3. Efektywne obliczenia"),
+    body(
+        "Macierz NC×NC <b>nie jest nigdy materializowana</b> "
+        "(dla N = 112², C = 40 miałaby ~2×10\u2077 wymiarów). "
+        "Korzystamy z rozkładów spektralnych:"
+    ),
+    gap(),
+    eq("Kx = V \u039bx V\u1d40,    Ky = V \u039by V\u1d40,    K_C = Vc \u039bC Vc\u1d40"),
+    gap(4),
+    body("Wartości własne złożonego składnika Kroneckerowskiego:"),
+    eq("\u03bbijk = \u03bbx_i · \u03bby_j · \u03bbC_k  +  \u03b5"),
+    gap(4),
+    body(
+        "Rozwiązanie A⁻¹v (gdzie A = (Kx⊗Ky)⊗K_C + εI) sprowadza się do "
+        "<b>sześciu operacji einsum</b>: transformacja do bazy eigenvektorów, "
+        "skalowanie przez 1/λijk, transformacja z powrotem."
+    ),
+    gap(4),
+    body("Człon niskorangowy obsługuje <b>tożsamość Woodbury</b>:"),
+    eq("K\u207b\u00b9e = A\u207b\u00b9e \u2212 A\u207b\u00b9U(I + U\u1d40A\u207b\u00b9U)\u207b\u00b9U\u1d40A\u207b\u00b9e"),
+    eq("log det K = log det A + log det(I + U\u1d40A\u207b\u00b9U)"),
+    gap(4),
+    body("<b>Złożoność:</b>"),
+    bullet("O(n³) — raz przy inicjalizacji (rozkład Kx, Ky)."),
+    bullet("O(C³) — per batch (rozkład K_C z embedingów markerów)."),
+    bullet("O(n²·C) — per obraz (solver A⁻¹, Woodbury)."),
+
+    # --- Stabilność ---
+    h2("4. Stabilność numeryczna"),
+    body("Dwa kluczowe zabiegi niezbędne do zbieżności:"),
+    bullet(
+        "<b>Normalizacja wierszy E</b>: bez niej K_C ma liczbę uwarunkowania ~10⁵–10⁶ "
+        "→ GP NLL = nan od pierwszej epoki."
+    ),
+    bullet(
+        "<b>float64 dla eigh(K_C)</b>: gdy C > D (wymiar projekcji), "
+        "K_C ma C−D powtarzających się wartości własnych dokładnie równych δ. "
+        "LAPACK w float32 nie zbiega — rzutowanie do float64 i z powrotem rozwiązuje problem."
+    ),
+
+    # --- Wyniki ---
+    h2("5. Wyniki wstępne"),
+    body(
+        "Po 24 epokach (ImVs-19, batch_size=8, λ_GP=0.1, lengthscale=5.0):"
+    ),
+    bullet("MAE: 0.127 → 0.030 (postępująca poprawa rekonstrukcji)."),
+    bullet(
+        "Pearson ρ(MAE, Var) ≈ 0.90–0.95 — model dobrze kalibruje niepewność: "
+        "wysoka predykowana wariancja koreluje z wysokim błędem rekonstrukcji."
+    ),
+    bullet("GP NLL stale ujemny i malejący — kowariancja markerów aktywnie się uczy."),
+    bullet("Liczba uwarunkowania K_C: ~1000–1500, min eigval = 0.01 — stabilna."),
+]
+
+doc.build(story)
+print(f"Saved: {OUTPUT}")
diff --git a/test_kc_error_corr.py b/test_kc_error_corr.py
new file mode 100644
index 0000000..0f5c33a
--- /dev/null
+++ b/test_kc_error_corr.py
@@ -0,0 +1,88 @@
+"""Does the learned K_C predict which markers have correlated leave-one-out errors?
+
+K_C models the cross-marker covariance of pixel residuals. Direct test: from the saved
+LOO reconstructions, compute the empirical cross-marker correlation of residual maps
+(recon - target), averaged over images, and correlate it against K_C — both full and
+with the dominant shared component removed. A permutation test over marker labels gives
+a null for the residual (structure-specific) comparison.
+"""
+
+import glob
+
+import numpy as np
+
+KC_NPZ = "/home/mzmyslowski/marcin_multiplex/logs/marker_covariance_ImVs-34_hn.npz"
+RECON_DIR = "/raid_encrypted/immucan/recons/immuvis-gp/immuvis_last_checkpoint-ImVs-34_loo"
+
+
+def residualize(m: np.ndarray) -> np.ndarray:
+    """Remove the leading (shared) eigen-component of a symmetric matrix."""
+    w, v = np.linalg.eigh(m)
+    return m - w[-1] * np.outer(v[:, -1], v[:, -1])
+
+
+def empirical_error_correlation(files: list[str], ref_names: list[str]) -> np.ndarray:
+    """Average per-image cross-marker correlation of residual maps (recon - target)."""
+    c = len(ref_names)
+    acc = np.zeros((c, c))
+    cnt = np.zeros((c, c))
+    for f in files:
+        d = np.load(f, allow_pickle=True)
+        names = list(d["marker_names"])
+        idx = [names.index(m) for m in ref_names]
+        resid = (d["recon"].astype(np.float64) - d["target"].astype(np.float64))[idx]
+        resid = resid.reshape(c, -1)
+        corr = np.corrcoef(resid)  # nan where a residual map is constant
+        good = np.isfinite(corr)
+        acc[good] += corr[good]
+        cnt[good] += 1
+    return acc / np.maximum(cnt, 1)
+
+
+def main() -> None:
+    kc = np.load(KC_NPZ, allow_pickle=True)
+    kc_names = list(kc["marker_names"])
+    k_full = kc["k_c"]
+
+    files = sorted(glob.glob(RECON_DIR + "/*.npz"))
+    ref_names = [m for m in kc_names if m in list(np.load(files[0], allow_pickle=True)["marker_names"])]
+    print(f"images: {len(files)} | markers aligned: {len(ref_names)}")
+
+    err = empirical_error_correlation(files, ref_names)
+
+    # Align K_C to the same marker order.
+    ik = [kc_names.index(m) for m in ref_names]
+    kc_a = k_full[np.ix_(ik, ik)]
+    off = ~np.eye(len(ref_names), dtype=bool)
+
+    r_full = np.corrcoef(kc_a[off], err[off])[0, 1]
+    kc_r, err_r = residualize(kc_a), residualize(err)
+    r_resid = np.corrcoef(kc_r[off], err_r[off])[0, 1]
+
+    # Permutation null for the residual comparison: shuffle marker labels of err.
+    rng_orders = [np.roll(np.arange(len(ref_names)), s) for s in range(1, len(ref_names))]
+    perm = []
+    for o in rng_orders:
+        er = err_r[np.ix_(o, o)]
+        perm.append(np.corrcoef(kc_r[off], er[off])[0, 1])
+    perm = np.array(perm)
+    p_val = (np.sum(np.abs(perm) >= abs(r_resid)) + 1) / (len(perm) + 1)
+
+    print(f"\ncorr(K_C full,     error-corr full)     = {r_full:+.3f}")
+    print(f"corr(K_C residual, error-corr residual) = {r_resid:+.3f}   (perm p = {p_val:.3f}, null |r| max {np.abs(perm).max():.3f})")
+
+    # Interpretable pairs: strongest residual K_C pairs and whether errors track them.
+    names = ref_names
+    pairs = [(names[i], names[j], kc_r[i, j], err_r[i, j]) for i in range(len(names)) for j in range(i + 1, len(names))]
+    pairs.sort(key=lambda p: p[2], reverse=True)
+    print("\nTop 12 K_C-grouped pairs -> their empirical LOO error correlation:")
+    print(f"  {'pair':<24} {'K_C_resid':>10} {'err_resid':>10}")
+    for a, b, kv, ev in pairs[:12]:
+        print(f"  {a+' - '+b:<24} {kv:>+10.3f} {ev:>+10.3f}")
+    print("\nBottom 6 (K_C anti-grouped) -> error correlation:")
+    for a, b, kv, ev in pairs[-6:]:
+        print(f"  {a+' - '+b:<24} {kv:>+10.3f} {ev:>+10.3f}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/test_kc_redundancy_mse.py b/test_kc_redundancy_mse.py
new file mode 100644
index 0000000..f6efd81
--- /dev/null
+++ b/test_kc_redundancy_mse.py
@@ -0,0 +1,60 @@
+"""Do markers with more K_C neighbours reconstruct better in LOO? (redundancy -> lower MSE)
+
+K_C measures marker similarity. Hypothesis: a marker that is similar to others in the panel
+is easy to impute from them when masked -> lower LOO error. We score each marker's redundancy
+from K_C (mean and max similarity to the rest of the panel) and correlate it with per-marker
+mean MSE from the LOO CSV. Pearson (scale-invariant recon quality) is reported alongside MSE
+to guard against the marker-intensity confound.
+"""
+
+import numpy as np
+import pandas as pd
+from scipy.stats import pearsonr, spearmanr
+
+KC_NPZ = "/home/mzmyslowski/marcin_multiplex/logs/marker_covariance_ImVs-34_hn.npz"
+CSV = "/raid_encrypted/immucan/results/with_reconstructs/immuvis_last_checkpoint-ImVs-34_loo.csv"
+
+
+def main() -> None:
+    kc = np.load(KC_NPZ, allow_pickle=True)
+    names = list(kc["marker_names"])
+    k = kc["k_c"].copy()
+    np.fill_diagonal(k, np.nan)  # ignore self
+
+    df = pd.read_csv(CSV)
+    per_marker = df.groupby("marker").agg(mse=("mse", "mean"), pearson=("pearson", "mean"), n=("mse", "size"))
+
+    rows = []
+    for i, m in enumerate(names):
+        if m not in per_marker.index:
+            continue
+        row = k[i]
+        rows.append(
+            {
+                "marker": m,
+                "kc_mean": np.nanmean(row),          # overall similarity to panel
+                "kc_max": np.nanmax(row),            # best single "twin"
+                "kc_nn05": int(np.nansum(row > 0.5)),  # count of strong neighbours
+                "mse": per_marker.loc[m, "mse"],
+                "pearson": per_marker.loc[m, "pearson"],
+            }
+        )
+    t = pd.DataFrame(rows)
+    print(f"markers matched: {len(t)}")
+
+    print("\nCorrelation of K_C redundancy score vs per-marker LOO metric:")
+    print(f"  {'score':<10} {'vs':<8} {'Spearman':>10} {'Pearson':>10}")
+    for score in ["kc_mean", "kc_max", "kc_nn05"]:
+        for target, sign in [("mse", "(want -)"), ("pearson", "(want +)")]:
+            rho = spearmanr(t[score], t[target]).correlation
+            r = pearsonr(t[score], t[target])[0]
+            print(f"  {score:<10} {target:<8} {rho:>+10.3f} {r:>+10.3f}   {sign}")
+
+    print("\nMost redundant markers (high kc_mean) — do they reconstruct better?")
+    print(t.sort_values("kc_mean", ascending=False)[["marker", "kc_mean", "kc_max", "kc_nn05", "mse", "pearson"]].head(8).to_string(index=False))
+    print("\nLeast redundant markers (low kc_mean):")
+    print(t.sort_values("kc_mean")[["marker", "kc_mean", "kc_max", "kc_nn05", "mse", "pearson"]].head(8).to_string(index=False))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/test_kc_redundancy_nmse.py b/test_kc_redundancy_nmse.py
new file mode 100644
index 0000000..0eadcb9
--- /dev/null
+++ b/test_kc_redundancy_nmse.py
@@ -0,0 +1,83 @@
+"""Redundancy vs SCALE-INDEPENDENT reconstruction quality in LOO.
+
+Separates 'easy because redundant' from 'easy because bright'. Instead of raw MSE we use
+per-(image, marker) NMSE = MSE / Var(target) (= 1 - R^2), plus Pearson(recon, target) —
+both invariant to the marker's dynamic range. Computed from the saved LOO NPZ reconstructions
+(recon, target), aggregated per marker with the median (robust), then correlated with K_C
+redundancy scores.
+"""
+
+import glob
+
+import numpy as np
+from scipy.stats import pearsonr, spearmanr
+
+KC_NPZ = "/home/mzmyslowski/marcin_multiplex/logs/marker_covariance_ImVs-34_hn.npz"
+RECON_DIR = "/raid_encrypted/immucan/recons/immuvis-gp/immuvis_last_checkpoint-ImVs-34_loo"
+VAR_FLOOR = 1e-6  # skip (image, marker) where the target is essentially flat
+
+
+def main() -> None:
+    kc = np.load(KC_NPZ, allow_pickle=True)
+    names = list(kc["marker_names"])
+    k = kc["k_c"].copy()
+    np.fill_diagonal(k, np.nan)
+
+    files = sorted(glob.glob(RECON_DIR + "/*.npz"))
+    ref = [m for m in names if m in list(np.load(files[0], allow_pickle=True)["marker_names"])]
+
+    nmse: dict[str, list[float]] = {m: [] for m in ref}
+    pear: dict[str, list[float]] = {m: [] for m in ref}
+    for f in files:
+        d = np.load(f, allow_pickle=True)
+        fn = list(d["marker_names"])
+        recon = d["recon"].astype(np.float64)
+        target = d["target"].astype(np.float64)
+        for m in ref:
+            c = fn.index(m)
+            t = target[c].ravel()
+            r = recon[c].ravel()
+            v = t.var()
+            if v < VAR_FLOOR:
+                continue
+            nmse[m].append(((r - t) ** 2).mean() / v)
+            if r.std() > 1e-8:
+                pear[m].append(np.corrcoef(r, t)[0, 1])
+
+    rows = []
+    for i, m in enumerate(names):
+        if m not in ref or not nmse[m]:
+            continue
+        row = k[i]
+        rows.append(
+            {
+                "marker": m,
+                "kc_max": np.nanmax(row),
+                "kc_mean": np.nanmean(row),
+                "kc_nn05": int(np.nansum(row > 0.5)),
+                "nmse": float(np.median(nmse[m])),          # scale-free (1 - R^2)
+                "r2": float(1 - np.median(nmse[m])),
+                "pearson": float(np.median(pear[m])) if pear[m] else np.nan,
+            }
+        )
+
+    markers = [r["marker"] for r in rows]
+    arr = {kk: np.array([r[kk] for r in rows]) for kk in rows[0] if kk != "marker"}
+    print(f"markers: {len(rows)}")
+    print("\nK_C redundancy vs SCALE-INDEPENDENT quality:")
+    print(f"  {'score':<9} {'vs':<9} {'Spearman':>9} {'Pearson':>9}   {'wanted':>7}")
+    for score in ["kc_max", "kc_mean", "kc_nn05"]:
+        for tgt, want in [("nmse", "-"), ("r2", "+"), ("pearson", "+")]:
+            rho = spearmanr(arr[score], arr[tgt]).correlation
+            rp = pearsonr(arr[score], arr[tgt])[0]
+            print(f"  {score:<9} {tgt:<9} {rho:>+9.3f} {rp:>+9.3f}   {want:>7}")
+
+    order = np.argsort(-arr["kc_max"])
+    print("\nBy kc_max (best twin) — high twin should mean low NMSE / high R^2 if effect is real:")
+    print(f"  {'marker':<16} {'kc_max':>7} {'nmse':>7} {'r2':>7} {'pearson':>8}")
+    for j in list(order[:8]) + list(order[-8:]):
+        print(f"  {markers[j]:<16} {arr['kc_max'][j]:>7.3f} {arr['nmse'][j]:>7.3f} {arr['r2'][j]:>7.3f} {arr['pearson'][j]:>8.3f}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_kronecker_marker.py b/tests/test_kronecker_marker.py
new file mode 100644
index 0000000..81bc259
--- /dev/null
+++ b/tests/test_kronecker_marker.py
@@ -0,0 +1,321 @@
+"""Tests for KroneckerMarkerCovariance numerical correctness."""
+
+import math
+import torch
+import pytest
+
+
+def _build_module(grid_size=4, marker_embed_dim=3, hyperkernel_model_dim=8, device="cpu"):
+    """Helper to build a KroneckerMarkerCovariance with small dims for testing."""
+    from multiplex_model.modules.gp_covariance import KroneckerMarkerCovariance
+
+    return KroneckerMarkerCovariance(
+        grid_size=grid_size,
+        marker_embed_dim=marker_embed_dim,
+        hyperkernel_model_dim=hyperkernel_model_dim,
+        kernel_jitter=1e-2,
+        marker_jitter=1e-2,
+        spatial_matern_kernel_nu=1.5,
+        spatial_matern_kernel_length_scale=5.0,
+        device=device,
+    )
+
+
+def test_A_solve_triple_recovers_identity():
+    """A^{-1} A v == v  for random v, using dense materialization as ground truth."""
+    torch.manual_seed(42)
+    n = 4
+    C = 3
+    N = n * n
+    NC = N * C
+
+    mod = _build_module(grid_size=n, marker_embed_dim=3, hyperkernel_model_dim=8)
+
+    # Build K_C from random marker embeddings
+    marker_emb = torch.randn(C, 8)
+    E = mod.embedding_projection(marker_emb)  # [C, 3]
+    K_C = E @ E.T + mod.marker_jitter * torch.eye(C)
+    lam_C, V_C = torch.linalg.eigh(K_C)
+
+    # Triple eigenvalues
+    triple_eigs = (
+        mod.kron_eigs.unsqueeze(-1) * lam_C.unsqueeze(0).unsqueeze(0)
+        + mod.kernel_jitter
+    )
+
+    # Build dense A for ground truth
+    # A = (K_x kron K_y) kron K_C + jitter * I
+    V = mod.V
+    lam = mod.lam
+    K1d = V @ torch.diag(lam) @ V.T
+    K_spatial = torch.kron(K1d, K1d)  # [N, N]
+    A_dense = torch.kron(K_spatial, K_C) + mod.kernel_jitter * torch.eye(NC)
+
+    # Random vector
+    v = torch.randn(NC)
+    Av = A_dense @ v
+
+    # Solve A^{-1} (A v) should recover v
+    recovered = mod._A_solve_triple(Av, V_C, triple_eigs)
+
+    torch.testing.assert_close(recovered, v, atol=1e-3, rtol=1e-3)
+
+
+def test_A_solve_triple_batched():
+    """_A_solve_triple with multiple right-hand sides [NC, m]."""
+    torch.manual_seed(42)
+    n = 4
+    C = 3
+    N = n * n
+    NC = N * C
+    m = 5
+
+    mod = _build_module(grid_size=n, marker_embed_dim=3, hyperkernel_model_dim=8)
+
+    marker_emb = torch.randn(C, 8)
+    E = mod.embedding_projection(marker_emb)
+    K_C = E @ E.T + mod.marker_jitter * torch.eye(C)
+    lam_C, V_C = torch.linalg.eigh(K_C)
+    triple_eigs = (
+        mod.kron_eigs.unsqueeze(-1) * lam_C.unsqueeze(0).unsqueeze(0)
+        + mod.kernel_jitter
+    )
+
+    V = mod.V
+    lam = mod.lam
+    K1d = V @ torch.diag(lam) @ V.T
+    K_spatial = torch.kron(K1d, K1d)
+    A_dense = torch.kron(K_spatial, K_C) + mod.kernel_jitter * torch.eye(NC)
+
+    v = torch.randn(NC, m)
+    Av = A_dense @ v
+    recovered = mod._A_solve_triple(Av, V_C, triple_eigs)
+
+    torch.testing.assert_close(recovered, v, atol=1e-3, rtol=1e-3)
+
+
+def test_log_prob_joint_matches_dense():
+    """log_prob_joint should match direct dense multivariate normal log-prob."""
+    torch.manual_seed(42)
+    n = 4
+    C = 3
+    N = n * n
+    NC = N * C
+
+    mod = _build_module(grid_size=n, marker_embed_dim=3, hyperkernel_model_dim=8)
+
+    marker_emb = torch.randn(C, 8)
+
+    mu_all = torch.randn(N, C)
+    U_all = torch.abs(torch.randn(N, C)) * 0.1 + 0.01  # positive sigma
+    targets = torch.randn(N, C)
+
+    # Our method
+    log_prob = mod.log_prob_joint(mu_all, U_all, targets, marker_emb)
+
+    # Dense ground truth — mirror the module's row-normalization
+    E = torch.nn.functional.normalize(mod.embedding_projection(marker_emb), p=2, dim=1)
+    K_C = E @ E.T + mod.marker_jitter * torch.eye(C)
+
+    V = mod.V
+    lam = mod.lam
+    K1d = V @ torch.diag(lam) @ V.T
+    K_spatial = torch.kron(K1d, K1d)
+    A_dense = torch.kron(K_spatial, K_C) + mod.kernel_jitter * torch.eye(NC)
+
+    # Build U_block [NC, C] in spatial-major order: row (i*C + c) = pixel i, marker c
+    U_block = torch.diag_embed(U_all).reshape(NC, C)  # [N,C] -> [N,C,C] -> [NC,C]
+
+    K_dense = A_dense + U_block @ U_block.T
+
+    # Dense log prob: -0.5 * (e^T K^{-1} e + log|K| + NC*log(2pi))
+    e = (targets - mu_all).reshape(-1)  # [NC] spatial-major: [pix0_ch0, pix0_ch1, ..., pixN_chC]
+    K_inv_e = torch.linalg.solve(K_dense, e)
+    mahal = e @ K_inv_e
+    log_det = torch.linalg.slogdet(K_dense)[1]
+    expected = -0.5 * (mahal + log_det + NC * math.log(2 * math.pi))
+
+    torch.testing.assert_close(log_prob, expected, atol=1e-3, rtol=1e-3)
+
+
+def test_compute_marker_correlation_shape_and_diagonal():
+    """compute_marker_correlation returns CxC with ones on diagonal."""
+    mod = _build_module(grid_size=4, marker_embed_dim=3, hyperkernel_model_dim=8)
+    marker_emb = torch.randn(5, 8)
+
+    corr = mod.compute_marker_correlation(marker_emb)
+    assert corr.shape == (5, 5)
+    torch.testing.assert_close(torch.diag(corr), torch.ones(5), atol=1e-5, rtol=1e-5)
+
+
+def test_hybrid_marker_loss_forward_shape_and_components():
+    """HybridKroneckerMarkerGPNLLLoss returns scalar loss and dict with expected keys."""
+    from multiplex_model.losses import HybridKroneckerMarkerGPNLLLoss
+
+    torch.manual_seed(42)
+    n = 4
+    B, C = 2, 3
+    H = W = n
+
+    mod = _build_module(grid_size=n, marker_embed_dim=3, hyperkernel_model_dim=8)
+    loss_fn = HybridKroneckerMarkerGPNLLLoss(
+        covariance_module=mod,
+        lambda_gp=0.1,
+        downscale_factor=1,
+        device="cpu",
+    )
+
+    target = torch.rand(B, C, H, W)
+    mu = torch.rand(B, C, H, W)
+    logvar = torch.randn(B, C, H, W) * 0.1
+    marker_embeddings = torch.randn(B, C, 8)  # [B, C, model_dim]
+
+    total_loss, loss_dict = loss_fn(target, mu, logvar, marker_embeddings)
+
+    assert total_loss.dim() == 0, "Loss should be scalar"
+    assert total_loss.requires_grad, "Loss must be differentiable"
+    assert "standard_nll" in loss_dict
+    assert "gp_nll" in loss_dict
+    assert "total_loss" in loss_dict
+
+    # Verify gradient flows through marker_embeddings
+    marker_embeddings_grad = torch.randn(B, C, 8, requires_grad=True)
+    total_loss2, _ = loss_fn(target, mu, logvar, marker_embeddings_grad)
+    total_loss2.backward()
+    assert marker_embeddings_grad.grad is not None, "Gradients must flow to marker embeddings"
+
+
+def test_hybrid_marker_loss_lambda_zero_equals_standard():
+    """With lambda_gp=0, HybridKroneckerMarkerGPNLLLoss should equal standard NLL."""
+    from multiplex_model.losses import HybridKroneckerMarkerGPNLLLoss
+
+    torch.manual_seed(42)
+    n = 4
+    B, C = 1, 3
+    H = W = n
+
+    mod = _build_module(grid_size=n, marker_embed_dim=3, hyperkernel_model_dim=8)
+    loss_fn = HybridKroneckerMarkerGPNLLLoss(
+        covariance_module=mod,
+        lambda_gp=0.0,
+        downscale_factor=1,
+        device="cpu",
+    )
+
+    target = torch.rand(B, C, H, W)
+    mu = torch.rand(B, C, H, W)
+    logvar = torch.randn(B, C, H, W) * 0.1
+    marker_embeddings = torch.randn(B, C, 8)
+
+    total_loss, loss_dict = loss_fn(target, mu, logvar, marker_embeddings)
+
+    # Standard NLL computed directly
+    var = torch.exp(logvar)
+    expected_nll = torch.mean((target - mu) ** 2 / (var + 1e-8) + logvar)
+
+    torch.testing.assert_close(total_loss, expected_nll, atol=1e-5, rtol=1e-5)
+
+
+def test_end_to_end_training_step():
+    """Simulate one training step: model forward -> extract embeddings -> loss -> backward."""
+    torch.manual_seed(42)
+    B, C_total, H, W = 2, 5, 16, 16
+    C_active = 4
+
+    from multiplex_model.modules import MultiplexAutoencoder
+    from multiplex_model.losses import HybridKroneckerMarkerGPNLLLoss
+    from multiplex_model.modules.gp_covariance import KroneckerMarkerCovariance
+
+    model = MultiplexAutoencoder(
+        num_channels=C_total,
+        encoder_config={
+            "ma_layers_blocks": [1],
+            "ma_embedding_dims": [8],
+            "pm_layers_blocks": [1],
+            "pm_embedding_dims": [16],
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+        decoder_config={
+            "decoded_embed_dim": 16,
+            "num_blocks": 1,
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+    )
+
+    # hyperkernel_model_dim = pm_embedding_dims[0] * kernel_size^2 * ma_embedding_dims[-1]
+    # = 16 * 1 * 8 = 128
+    hyperkernel_model_dim = 16 * 1 * 8
+
+    gp_module = KroneckerMarkerCovariance(
+        grid_size=H,
+        marker_embed_dim=8,
+        hyperkernel_model_dim=hyperkernel_model_dim,
+        kernel_jitter=1e-2,
+        marker_jitter=1e-2,
+        device="cpu",
+    )
+
+    loss_fn = HybridKroneckerMarkerGPNLLLoss(
+        covariance_module=gp_module,
+        lambda_gp=0.1,
+        downscale_factor=1,
+        device="cpu",
+    )
+
+    optimizer = torch.optim.AdamW(
+        list(model.parameters()) + list(gp_module.parameters()),
+        lr=1e-3,
+    )
+
+    # Simulate forward pass
+    img = torch.rand(B, C_active, H, W)
+    channel_ids = torch.arange(C_active).unsqueeze(0).expand(B, -1)
+    active_ids = channel_ids.clone()
+
+    output = model(img, active_ids, channel_ids)["output"]
+    mi, logvar = output.unbind(dim=-1)
+    mi = torch.sigmoid(mi)
+    logvar = torch.clamp(logvar, -15.0, 15.0)
+
+    # Extract marker embeddings
+    marker_emb = model.encoder.hyperkernel.hyperkernel_weights(channel_ids)
+
+    # Compute loss
+    loss, loss_dict = loss_fn(img, mi, logvar, marker_emb)
+
+    # Backward
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    # Verify gradients exist
+    assert model.encoder.hyperkernel.hyperkernel_weights.weight.grad is not None
+    assert gp_module.embedding_projection.weight.grad is not None
+    assert loss.isfinite(), f"Loss is not finite: {loss.item()}"
+
+    print(f"End-to-end smoke test passed. Loss: {loss.item():.4f}")
+
+
+def test_log_prob_joint_C_greater_than_embed_dim():
+    """C > marker_embed_dim: K_C has repeated eigenvalues — float64 eigh path must not diverge."""
+    from multiplex_model.modules.gp_covariance import KroneckerMarkerCovariance
+
+    C, marker_embed_dim, grid_size = 10, 4, 8
+    N = grid_size * grid_size
+    mod = KroneckerMarkerCovariance(
+        grid_size=grid_size,
+        marker_embed_dim=marker_embed_dim,
+        hyperkernel_model_dim=16,
+        kernel_jitter=1e-2,
+        marker_jitter=1e-2,
+        device="cpu",
+    )
+    torch.manual_seed(0)
+    targets = torch.randn(N, C)
+    mu = torch.randn(N, C)
+    sigma = torch.ones(N, C) * 0.5
+    marker_emb = torch.randn(C, 16)
+
+    lp = mod.log_prob_joint(mu, sigma, targets, marker_emb)
+
+    assert lp.isfinite(), f"log_prob_joint non-finite with C={C} > marker_embed_dim={marker_embed_dim}: {lp.item()}"
diff --git a/tests/test_training_integration.py b/tests/test_training_integration.py
new file mode 100644
index 0000000..5f5c6f4
--- /dev/null
+++ b/tests/test_training_integration.py
@@ -0,0 +1,604 @@
+"""Integration tests for training loop, validation loop, and logging infrastructure.
+
+These tests cover the scaffolding layer (masking, script functions, logging) that
+the numerical unit tests in test_kronecker_marker.py do not exercise.
+"""
+
+import math
+import sys
+import os
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import torch
+import pytest
+
+# Make the training script importable as a module (functions only, __main__ is guarded)
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+# ---------------------------------------------------------------------------
+# Shared helpers
+# ---------------------------------------------------------------------------
+
+def _build_tiny_setup(grid_size=8, C_total=4):
+    """Build a tiny model + GP module + loss function for integration tests."""
+    from multiplex_model.modules import MultiplexAutoencoder
+    from multiplex_model.modules.gp_covariance import KroneckerMarkerCovariance
+    from multiplex_model.losses import HybridKroneckerMarkerGPNLLLoss
+
+    model = MultiplexAutoencoder(
+        num_channels=C_total,
+        encoder_config={
+            "ma_layers_blocks": [1],
+            "ma_embedding_dims": [8],
+            "pm_layers_blocks": [1],
+            "pm_embedding_dims": [16],
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+        decoder_config={
+            "decoded_embed_dim": 16,
+            "num_blocks": 1,
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+    )
+
+    # hyperkernel_model_dim = pm_embedding_dims[0] * kernel_size^2 * ma_embedding_dims[-1]
+    hyperkernel_model_dim = 16 * 1 * 8
+
+    gp_module = KroneckerMarkerCovariance(
+        grid_size=grid_size,
+        marker_embed_dim=8,
+        hyperkernel_model_dim=hyperkernel_model_dim,
+        kernel_jitter=1e-2,
+        marker_jitter=1e-2,
+        device="cpu",
+    )
+
+    loss_fn = HybridKroneckerMarkerGPNLLLoss(
+        covariance_module=gp_module,
+        lambda_gp=0.1,
+        downscale_factor=1,
+        device="cpu",
+    )
+
+    return model, gp_module, loss_fn
+
+
+def _make_fake_dataloader(B=2, C=4, H=8, W=8, num_batches=3):
+    """DataLoader yielding (img, channel_ids, panel_idx, img_path) like the real one."""
+
+    class FakeDataset(torch.utils.data.Dataset):
+        def __init__(self, n):
+            self.n = n
+
+        def __len__(self):
+            return self.n
+
+        def __getitem__(self, idx):
+            img = torch.rand(C, H, W)
+            channel_ids = torch.arange(C)
+            panel_idx = torch.tensor(0)
+            img_path = f"fake/path/{idx}.tiff"
+            return img, channel_ids, panel_idx, img_path
+
+    return torch.utils.data.DataLoader(FakeDataset(B * num_batches), batch_size=B)
+
+
+# ---------------------------------------------------------------------------
+# Test 1: Validation loop smoke test
+# ---------------------------------------------------------------------------
+
+def test_validation_loop_runs():
+    """test_masked_gp runs end-to-end without error and returns finite metrics."""
+    from train_masked_model_gp import test_masked_gp
+
+    torch.manual_seed(42)
+    H = W = 8
+    C = 4
+    B = 2
+
+    model, gp_module, loss_fn = _build_tiny_setup(grid_size=H, C_total=C)
+    dataloader = _make_fake_dataloader(B=B, C=C, H=H, W=W, num_batches=3)
+    marker_names_map = {i: f"marker_{i}" for i in range(C)}
+
+    val_metrics = test_masked_gp(
+        model=model,
+        test_dataloader=dataloader,
+        device="cpu",
+        epoch=0,
+        gp_covariance_module=gp_module,
+        gp_loss_fn=loss_fn,
+        marker_names_map=marker_names_map,
+        num_plots=1,
+        spatial_masking_ratio=0.5,
+        fully_masked_channels_max_frac=0.25,
+        mask_patch_size=2,
+        use_gp_loss=True,
+        use_marker_covariance=True,
+    )
+
+    for key in ("val_loss", "val_mae", "val_mse", "val_standard_nll", "val_gp_nll"):
+        assert key in val_metrics, f"Missing key: {key}"
+        assert math.isfinite(val_metrics[key]), f"{key} is not finite: {val_metrics[key]}"
+
+
+# ---------------------------------------------------------------------------
+# Test 2: Config parsing and module instantiation
+# ---------------------------------------------------------------------------
+
+def test_config_fields_and_module_instantiation():
+    """TrainingConfig accepts new marker covariance fields; KroneckerMarkerCovariance
+    instantiates correctly and its parameters are on the right device."""
+    from multiplex_model.utils.configuration import TrainingConfig
+    from multiplex_model.modules.gp_covariance import KroneckerMarkerCovariance
+
+    # Verify new fields exist with correct defaults
+    defaults = TrainingConfig.model_fields
+    assert "use_marker_covariance" in defaults
+    assert "marker_embed_dim" in defaults
+    assert "marker_jitter" in defaults
+
+    assert defaults["use_marker_covariance"].default is False
+    assert defaults["marker_embed_dim"].default == 32
+    assert defaults["marker_jitter"].default == pytest.approx(1e-2)
+
+    # Instantiate the module and verify parameters are on device
+    gp_module = KroneckerMarkerCovariance(
+        grid_size=16,
+        marker_embed_dim=32,
+        hyperkernel_model_dim=128,
+        kernel_jitter=1e-2,
+        marker_jitter=1e-2,
+        device="cpu",
+    )
+    gp_module = gp_module.to("cpu")
+
+    for name, param in gp_module.named_parameters():
+        assert param.device.type == "cpu", f"Parameter {name} is on {param.device}, expected cpu"
+
+    # Verify the projection layer has the right shape
+    assert gp_module.embedding_projection.in_features == 128
+    assert gp_module.embedding_projection.out_features == 32
+
+
+# ---------------------------------------------------------------------------
+# Test 3: Logging function signature
+# ---------------------------------------------------------------------------
+
+def test_log_validation_images_accepts_name_suffix():
+    """log_validation_images accepts name_suffix kwarg without TypeError."""
+    from multiplex_model.utils.train_logging import log_validation_images
+
+    fig, _ = plt.subplots(1, 1, figsize=(2, 2))
+    # _experiment is None in tests so nothing is actually logged — just check no TypeError
+    log_validation_images(
+        fig=fig,
+        panel_idx=0,
+        img_path="fake/path.tiff",
+        epoch=0,
+        masked_channels_names="marker_0",
+        img_idx=0,
+        name_suffix="_sigma",
+    )
+    plt.close(fig)
+
+
+def test_log_validation_images_default_no_suffix():
+    """log_validation_images still works without name_suffix (backward compat)."""
+    from multiplex_model.utils.train_logging import log_validation_images
+
+    fig, _ = plt.subplots(1, 1, figsize=(2, 2))
+    log_validation_images(
+        fig=fig,
+        panel_idx=0,
+        img_path="fake/path.tiff",
+        epoch=0,
+        masked_channels_names="marker_0",
+        img_idx=0,
+    )
+    plt.close(fig)
+
+
+# ---------------------------------------------------------------------------
+# Test 4: Training step with masking via script functions
+# ---------------------------------------------------------------------------
+
+def test_training_step_with_channel_and_spatial_masking():
+    """Full training step: channel mask → spatial mask → forward → embed extract → loss → backward.
+
+    Verifies that img, channel_ids, and marker_emb shapes are all consistent
+    after apply_channel_masking reduces the channel set.
+    """
+    from multiplex_model.utils.masking import apply_channel_masking, apply_spatial_masking
+
+    torch.manual_seed(42)
+    H = W = 8
+    C_total = 6
+    B = 2
+
+    model, gp_module, loss_fn = _build_tiny_setup(grid_size=H, C_total=C_total)
+    optimizer = torch.optim.AdamW(
+        list(model.parameters()) + list(gp_module.parameters()), lr=1e-3
+    )
+
+    img = torch.rand(B, C_total, H, W)
+    channel_ids = torch.arange(C_total).unsqueeze(0).expand(B, -1).contiguous()
+
+    # Channel masking — both img and channel_ids are reduced to the active subset
+    img, channel_ids, masked_img, active_channel_ids = apply_channel_masking(
+        img,
+        channel_ids,
+        min_channels_frac=0.5,
+        fully_masked_channels_max_frac=0.25,
+        apply_channel_subset_sampling=True,
+    )
+
+    # Spatial masking
+    masked_img, _ = apply_spatial_masking(masked_img, spatial_masking_ratio=0.5, mask_patch_size=2)
+
+    # Forward pass
+    output = model(masked_img, active_channel_ids, channel_ids)["output"]
+    mi, logvar = output.unbind(dim=-1)
+    mi = torch.sigmoid(mi)
+    logvar = torch.clamp(logvar, -15.0, 15.0)
+
+    # Embedding extraction — channel_ids now matches the reduced img
+    marker_emb = model.encoder.hyperkernel.hyperkernel_weights(channel_ids)
+
+    C_active = img.shape[1]
+    assert mi.shape == img.shape, f"mi {mi.shape} != img {img.shape}"
+    assert marker_emb.shape[:2] == (B, C_active), (
+        f"marker_emb {marker_emb.shape} inconsistent with img channel count {C_active}"
+    )
+
+    # Loss + backward
+    loss, loss_dict = loss_fn(img, mi, logvar, marker_emb)
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    assert loss.isfinite(), f"Loss is not finite: {loss.item()}"
+    assert set(loss_dict.keys()) == {"standard_nll", "gp_nll", "total_loss"}
+
+
+# ---------------------------------------------------------------------------
+# Test 5: EncoderConfig mask_token fields
+# ---------------------------------------------------------------------------
+
+
+def test_encoder_config_accepts_mask_token_fields():
+    from multiplex_model.utils.configuration import EncoderConfig
+
+    cfg = EncoderConfig(
+        ma_layers_blocks=[1],
+        ma_embedding_dims=[8],
+        pm_layers_blocks=[1],
+        pm_embedding_dims=[16],
+        hyperkernel={"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        use_mask_token=True,
+        mask_token_init=0.5,
+    )
+    assert cfg.use_mask_token is True
+    assert cfg.mask_token_init == 0.5
+
+
+def test_encoder_config_mask_token_defaults():
+    from multiplex_model.utils.configuration import EncoderConfig
+
+    cfg = EncoderConfig(
+        ma_layers_blocks=[1],
+        ma_embedding_dims=[8],
+        pm_layers_blocks=[1],
+        pm_embedding_dims=[16],
+        hyperkernel={"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+    )
+    assert cfg.use_mask_token is False
+    assert cfg.mask_token_init == 0.0
+
+
+# ---------------------------------------------------------------------------
+# Test 6: Learnable mask token in MultiplexImageEncoder
+# ---------------------------------------------------------------------------
+
+
+def test_encoder_mask_token_is_none_when_disabled():
+    from multiplex_model.modules.immuvis import MultiplexImageEncoder
+
+    enc = MultiplexImageEncoder(
+        num_channels=4,
+        ma_layers_blocks=[1],
+        ma_embedding_dims=[8],
+        pm_layers_blocks=[1],
+        pm_embedding_dims=[16],
+        hyperkernel_config={"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+    )
+    assert enc.mask_token is None
+
+
+def test_encoder_mask_token_is_parameter_when_enabled():
+    import torch.nn as nn
+
+    from multiplex_model.modules.immuvis import MultiplexImageEncoder
+
+    enc = MultiplexImageEncoder(
+        num_channels=4,
+        ma_layers_blocks=[1],
+        ma_embedding_dims=[8],
+        pm_layers_blocks=[1],
+        pm_embedding_dims=[16],
+        hyperkernel_config={"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        use_mask_token=True,
+        mask_token_init=0.5,
+    )
+    assert isinstance(enc.mask_token, nn.Parameter)
+    assert enc.mask_token.item() == pytest.approx(0.5)
+
+
+def test_encoder_forward_applies_mask_token_to_masked_pixels():
+    from multiplex_model.modules.immuvis import MultiplexImageEncoder
+
+    torch.manual_seed(0)
+    B, C, H, W = 1, 2, 4, 4
+    enc = MultiplexImageEncoder(
+        num_channels=C,
+        ma_layers_blocks=[1],
+        ma_embedding_dims=[8],
+        pm_layers_blocks=[1],
+        pm_embedding_dims=[16],
+        hyperkernel_config={"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        use_mask_token=True,
+        mask_token_init=99.0,
+    )
+    x = torch.zeros(B, C, H, W)
+    spatial_mask = torch.zeros(B, C, H, W, dtype=torch.bool)
+    spatial_mask[:, :, 0, 0] = True
+
+    with torch.no_grad():
+        token_val = enc.mask_token.to(dtype=x.dtype)
+        x_after = torch.where(spatial_mask, token_val, x)
+    assert x_after[:, :, 0, 0].allclose(torch.tensor(99.0))
+    assert x_after[:, :, 1, 1].allclose(torch.tensor(0.0))
+
+    enc_indices = torch.arange(C).unsqueeze(0).expand(B, -1)
+    out = enc(x, enc_indices, spatial_mask=spatial_mask)
+    assert "output" in out
+
+
+# ---------------------------------------------------------------------------
+# Test 7: MultiplexAutoencoder spatial_mask and architecture config
+# ---------------------------------------------------------------------------
+
+
+def test_autoencoder_encode_accepts_spatial_mask():
+    import torch
+    from multiplex_model.modules import MultiplexAutoencoder
+
+    B, C, H, W = 2, 4, 8, 8
+    model = MultiplexAutoencoder(
+        num_channels=C,
+        encoder_config={
+            "ma_layers_blocks": [1],
+            "ma_embedding_dims": [8],
+            "pm_layers_blocks": [1],
+            "pm_embedding_dims": [16],
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+        decoder_config={
+            "decoded_embed_dim": 16,
+            "num_blocks": 1,
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+    )
+    x = torch.rand(B, C, H, W)
+    enc_ids = torch.arange(C).unsqueeze(0).expand(B, -1)
+    spatial_mask = torch.zeros(B, C, H, W, dtype=torch.bool)
+    out = model.encode(x, enc_ids, spatial_mask=spatial_mask)
+    assert "output" in out
+
+
+def test_autoencoder_forward_accepts_spatial_mask():
+    import torch
+    from multiplex_model.modules import MultiplexAutoencoder
+
+    B, C, H, W = 2, 4, 8, 8
+    model = MultiplexAutoencoder(
+        num_channels=C,
+        encoder_config={
+            "ma_layers_blocks": [1],
+            "ma_embedding_dims": [8],
+            "pm_layers_blocks": [1],
+            "pm_embedding_dims": [16],
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+        decoder_config={
+            "decoded_embed_dim": 16,
+            "num_blocks": 1,
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+    )
+    x = torch.rand(B, C, H, W)
+    enc_ids = torch.arange(C).unsqueeze(0).expand(B, -1)
+    dec_ids = enc_ids
+    spatial_mask = torch.zeros(B, C, H, W, dtype=torch.bool)
+    out = model(x, enc_ids, dec_ids, spatial_mask=spatial_mask)
+    assert "output" in out
+
+
+def test_autoencoder_get_architecture_config_roundtrip():
+    import torch
+    from multiplex_model.modules import MultiplexAutoencoder
+
+    C = 4
+    model = MultiplexAutoencoder(
+        num_channels=C,
+        encoder_config={
+            "ma_layers_blocks": [1],
+            "ma_embedding_dims": [8],
+            "pm_layers_blocks": [1],
+            "pm_embedding_dims": [16],
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+        decoder_config={
+            "decoded_embed_dim": 16,
+            "num_blocks": 1,
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+    )
+    cfg = model.get_architecture_config()
+    assert cfg["num_channels"] == C
+    assert "encoder_config" in cfg
+    assert "decoder_config" in cfg
+
+    model2 = MultiplexAutoencoder(**cfg)
+    assert model2.num_channels == C
+
+
+def test_autoencoder_load_from_checkpoint_roundtrip():
+    import torch
+    from multiplex_model.modules import MultiplexAutoencoder
+
+    C = 4
+    model = MultiplexAutoencoder(
+        num_channels=C,
+        encoder_config={
+            "ma_layers_blocks": [1],
+            "ma_embedding_dims": [8],
+            "pm_layers_blocks": [1],
+            "pm_embedding_dims": [16],
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+        decoder_config={
+            "decoded_embed_dim": 16,
+            "num_blocks": 1,
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+    )
+    fake_checkpoint = {
+        "model_state_dict": model.state_dict(),
+        "model_config": model.get_architecture_config(),
+    }
+    loaded = MultiplexAutoencoder.load_from_checkpoint(fake_checkpoint)
+    assert loaded.num_channels == C
+    for (k1, v1), (k2, v2) in zip(model.state_dict().items(), loaded.state_dict().items()):
+        assert k1 == k2
+        assert v1.allclose(v2)
+
+
+# ---------------------------------------------------------------------------
+# Test 8: log_training_metrics mask_token parameter
+# ---------------------------------------------------------------------------
+
+
+def test_log_training_metrics_accepts_mask_token():
+    import inspect
+    from multiplex_model.utils.train_logging import log_training_metrics
+
+    sig = inspect.signature(log_training_metrics)
+    assert "mask_token" in sig.parameters, "log_training_metrics must accept mask_token kwarg"
+    param = sig.parameters["mask_token"]
+    assert param.default is None, "mask_token should default to None"
+
+    # Calling with mask_token must not raise TypeError
+    log_training_metrics(
+        loss=0.5,
+        lr=1e-3,
+        mu=0.5,
+        logvar=-1.0,
+        mae=0.1,
+        mse=0.01,
+        step=0,
+        mask_token=0.123,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Test 9: learnmask+GP validation loop smoke test
+# ---------------------------------------------------------------------------
+
+def test_learnmask_gp_validation_loop_runs():
+    """test_masked_learnmask_gp runs end-to-end without error and returns finite metrics."""
+    from train_masked_model_learnmask_gp import test_masked_learnmask_gp
+    from multiplex_model.modules import MultiplexAutoencoder
+
+    torch.manual_seed(42)
+    H = W = 8
+    C = 4
+    B = 2
+
+    # Build model with use_mask_token=True
+    hyperkernel_model_dim = 16 * 1 * 8
+    model = MultiplexAutoencoder(
+        num_channels=C,
+        encoder_config={
+            "ma_layers_blocks": [1],
+            "ma_embedding_dims": [8],
+            "pm_layers_blocks": [1],
+            "pm_embedding_dims": [16],
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+            "use_mask_token": True,
+            "mask_token_init": 0.0,
+        },
+        decoder_config={
+            "decoded_embed_dim": 16,
+            "num_blocks": 1,
+            "hyperkernel_config": {"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        },
+    )
+    _, gp_module, loss_fn = _build_tiny_setup(grid_size=H, C_total=C)
+    dataloader = _make_fake_dataloader(B=B, C=C, H=H, W=W, num_batches=3)
+    marker_names_map = {i: f"marker_{i}" for i in range(C)}
+
+    val_metrics = test_masked_learnmask_gp(
+        model=model,
+        test_dataloader=dataloader,
+        device="cpu",
+        epoch=0,
+        gp_covariance_module=gp_module,
+        gp_loss_fn=loss_fn,
+        marker_names_map=marker_names_map,
+        num_plots=1,
+        spatial_masking_ratio=0.5,
+        fully_masked_channels_max_frac=0.25,
+        mask_patch_size=2,
+        use_gp_loss=True,
+    )
+
+    for key in ("val_loss", "val_mae", "val_mse", "val_standard_nll", "val_gp_nll"):
+        assert key in val_metrics, f"Missing key: {key}"
+        assert math.isfinite(val_metrics[key]), f"{key} is not finite: {val_metrics[key]}"
+
+
+# ---------------------------------------------------------------------------
+# Test 10: mask token gradient flow
+# ---------------------------------------------------------------------------
+
+def test_mask_token_gradient_flows():
+    """Backward pass propagates gradient to mask_token parameter."""
+    from multiplex_model.modules.immuvis import MultiplexImageEncoder
+
+    torch.manual_seed(0)
+    B, C, H, W = 1, 2, 4, 4
+    enc = MultiplexImageEncoder(
+        num_channels=C,
+        ma_layers_blocks=[1],
+        ma_embedding_dims=[8],
+        pm_layers_blocks=[1],
+        pm_embedding_dims=[16],
+        hyperkernel_config={"kernel_size": 1, "padding": 0, "stride": 1, "use_bias": True},
+        use_mask_token=True,
+        mask_token_init=0.0,
+    )
+
+    x = torch.rand(B, C, H, W)
+    spatial_mask = torch.zeros(B, C, H, W, dtype=torch.bool)
+    spatial_mask[:, :, :2, :2] = True
+    enc_indices = torch.arange(C).unsqueeze(0).expand(B, -1)
+
+    out = enc(x, enc_indices, spatial_mask=spatial_mask)
+    loss = out["output"].sum()
+    loss.backward()
+
+    assert enc.mask_token is not None
+    assert enc.mask_token.grad is not None, "mask_token has no gradient — not in computation graph"
diff --git a/train.sh b/train.sh
new file mode 100755
index 0000000..bee8160
--- /dev/null
+++ b/train.sh
@@ -0,0 +1,43 @@
+#!/bin/bash
+#SBATCH --partition=common
+#SBATCH --qos=mzmyslowski
+#SBATCH --nodelist=szary
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=50G
+#SBATCH --gres=gpu:1
+#SBATCH --time=7-00:00:00
+#SBATCH --job-name=train
+#SBATCH --output=logs/train_%j.out
+#SBATCH --error=logs/train_%j.err
+
+set -e
+
+if [ -z "$1" ]; then
+    echo "Usage: sbatch train.sh <config_file> [gp]"
+    echo "  config_file: path to YAML config"
+    echo "  gp: pass 'gp' as second arg to use GP training script"
+    exit 1
+fi
+
+config_file=$1
+use_gp=${2:-""}
+
+mkdir -p logs
+
+export COMET_API_KEY=$(grep COMET_API_KEY ~/.bashrc | cut -d= -f2)
+
+. ~/venv/bin/activate
+
+if [ "$use_gp" = "gp" ]; then
+    echo "Starting GP training with config: $config_file"
+    python3 train_masked_model_gp.py "$config_file"
+elif [ "$use_gp" = "learnmask" ]; then
+    echo "Starting learnmask training with config: $config_file"
+    python3 train_masked_model_learnmask.py "$config_file"
+elif [ "$use_gp" = "learnmask_gp" ]; then
+    echo "Starting learnmask+GP training with config: $config_file"
+    python3 train_masked_model_learnmask_gp.py "$config_file"
+else
+    echo "Starting standard training with config: $config_file"
+    python3 train_masked_model.py "$config_file"
+fi
diff --git a/train_masked_config.yaml b/train_masked_config.yaml
index ddb2db7..fa275ca 100644
--- a/train_masked_config.yaml
+++ b/train_masked_config.yaml
@@ -47,7 +47,7 @@ save_checkpoint_freq: 5
 beta: 1.0
 
 # Comet.ml logging configuration
-tags: [...]
-comet_project: ...
-comet_workspace: null  # optional, can also be set via COMET_WORKSPACE env var
+tags: []
+comet_project: multiplex-image-model
+comet_workspace: micha-zmys-owski  # optional, can also be set via COMET_WORKSPACE env var
 comet_api_key: null  # optional, can also be set via COMET_API_KEY env var
diff --git a/train_masked_gp_config.yaml b/train_masked_gp_config.yaml
index 206f063..8d7d85f 100644
--- a/train_masked_gp_config.yaml
+++ b/train_masked_gp_config.yaml
@@ -5,9 +5,10 @@
 # GP LOSS CONFIGURATION
 # ============================================================================
 use_gp_loss: true                    # Enable/disable GP loss
-lambda_gp: 0.0                      # Weight for GP loss (0.0 = only standard, 1.0 = only GP)
+use_kronecker_gp: true               # Use Kronecker (~40x faster than CG)
+lambda_gp: 0.1                      # Weight for GP loss (0.0 = only standard, 1.0 = only GP)
 gp_kernel_jitter: 1e-2               # Diagonal noise for numerical stability
-gp_lengthscale: 0.1                  # Spatial correlation length scale
+gp_lengthscale: 5.0                  # Spatial correlation length scale
 gp_max_cg_iterations: 50             # Max conjugate gradient iterations
 gp_downscale_factor: 1               # Spatial downsampling (1=none, 2=half, 4=quarter)
 gp_learn_lengthscale: true          # Whether to learn kernel lengthscale
@@ -45,7 +46,7 @@ panel_config: configs/all_panels_config.yaml
 tokenizer_config: configs/all_markers_tokenizer.yaml
 input_image_size: [112, 112]
 num_workers: 8
-batch_size: 1
+batch_size: 8
 
 # Training configuration
 device: cuda
@@ -53,19 +54,20 @@ lr: 5e-4
 final_lr: 1e-5
 weight_decay: 0.0001
 gradient_accumulation_steps: 1
-epochs: 10
-frac_warmup_steps: 0.1
+epochs: 200
+frac_warmup_steps: 0.01
 min_channels_frac: 0.75
 spatial_masking_ratio: 0.6
 fully_masked_channels_max_frac: 0.5
 mask_patch_size: 8
-from_checkpoint: null
+from_checkpoint: checkpoints/last_checkpoint-ImVs-12.pth
+reset_lr_schedule: true              # fresh cosine cycle from trained weights
 checkpoints_dir: checkpoints
 save_checkpoint_freq: 5
 beta: 0.5
 
 # Comet.ml logging configuration
-tags: ['SZARY', 'GP', 'times', 'lambda 0.1']
-comet_project: ...
-comet_workspace: ...  # optional, can also be set via COMET_WORKSPACE env var
-comet_api_key: ...  # optional, can also be set via COMET_API_KEY env var
+tags: ['SZARY', 'GP', 'kronecker', 'lambda 0.1', 'run2']
+comet_project: multiplex-image-model
+comet_workspace: micha-zmys-owski
+comet_api_key: null  # set via COMET_API_KEY env var
diff --git a/train_masked_gp_marker_config.yaml b/train_masked_gp_marker_config.yaml
new file mode 100644
index 0000000..0259789
--- /dev/null
+++ b/train_masked_gp_marker_config.yaml
@@ -0,0 +1,74 @@
+# Configuration for training with Kronecker Marker GP loss
+# Extends standard GP config with marker covariance from Hyperkernel embeddings
+
+# ============================================================================
+# GP LOSS CONFIGURATION
+# ============================================================================
+use_gp_loss: true
+use_kronecker_gp: true
+use_marker_covariance: true          # Enable marker covariance (K_C from embeddings)
+marker_embed_dim: 32                  # Projection dim for embedding -> K_C
+marker_jitter: 1e-2                   # Jitter for K_C numerical stability
+lambda_gp: 0.1
+gp_kernel_jitter: 1e-2
+gp_lengthscale: 5.0
+gp_max_cg_iterations: 50
+gp_downscale_factor: 1
+gp_learn_lengthscale: false           # Not applicable for Kronecker
+
+# ============================================================================
+# STANDARD TRAINING CONFIGURATION
+# ============================================================================
+encoder:
+    ma_layers_blocks: [4,]
+    ma_embedding_dims: [16,]
+    pm_layers_blocks: [4, 4, 4]
+    pm_embedding_dims: [128, 256, 512]
+    use_latent_norm: true
+
+    hyperkernel:
+      kernel_size: 1
+      padding: 0
+      stride: 1
+      use_bias: true
+
+decoder:
+  decoded_embed_dim: 384
+  num_blocks: 1
+
+  hyperkernel:
+      kernel_size: 1
+      padding: 0
+      stride: 1
+      use_bias: true
+
+# Data configuration
+panel_config: configs/all_panels_config.yaml
+tokenizer_config: configs/all_markers_tokenizer.yaml
+input_image_size: [112, 112]
+num_workers: 8
+batch_size: 8
+
+# Training configuration
+device: cuda
+lr: 5e-4
+final_lr: 1e-5
+weight_decay: 0.0001
+gradient_accumulation_steps: 1
+epochs: 200
+frac_warmup_steps: 0.01
+min_channels_frac: 0.75
+spatial_masking_ratio: 0.6
+fully_masked_channels_max_frac: 0.5
+mask_patch_size: 8
+from_checkpoint: null
+reset_lr_schedule: false
+checkpoints_dir: checkpoints
+save_checkpoint_freq: 5
+beta: 0.5
+
+# Comet.ml logging configuration
+tags: ['SZARY', 'GP', 'kronecker', 'marker-covariance']
+comet_project: multiplex-image-model
+comet_workspace: micha-zmys-owski
+comet_api_key: null
diff --git a/train_masked_learnmask_config.yaml b/train_masked_learnmask_config.yaml
new file mode 100644
index 0000000..c9f0b51
--- /dev/null
+++ b/train_masked_learnmask_config.yaml
@@ -0,0 +1,71 @@
+# Configuration for training with learnable mask token (beta-NLL loss, no GP)
+# Spatial masking uses a learnable scalar token instead of zero-fill
+
+# ============================================================================
+# ENCODER / DECODER ARCHITECTURE
+# ============================================================================
+encoder:
+    ma_layers_blocks: [4,]
+    ma_embedding_dims: [16,]
+    pm_layers_blocks: [4, 4, 4]
+    pm_embedding_dims: [128, 256, 512]
+    use_latent_norm: true
+    use_mask_token: true
+    mask_token_init: 0.0
+
+    hyperkernel:
+      kernel_size: 1
+      padding: 0
+      stride: 1
+      use_bias: true
+
+decoder:
+  decoded_embed_dim: 384
+  num_blocks: 1
+
+  hyperkernel:
+      kernel_size: 1
+      padding: 0
+      stride: 1
+      use_bias: true
+
+# ============================================================================
+# DATA CONFIGURATION
+# ============================================================================
+panel_config: configs/all_panels_config.yaml
+tokenizer_config: configs/all_markers_tokenizer.yaml
+input_image_size: [112, 112]
+num_workers: 8
+batch_size: 8
+
+# ============================================================================
+# TRAINING CONFIGURATION
+# ============================================================================
+device: cuda
+lr: 5e-4
+final_lr: 1e-5
+weight_decay: 0.0001
+gradient_accumulation_steps: 1
+epochs: 200
+frac_warmup_steps: 0.01
+min_channels_frac: 0.75
+spatial_masking_ratio: 0.6
+fully_masked_channels_max_frac: 0.5
+mask_patch_size: 8
+beta: 0.5
+
+# ============================================================================
+# CHECKPOINT CONFIGURATION
+# ============================================================================
+from_checkpoint: checkpoints/last_checkpoint-ImVs-30.pth
+reset_lr_schedule: false
+checkpoints_dir: checkpoints
+save_checkpoint_freq: 5
+
+# ============================================================================
+# COMET.ML LOGGING
+# ============================================================================
+tags: ['SZARY', 'learnmask', 'mask-token']
+comet_project: multiplex-image-model
+comet_workspace: micha-zmys-owski
+comet_api_key: null
diff --git a/train_masked_learnmask_gp_config.yaml b/train_masked_learnmask_gp_config.yaml
new file mode 100644
index 0000000..f038afd
--- /dev/null
+++ b/train_masked_learnmask_gp_config.yaml
@@ -0,0 +1,86 @@
+# Configuration: learnable mask token + Kronecker marker covariance GP loss
+# Combines (K_x ⊗ K_y) ⊗ K_C marker covariance with learnable spatial mask token
+
+# ============================================================================
+# GP LOSS CONFIGURATION
+# ============================================================================
+use_gp_loss: true
+use_kronecker_gp: true
+use_marker_covariance: true
+marker_embed_dim: 32
+marker_jitter: 1.0e-2
+lambda_gp: 0.1
+gp_kernel_jitter: 1.0e-2
+gp_lengthscale: 5.0
+gp_downscale_factor: 1
+gp_max_cg_iterations: 50
+gp_learn_lengthscale: false
+
+# ============================================================================
+# ENCODER / DECODER ARCHITECTURE
+# ============================================================================
+encoder:
+    ma_layers_blocks: [4,]
+    ma_embedding_dims: [16,]
+    pm_layers_blocks: [4, 4, 4]
+    pm_embedding_dims: [128, 256, 512]
+    use_latent_norm: true
+    use_mask_token: true
+    mask_token_init: 0.0
+
+    hyperkernel:
+      kernel_size: 1
+      padding: 0
+      stride: 1
+      use_bias: true
+
+decoder:
+  decoded_embed_dim: 384
+  num_blocks: 1
+
+  hyperkernel:
+      kernel_size: 1
+      padding: 0
+      stride: 1
+      use_bias: true
+
+# ============================================================================
+# DATA CONFIGURATION
+# ============================================================================
+panel_config: configs/all_panels_config.yaml
+tokenizer_config: configs/all_markers_tokenizer.yaml
+input_image_size: [112, 112]
+num_workers: 8
+batch_size: 8
+
+# ============================================================================
+# TRAINING CONFIGURATION
+# ============================================================================
+device: cuda
+lr: 5.0e-4
+final_lr: 1.0e-5
+weight_decay: 0.0001
+gradient_accumulation_steps: 1
+epochs: 200
+frac_warmup_steps: 0.01
+min_channels_frac: 0.75
+spatial_masking_ratio: 0.6
+fully_masked_channels_max_frac: 0.5
+mask_patch_size: 8
+beta: 0.5
+
+# ============================================================================
+# CHECKPOINT CONFIGURATION
+# ============================================================================
+from_checkpoint: null
+reset_lr_schedule: false
+checkpoints_dir: checkpoints
+save_checkpoint_freq: 5
+
+# ============================================================================
+# COMET.ML LOGGING
+# ============================================================================
+tags: ['SZARY', 'learnmask', 'mask-token', 'GP', 'kronecker', 'marker-covariance']
+comet_project: multiplex-image-model
+comet_workspace: micha-zmys-owski
+comet_api_key: null
diff --git a/train_masked_model.py b/train_masked_model.py
index 0d230ee..730fd2c 100644
--- a/train_masked_model.py
+++ b/train_masked_model.py
@@ -1,3 +1,4 @@
+import math
 import os
 import sys
 
@@ -32,6 +33,7 @@
     get_scheduler_with_warmup,
     init_experiment,
     log_training_metrics,
+    log_validation_batch_metrics,
     log_validation_images,
     log_validation_metrics,
     plot_reconstructs_with_masks,
@@ -173,6 +175,7 @@ def test_masked(
     all_latents = []
     all_channel_variances = []
     all_channel_maes = []
+    all_channel_mses = []
 
     with torch.no_grad():
         for idx, (img, channel_ids, panel_idx, img_path) in enumerate(
@@ -207,8 +210,19 @@ def test_masked(
                 dim=(0, 2, 3)
             )  # Mean variance per channel
             mae_per_channel = torch.abs(img - mi).mean(dim=(0, 2, 3))  # MAE per channel
+            mse_per_channel = torch.square(img - mi).mean(dim=(0, 2, 3))  # MSE per channel
             all_channel_variances.append(variance_per_channel.cpu())
             all_channel_maes.append(mae_per_channel.cpu())
+            all_channel_mses.append(mse_per_channel.cpu())
+
+            batch_var_mse_corr = torch.corrcoef(
+                torch.stack([variance_per_channel.cpu(), mse_per_channel.cpu()])
+            )[0, 1].item()
+            if math.isfinite(batch_var_mse_corr):
+                log_validation_batch_metrics(
+                    variance_mse_correlation_per_batch=batch_var_mse_corr,
+                    step=epoch * len(test_dataloader) + idx,
+                )
 
             loss = nll_loss(img, mi, logvar)
             running_loss += loss.item()
@@ -246,15 +260,18 @@ def test_masked(
     val_mae = running_mae / len(test_dataloader)
     val_mse = running_mse / len(test_dataloader)
 
-    all_latents = torch.cat(all_latents)
-    rankme = RankMe(all_latents)
+    latents = torch.cat(all_latents)
+    rankme = RankMe(latents)
 
-    # Calculate Pearson correlation between predicted variances and MAEs per channel
-    all_channel_variances = torch.cat(all_channel_variances)
-    all_channel_maes = torch.cat(all_channel_maes)
-    # Calculate Pearson correlation using flattened data across all batches
+    # Calculate Pearson correlation between predicted variances and MAEs/MSEs per channel
+    all_variances = torch.cat(all_channel_variances)
+    all_maes = torch.cat(all_channel_maes)
+    all_mses = torch.cat(all_channel_mses)
     variance_mae_corr = torch.corrcoef(
-        torch.stack([all_channel_variances.flatten(), all_channel_maes.flatten()])
+        torch.stack([all_variances.flatten(), all_maes.flatten()])
+    )[0, 1].item()
+    variance_mse_corr = torch.corrcoef(
+        torch.stack([all_variances.flatten(), all_mses.flatten()])
     )[0, 1].item()
 
     val_metrics = {
@@ -263,6 +280,7 @@ def test_masked(
         "val_mse": val_mse,
         "latent_rankme": rankme,
         "variance_mae_correlation": variance_mae_corr,
+        "variance_mse_correlation": variance_mse_corr,
         "epoch": epoch,
     }
 
@@ -273,6 +291,7 @@ def test_masked(
     print(f"MAE: {val_mae:.6f}")
     print(f"MSE: {val_mse:.6f}")
     print(f"Pearson MAE vs Var: {variance_mae_corr:.4f}")
+    print(f"Pearson MSE vs Var: {variance_mse_corr:.4f}")
     print("=" * 90)
     print()
 
diff --git a/train_masked_model_gp.py b/train_masked_model_gp.py
index c120afc..c8283a0 100644
--- a/train_masked_model_gp.py
+++ b/train_masked_model_gp.py
@@ -5,6 +5,8 @@
 log-likelihood loss that models spatial correlations using the GP covariance module.
 """
 
+import logging
+import math
 import os
 import sys
 
@@ -30,11 +32,13 @@
 from multiplex_model.losses import (
     HybridGPNLLLoss,
     HybridKroneckerGPNLLLoss,
+    HybridKroneckerMarkerGPNLLLoss,
     RankMe,
     beta_nll_loss,
     nll_loss,
 )
 from multiplex_model.modules.gp_covariance import (
+    KroneckerMarkerCovariance,
     KroneckerPlusSpatialCovariance,
     LowRankTimesSpatialCovariance,
 )
@@ -49,11 +53,15 @@
     get_scheduler_with_warmup,
     init_experiment,
     log_training_metrics,
+    log_validation_batch_metrics,
     log_validation_images,
     log_validation_metrics,
     plot_reconstructs_with_masks,
+    plot_reconstructs_with_uncertainty,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def train_masked_gp(
     model,
@@ -65,6 +73,7 @@ def train_masked_gp(
     marker_names_map,
     gp_covariance_module=None,
     use_gp_loss=True,
+    use_marker_covariance=False,
     lambda_gp=0.1,
     gp_max_cg_iterations=50,
     gp_downscale_factor=1,
@@ -117,7 +126,15 @@ def train_masked_gp(
     # Initialize GP loss if enabled
     gp_loss_fn = None
     if use_gp_loss and gp_covariance_module is not None:
-        if isinstance(gp_covariance_module, KroneckerPlusSpatialCovariance):
+        if isinstance(gp_covariance_module, KroneckerMarkerCovariance):
+            gp_loss_fn = HybridKroneckerMarkerGPNLLLoss(
+                covariance_module=gp_covariance_module,
+                lambda_gp=lambda_gp,
+                downscale_factor=gp_downscale_factor,
+                device=device,
+            )
+            print(f"Using Kronecker Marker GP loss with lambda_gp={lambda_gp}")
+        elif isinstance(gp_covariance_module, KroneckerPlusSpatialCovariance):
             gp_loss_fn = HybridKroneckerGPNLLLoss(
                 covariance_module=gp_covariance_module,
                 lambda_gp=lambda_gp,
@@ -171,8 +188,11 @@ def train_masked_gp(
             logvar = ClampWithGrad.apply(logvar, -15.0, 15.0)
 
             if use_gp_loss and gp_loss_fn is not None:
-                # Use hybrid GP loss
-                loss, loss_dict = gp_loss_fn(img, mi, logvar)
+                if use_marker_covariance:
+                    marker_emb = model.encoder.hyperkernel.hyperkernel_weights(channel_ids)
+                    loss, loss_dict = gp_loss_fn(img, mi, logvar, marker_emb)
+                else:
+                    loss, loss_dict = gp_loss_fn(img, mi, logvar)
                 
                 # Track loss components
                 for key in loss_dict:
@@ -182,6 +202,11 @@ def train_masked_gp(
                 loss = beta_nll_loss(img, mi, logvar, beta=beta)
                 epoch_loss_components["total_loss"].append(loss.item())
 
+            if not loss.isfinite():
+                logger.warning("Non-finite loss at step %d epoch %d, skipping batch", batch_idx, epoch)
+                optimizer.zero_grad()
+                continue
+
             scaler.scale(loss / gradient_accumulation_steps).backward()
 
             if (batch_idx + 1) % gradient_accumulation_steps == 0:
@@ -234,6 +259,7 @@ def train_masked_gp(
             mask_patch_size=mask_patch_size,
             marker_names_map=marker_names_map,
             use_gp_loss=use_gp_loss,
+            use_marker_covariance=use_marker_covariance,
         )
 
         # Save checkpoint
@@ -242,6 +268,7 @@ def train_masked_gp(
             "optimizer_state_dict": optimizer.state_dict(),
             "scheduler_state_dict": scheduler.state_dict(),
             "epoch": epoch,
+            "total_steps": total_steps,
         }
         if gp_covariance_module is not None:
              checkpoint["gp_covariance_state_dict"] = gp_covariance_module.state_dict()
@@ -274,6 +301,7 @@ def test_masked_gp(
     fully_masked_channels_max_frac=0.5,
     mask_patch_size=8,
     use_gp_loss=True,
+    use_marker_covariance=False,
 ):
     """
     Validation loop with optional GP loss evaluation.
@@ -294,6 +322,7 @@ def test_masked_gp(
     all_latents = []
     all_channel_variances = []
     all_channel_maes = []
+    all_channel_mses = []
 
     with torch.no_grad():
         for idx, (img, channel_ids, panel_idx, img_path) in enumerate(
@@ -319,6 +348,7 @@ def test_masked_gp(
             output = model.decode(latent, channel_ids)
             mi, logvar = output.unbind(dim=-1)
             mi = torch.sigmoid(mi)
+            logvar = torch.clamp(logvar, -15.0, 15.0)
 
             latent = normalize(latent.mean(dim=(2, 3)), p=2, dim=1)
             all_latents.append(latent.cpu())
@@ -326,14 +356,37 @@ def test_masked_gp(
             # Per-channel statistics
             variance_per_channel = torch.exp(logvar).mean(dim=(0, 2, 3))
             mae_per_channel = torch.abs(img - mi).mean(dim=(0, 2, 3))
+            mse_per_channel = torch.square(img - mi).mean(dim=(0, 2, 3))
             all_channel_variances.append(variance_per_channel.cpu())
             all_channel_maes.append(mae_per_channel.cpu())
+            all_channel_mses.append(mse_per_channel.cpu())
+
+            batch_var_mse_corr = torch.corrcoef(
+                torch.stack([variance_per_channel.cpu(), mse_per_channel.cpu()])
+            )[0, 1].item()
+            if math.isfinite(batch_var_mse_corr):
+                log_validation_batch_metrics(
+                    variance_mse_correlation_per_batch=batch_var_mse_corr,
+                    step=epoch * len(test_dataloader) + idx,
+                )
 
             # Compute loss
             if use_gp_loss and gp_loss_fn is not None:
-                loss, loss_dict = gp_loss_fn(img, mi, logvar)
+                if use_marker_covariance:
+                    marker_emb = model.encoder.hyperkernel.hyperkernel_weights(channel_ids)
+                    loss, loss_dict = gp_loss_fn(img, mi, logvar, marker_emb)
+                else:
+                    loss, loss_dict = gp_loss_fn(img, mi, logvar)
                 running_standard_nll += loss_dict["standard_nll"]
                 running_gp_nll += loss_dict["gp_nll"]
+                if use_marker_covariance and idx == 0:
+                    _, _, K_C = gp_covariance_module._compute_marker_eigen(marker_emb[0])
+                    eigvals = torch.linalg.eigvalsh(K_C)
+                    print(
+                        f"  Marker cov diagnostics — "
+                        f"min_eigval: {eigvals.min().item():.4f}, "
+                        f"condition_number: {(eigvals.max() / eigvals.min()).item():.2f}"
+                    )
             else:
                 loss = nll_loss(img, mi, logvar)
             
@@ -367,20 +420,44 @@ def test_masked_gp(
                     masked_channels_names=masked_channels_names,
                     img_idx=idx,
                 )
+
+                sigma = torch.exp(0.5 * logvar)
+                uncertainty_img = plot_reconstructs_with_uncertainty(
+                    img,
+                    mi,
+                    sigma,
+                    channel_ids,
+                    unactive_channels,
+                    markers_names_map=marker_names_map,
+                    ncols=9,
+                )
+                log_validation_images(
+                    fig=uncertainty_img,
+                    panel_idx=panel_idx[0],
+                    img_path=img_path[0],
+                    epoch=epoch,
+                    masked_channels_names=masked_channels_names,
+                    img_idx=idx,
+                    name_suffix="_sigma",
+                )
                 plt.close("all")
 
     val_loss = running_loss / len(test_dataloader)
     val_mae = running_mae / len(test_dataloader)
     val_mse = running_mse / len(test_dataloader)
 
-    all_latents = torch.cat(all_latents)
-    rankme = RankMe(all_latents)
+    latents = torch.cat(all_latents)
+    rankme = RankMe(latents)
 
-    # Variance-MAE correlation
-    all_channel_variances = torch.cat(all_channel_variances)
-    all_channel_maes = torch.cat(all_channel_maes)
+    # Variance-MAE/MSE correlation
+    all_variances = torch.cat(all_channel_variances)
+    all_maes = torch.cat(all_channel_maes)
+    all_mses = torch.cat(all_channel_mses)
     variance_mae_corr = torch.corrcoef(
-        torch.stack([all_channel_variances.flatten(), all_channel_maes.flatten()])
+        torch.stack([all_variances.flatten(), all_maes.flatten()])
+    )[0, 1].item()
+    variance_mse_corr = torch.corrcoef(
+        torch.stack([all_variances.flatten(), all_mses.flatten()])
     )[0, 1].item()
 
     val_metrics = {
@@ -389,6 +466,7 @@ def test_masked_gp(
         "val_mse": val_mse,
         "latent_rankme": rankme,
         "variance_mae_correlation": variance_mae_corr,
+        "variance_mse_correlation": variance_mse_corr,
         "epoch": epoch,
     }
     
@@ -406,6 +484,7 @@ def test_masked_gp(
     print(f"MAE: {val_mae:.6f}")
     print(f"MSE: {val_mse:.6f}")
     print(f"Pearson MAE vs Var: {variance_mae_corr:.4f}")
+    print(f"Pearson MSE vs Var: {variance_mse_corr:.4f}")
     print("=" * 90)
     print()
 
@@ -507,6 +586,9 @@ def test_masked_gp(
     gp_max_cg_iterations = getattr(config, "gp_max_cg_iterations", 50)
     gp_downscale_factor  = getattr(config, "gp_downscale_factor", 1)
     gp_learn_lengthscale = getattr(config, "gp_learn_lengthscale", False)
+    use_marker_covariance = getattr(config, "use_marker_covariance", False)
+    marker_embed_dim      = getattr(config, "marker_embed_dim", 32)
+    marker_jitter         = getattr(config, "marker_jitter", 1e-2)
 
     print("\nGP Loss Configuration:")
     print(f"  Use GP Loss:         {use_gp_loss}")
@@ -516,7 +598,10 @@ def test_masked_gp(
     print(f"  Lengthscale:         {gp_lengthscale}")
     print(f"  Max CG Iterations:   {gp_max_cg_iterations}  (ignored for Kronecker)")
     print(f"  Downscale Factor:    {gp_downscale_factor}")
-    print(f"  Learn Lengthscale:   {gp_learn_lengthscale}\n")
+    print(f"  Learn Lengthscale:   {gp_learn_lengthscale}  (ignored for Kronecker)")
+    print(f"  Use Marker Cov:      {use_marker_covariance}")
+    print(f"  Marker Embed Dim:    {marker_embed_dim}")
+    print(f"  Marker Jitter:       {marker_jitter}\n")
 
     # Initialize GP covariance module
     gp_covariance_module = None
@@ -525,11 +610,35 @@ def test_masked_gp(
         H_gp = H // gp_downscale_factor
         W_gp = W // gp_downscale_factor
 
-        if use_kronecker_gp:
-            # Kronecker requires square images after downscaling
+        if use_kronecker_gp and use_marker_covariance:
             assert H_gp == W_gp, (
                 f"Kronecker GP requires square spatial grid, "
-                f"got {H_gp}×{W_gp}. Adjust input_image_size or gp_downscale_factor."
+                f"got {H_gp}x{W_gp}. Adjust input_image_size or gp_downscale_factor."
+            )
+            # Compute hyperkernel_model_dim from encoder config
+            hk_cfg = config.encoder_config
+            if len(hk_cfg.ma_layers_blocks) == 0:
+                hk_input_dim = 1
+            else:
+                hk_input_dim = hk_cfg.ma_embedding_dims[-1]
+            hk_embed_dim = hk_cfg.pm_embedding_dims[0]
+            hk_kernel_size = hk_cfg.hyperkernel_config.kernel_size
+            hyperkernel_model_dim = hk_embed_dim * (hk_kernel_size ** 2) * hk_input_dim
+
+            gp_covariance_module = KroneckerMarkerCovariance(
+                grid_size=H_gp,
+                marker_embed_dim=marker_embed_dim,
+                hyperkernel_model_dim=hyperkernel_model_dim,
+                kernel_jitter=gp_kernel_jitter,
+                marker_jitter=marker_jitter,
+                spatial_matern_kernel_length_scale=gp_lengthscale,
+                device=device,
+            )
+            gp_covariance_module = gp_covariance_module.to(device)
+        elif use_kronecker_gp:
+            assert H_gp == W_gp, (
+                f"Kronecker GP requires square spatial grid, "
+                f"got {H_gp}x{W_gp}. Adjust input_image_size or gp_downscale_factor."
             )
             gp_covariance_module = KroneckerPlusSpatialCovariance(
                 grid_size=H_gp,
@@ -553,17 +662,42 @@ def test_masked_gp(
                 device=device,
             )
 
+    # Load checkpoint early to recover total_steps for scheduler reconstruction
+    start_epoch = 0
+    checkpoint = None
+    if config.resolve_checkpoint():
+        print(f"Loading model from checkpoint: {config.from_checkpoint}")
+        checkpoint = torch.load(config.from_checkpoint, map_location=device)
+        model.load_state_dict(checkpoint["model_state_dict"])
+        if gp_covariance_module is not None:
+            if "gp_covariance_state_dict" in checkpoint:
+                gp_covariance_module.load_state_dict(checkpoint["gp_covariance_state_dict"])
+            else:
+                logger.warning(
+                    "Checkpoint missing 'gp_covariance_state_dict' — "
+                    "KroneckerMarkerCovariance starts from random init"
+                )
+        start_epoch = checkpoint["epoch"] + 1
+
     # Optimizer and scheduler
-    total_steps = (
-        len(train_dataloader) * config.epochs // config.gradient_accumulation_steps
-    )
+    # When resuming normally, use saved total_steps so scheduler boundaries match original run.
+    # When reset_lr_schedule=True, recalculate from remaining epochs for a fresh cosine cycle
+    # that covers exactly the new training run (config.epochs - start_epoch epochs).
+    if checkpoint is not None and "total_steps" in checkpoint and not config.reset_lr_schedule:
+        total_steps = checkpoint["total_steps"]
+    else:
+        remaining_epochs = config.epochs - start_epoch
+        total_steps = len(train_dataloader) * remaining_epochs // config.gradient_accumulation_steps
     num_warmup_steps = int(total_steps * config.frac_warmup_steps)
     num_annealing_steps = total_steps - num_warmup_steps
 
     # Include GP covariance parameters in optimization if learnable
     params_to_optimize = list(model.parameters())
-    if use_gp_loss and gp_learn_lengthscale and gp_covariance_module is not None:
-        params_to_optimize += list(gp_covariance_module.parameters())
+    if use_gp_loss and gp_covariance_module is not None:
+        if use_marker_covariance:
+            params_to_optimize += list(gp_covariance_module.parameters())
+        elif gp_learn_lengthscale:
+            params_to_optimize += list(gp_covariance_module.parameters())
 
     optimizer = optim.AdamW(
         params_to_optimize,
@@ -579,6 +713,10 @@ def test_masked_gp(
         type="cosine",
     )
 
+    if checkpoint is not None and not config.reset_lr_schedule:
+        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+        scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
+
     # Initialize experiment tracking
     comet_config = config.model_dump()
     comet_config.update({
@@ -590,21 +728,12 @@ def test_masked_gp(
         "gp_max_cg_iterations": gp_max_cg_iterations,
         "gp_downscale_factor":  gp_downscale_factor,
         "gp_learn_lengthscale": gp_learn_lengthscale,
+        "use_marker_covariance": use_marker_covariance,
+        "marker_embed_dim":      marker_embed_dim,
+        "marker_jitter":         marker_jitter,
     })
     init_experiment(comet_config)
 
-    # Load checkpoint if specified
-    start_epoch = 0
-    if config.resolve_checkpoint():
-        print(f"Loading model from checkpoint: {config.from_checkpoint}")
-        checkpoint = torch.load(config.from_checkpoint, map_location=device)
-        model.load_state_dict(checkpoint["model_state_dict"])
-        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
-        scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
-        if gp_covariance_module is not None and "gp_covariance_state_dict" in checkpoint:
-             gp_covariance_module.load_state_dict(checkpoint["gp_covariance_state_dict"])
-        start_epoch = checkpoint["epoch"] + 1
-
     # Train the model
     train_masked_gp(
         model,
@@ -616,6 +745,7 @@ def test_masked_gp(
         marker_names_map=INV_TOKENIZER,
         gp_covariance_module=gp_covariance_module,
         use_gp_loss=use_gp_loss,
+        use_marker_covariance=use_marker_covariance,
         lambda_gp=lambda_gp,
         gp_max_cg_iterations=gp_max_cg_iterations,
         gp_downscale_factor=gp_downscale_factor,
diff --git a/train_masked_model_learnmask.py b/train_masked_model_learnmask.py
new file mode 100644
index 0000000..920ff82
--- /dev/null
+++ b/train_masked_model_learnmask.py
@@ -0,0 +1,440 @@
+import math
+import os
+import sys
+from typing import Any
+
+import comet_ml  # noqa: F401
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.optim as optim
+from ruamel.yaml import YAML
+from torch.amp import GradScaler, autocast
+from torch.nn.functional import normalize
+from torch.utils.data import DataLoader
+from torchvision.transforms import (
+    Compose,
+    RandomCrop,
+    RandomHorizontalFlip,
+    RandomRotation,
+)
+from torchvision.transforms.functional import InterpolationMode
+from tqdm import tqdm
+
+from multiplex_model.data import DatasetFromTIFF, PanelBatchSampler, TestCrop
+from multiplex_model.losses import RankMe, beta_nll_loss, nll_loss
+from multiplex_model.modules import MultiplexAutoencoder
+from multiplex_model.utils import (
+    ClampWithGrad,
+    TrainingConfig,
+    apply_channel_masking,
+    get_pixel_mask,
+    finish_experiment,
+    get_run_name,
+    get_scheduler_with_warmup,
+    init_experiment,
+    log_training_metrics,
+    log_validation_images,
+    log_validation_metrics,
+    plot_reconstructs_with_masks,
+)
+
+
+def train_masked(
+    model,
+    optimizer,
+    scheduler,
+    train_dataloader,
+    val_dataloader,
+    device,
+    marker_names_map,
+    epochs=10,
+    gradient_accumulation_steps=1,
+    beta=1.0,
+    min_channels_frac=0.75,
+    fully_masked_channels_max_frac=0.5,
+    spatial_masking_ratio=0.6,
+    mask_patch_size=8,
+    start_epoch=0,
+    save_checkpoint_every=5,
+    checkpoints_path="checkpoints",
+):
+    """Train a masked autoencoder (decode the remaining channels) with the given parameters."""
+    model.train()
+    scaler = GradScaler()
+    run_name = get_run_name()
+
+    if not os.path.exists(checkpoints_path):
+        os.makedirs(checkpoints_path, exist_ok=True)
+        print(f"Created checkpoints directory at {checkpoints_path}")
+
+    step = start_epoch * (len(train_dataloader) // gradient_accumulation_steps)
+    for epoch in range(start_epoch, epochs):
+        model.train()
+        for batch_idx, (img, channel_ids, panel_idx, img_path) in enumerate(
+            tqdm(train_dataloader, desc=f"Epoch {epoch}")
+        ):
+            img = img.to(device, dtype=torch.float32)
+            channel_ids = channel_ids.to(device, dtype=torch.long)
+
+            # Apply channel masking with channel subset sampling
+            img, channel_ids, masked_img, active_channel_ids = apply_channel_masking(
+                img,
+                channel_ids,
+                min_channels_frac,
+                fully_masked_channels_max_frac,
+                apply_channel_subset_sampling=True,
+            )
+
+            # Apply spatial masking
+            pixel_mask = get_pixel_mask(
+                masked_img, spatial_masking_ratio, mask_patch_size
+            )
+
+            with autocast(device_type="cuda", dtype=torch.bfloat16):
+                output = model(masked_img, active_channel_ids, channel_ids, spatial_mask=pixel_mask)["output"]
+                mi, logvar = output.unbind(dim=-1)
+                mi = torch.sigmoid(mi)
+                logvar = ClampWithGrad.apply(logvar, -15.0, 15.0)
+
+                loss = beta_nll_loss(img, mi, logvar, beta=beta)
+
+            scaler.scale(loss / gradient_accumulation_steps).backward()
+
+            if (batch_idx + 1) % gradient_accumulation_steps == 0:
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad()
+                scheduler.step()
+                mask_token = model.encoder.mask_token.item() if model.encoder.mask_token is not None else None
+
+                log_training_metrics(
+                    loss=loss.item(),
+                    lr=scheduler.get_last_lr()[0],
+                    mu=mi.mean().item(),
+                    logvar=logvar.mean().item(),
+                    mae=torch.abs(img - mi).mean().item(),
+                    mse=torch.square(img - mi).mean().item(),
+                    step=step,
+                    mask_token=mask_token,
+                )
+                step += 1
+
+        test_masked(
+            model,
+            val_dataloader,
+            device,
+            epoch,
+            spatial_masking_ratio=spatial_masking_ratio,
+            fully_masked_channels_max_frac=fully_masked_channels_max_frac,
+            mask_patch_size=mask_patch_size,
+            marker_names_map=marker_names_map,
+        )
+
+        checkpoint = {
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "scheduler_state_dict": scheduler.state_dict(),
+            "epoch": epoch,
+        }
+        if hasattr(model, "get_architecture_config"):
+            checkpoint["model_config"] = model.get_architecture_config()
+        if (epoch + 1) % save_checkpoint_every == 0:
+            torch.save(
+                checkpoint,
+                f"{checkpoints_path}/checkpoint-{run_name}-epoch_{epoch}.pth",
+            )
+        torch.save(checkpoint, f"{checkpoints_path}/last_checkpoint-{run_name}.pth")
+
+    final_model_path = f"{checkpoints_path}/final_model-{run_name}.pth"
+    print(f"Training completed. Saving final model at {final_model_path}...")
+    checkpoint = {
+        "model_state_dict": model.state_dict(),
+    }
+    if hasattr(model, "get_architecture_config"):
+        checkpoint["model_config"] = model.get_architecture_config()
+    torch.save(checkpoint, final_model_path)
+
+
+def test_masked(
+    model,
+    test_dataloader,
+    device,
+    epoch,
+    marker_names_map,
+    num_plots=4,
+    spatial_masking_ratio=0.6,
+    fully_masked_channels_max_frac=0.5,
+    mask_patch_size=8,
+):
+    model.eval()
+    running_loss = 0.0
+    running_mae = 0.0
+    running_mse = 0.0
+    plot_indices = np.random.choice(
+        np.arange(len(test_dataloader)), size=num_plots, replace=False
+    )
+    plot_indices = set(plot_indices)
+
+    all_latents: list[torch.Tensor] = []
+    all_channel_variances: list[torch.Tensor] = []
+    all_channel_maes: list[torch.Tensor] = []
+
+    with torch.no_grad():
+        for idx, (img, channel_ids, panel_idx, img_path) in enumerate(
+            tqdm(test_dataloader, desc=f"Testing epoch {epoch}")
+        ):
+            img = img.to(device, dtype=torch.float32)
+            channel_ids = channel_ids.to(device, dtype=torch.long)
+
+            # Apply channel masking (only full channel masking for validation, no channel dropping)
+            _, _, masked_img, active_channel_ids = apply_channel_masking(
+                img,
+                channel_ids,
+                fully_masked_channels_max_frac=fully_masked_channels_max_frac,
+                apply_channel_subset_sampling=False,
+            )
+
+            # Apply spatial masking
+            pixel_mask = get_pixel_mask(masked_img, spatial_masking_ratio, mask_patch_size)
+
+            latent = model.encode(masked_img, active_channel_ids, spatial_mask=pixel_mask)["output"]
+            output = model.decode(latent, channel_ids)
+            mi, logvar = output.unbind(dim=-1)
+            mi = torch.sigmoid(mi)
+
+            latent = normalize(latent.mean(dim=(2, 3)), p=2, dim=1)
+            all_latents.append(latent.cpu())
+
+            # Accumulate per-channel statistics for correlation analysis
+            variance_per_channel = torch.exp(logvar).mean(
+                dim=(0, 2, 3)
+            )  # Mean variance per channel
+            mae_per_channel = torch.abs(img - mi).mean(dim=(0, 2, 3))  # MAE per channel
+            all_channel_variances.append(variance_per_channel.cpu())
+            all_channel_maes.append(mae_per_channel.cpu())
+
+            loss = nll_loss(img, mi, logvar)
+            running_loss += loss.item()
+            running_mae += torch.abs(img - mi).mean().item()
+            running_mse += torch.square(img - mi).mean().item()
+
+            if idx in plot_indices:
+                unactive_channels = [
+                    i for i in channel_ids[0] if i not in active_channel_ids[0]
+                ]
+                masked_channels_names = " | ".join(
+                    [marker_names_map[i.item()] for i in unactive_channels]
+                )
+
+                reconstr_img = plot_reconstructs_with_masks(
+                    img,
+                    mi,
+                    pixel_mask,
+                    channel_ids,
+                    unactive_channels,
+                    markers_names_map=marker_names_map,
+                    ncols=9,
+                )
+                log_validation_images(
+                    fig=reconstr_img,
+                    panel_idx=panel_idx[0],
+                    img_path=img_path[0],
+                    epoch=epoch,
+                    masked_channels_names=masked_channels_names,
+                    img_idx=idx,
+                )
+                plt.close("all")
+
+    val_loss = running_loss / len(test_dataloader)
+    val_mae = running_mae / len(test_dataloader)
+    val_mse = running_mse / len(test_dataloader)
+
+    latents_cat = torch.cat(all_latents)
+    rankme = RankMe(latents_cat)
+
+    # Calculate Pearson correlation between predicted variances and MAEs per channel
+    channel_variances_cat = torch.cat(all_channel_variances)
+    channel_maes_cat = torch.cat(all_channel_maes)
+    # Calculate Pearson correlation using flattened data across all batches
+    variance_mae_corr = torch.corrcoef(
+        torch.stack([channel_variances_cat.flatten(), channel_maes_cat.flatten()])
+    )[0, 1].item()
+    if not math.isfinite(variance_mae_corr):
+        variance_mae_corr = float("nan")
+
+    val_metrics = {
+        "val_loss": val_loss,
+        "val_mae": val_mae,
+        "val_mse": val_mse,
+        "latent_rankme": rankme,
+        "variance_mae_correlation": variance_mae_corr,
+        "epoch": epoch,
+    }
+
+    log_validation_metrics(**val_metrics)
+
+    print(f"{'=' * 40} EPOCH {epoch + 1} {'=' * 40}")
+    print(f"NLL: {val_loss:.4f}")
+    print(f"MAE: {val_mae:.6f}")
+    print(f"MSE: {val_mse:.6f}")
+    print(f"Pearson MAE vs Var: {variance_mae_corr:.4f}")
+    print("=" * 90)
+    print()
+
+    return val_metrics
+
+
+if __name__ == "__main__":
+    # Load the configuration file
+    config_path = sys.argv[1]
+    yaml = YAML(typ="safe")
+    with open(config_path, "r") as f:
+        raw_config = yaml.load(f)
+
+    # Validate configuration using Pydantic model
+    config = TrainingConfig(**raw_config)
+
+    device = config.device
+    print(f"Using device: {device}")
+
+    SIZE = config.input_image_size
+    BATCH_SIZE = config.batch_size
+    NUM_WORKERS = config.num_workers
+
+    PANEL_CONFIG = YAML().load(open(config.panel_config))
+    TOKENIZER = YAML().load(open(config.tokenizer_config))
+    INV_TOKENIZER = {v: k for k, v in TOKENIZER.items()}
+
+    train_transform = Compose(
+        [
+            RandomRotation(180, interpolation=InterpolationMode.BILINEAR),
+            RandomCrop(SIZE),
+            RandomHorizontalFlip(),
+        ]
+    )
+
+    test_transform = TestCrop(SIZE[0])
+
+    train_dataset = DatasetFromTIFF(
+        panels_config=PANEL_CONFIG,
+        split="train",
+        marker_tokenizer=TOKENIZER,
+        transform=train_transform,
+        use_preprocessing=False,
+        use_median_denoising=False,
+        use_butterworth_filter=True,
+        use_minmax_normalization=False,
+        use_clip_normalization=True,
+        file_extension="npy",
+    )
+
+    test_dataset = DatasetFromTIFF(
+        panels_config=PANEL_CONFIG,
+        split="test",
+        marker_tokenizer=TOKENIZER,
+        transform=test_transform,
+        use_preprocessing=False,
+        use_median_denoising=False,
+        use_butterworth_filter=True,
+        use_minmax_normalization=False,
+        use_clip_normalization=True,
+        file_extension="npy",
+    )
+
+    train_batch_sampler = PanelBatchSampler(train_dataset, BATCH_SIZE)
+    test_batch_sampler = PanelBatchSampler(test_dataset, BATCH_SIZE, shuffle=False)
+
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_sampler=train_batch_sampler,
+        num_workers=NUM_WORKERS,
+        pin_memory=True,
+        persistent_workers=True,
+        prefetch_factor=4,
+    )
+    test_dataloader = DataLoader(
+        test_dataset,
+        batch_sampler=test_batch_sampler,
+        num_workers=NUM_WORKERS,
+        pin_memory=True,
+        persistent_workers=True,
+        prefetch_factor=4,
+    )
+
+    # Build model configuration
+    num_channels = len(TOKENIZER)
+    model_config: dict[str, Any] = {
+        "num_channels": num_channels,
+        "encoder_config": config.encoder_config.model_dump(),
+        "decoder_config": config.decoder_config.model_dump(),
+    }
+
+    # Load checkpoint if specified
+    start_epoch = 0
+    checkpoint = None
+    if config.resolve_checkpoint():
+        assert config.from_checkpoint is not None
+        print(f"Loading model from checkpoint: {config.from_checkpoint}")
+        checkpoint = torch.load(config.from_checkpoint, map_location=device)
+        model = MultiplexAutoencoder.load_from_checkpoint(
+            checkpoint,
+            model_config=model_config,
+        ).to(device)
+        start_epoch = checkpoint.get("epoch", -1) + 1
+    else:
+        model = MultiplexAutoencoder(**model_config).to(device)
+
+    # Setup optimizer and scheduler
+    total_steps = (
+        len(train_dataloader) * config.epochs // config.gradient_accumulation_steps
+    )
+    num_warmup_steps = int(total_steps * config.frac_warmup_steps)
+    num_annealing_steps = total_steps - num_warmup_steps
+
+    optimizer = optim.AdamW(
+        model.parameters(), lr=config.peak_lr, weight_decay=config.weight_decay
+    )
+    scheduler = get_scheduler_with_warmup(
+        optimizer,
+        num_warmup_steps,
+        num_annealing_steps,
+        final_lr=config.final_lr,
+        peak_lr=config.peak_lr,
+        type="cosine",
+    )
+
+    if checkpoint is not None:
+        if "optimizer_state_dict" in checkpoint:
+            optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+        if "scheduler_state_dict" in checkpoint:
+            scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
+
+    # Initialize Comet.ml experiment
+    comet_config = config.model_dump()
+    init_experiment(comet_config)
+
+    # Train the model
+    train_masked(
+        model,
+        optimizer,
+        scheduler,
+        train_dataloader,
+        test_dataloader,
+        device,
+        marker_names_map=INV_TOKENIZER,
+        epochs=config.epochs,
+        start_epoch=start_epoch,
+        gradient_accumulation_steps=config.gradient_accumulation_steps,
+        min_channels_frac=config.min_channels_frac,
+        spatial_masking_ratio=config.spatial_masking_ratio,
+        fully_masked_channels_max_frac=config.fully_masked_channels_max_frac,
+        mask_patch_size=config.mask_patch_size,
+        save_checkpoint_every=config.save_checkpoint_freq,
+        checkpoints_path=config.checkpoints_dir,
+        beta=config.beta,
+    )
+
+    finish_experiment()
diff --git a/train_masked_model_learnmask_gp.py b/train_masked_model_learnmask_gp.py
new file mode 100644
index 0000000..5616174
--- /dev/null
+++ b/train_masked_model_learnmask_gp.py
@@ -0,0 +1,634 @@
+"""Training script combining learnable spatial mask token with Kronecker marker GP loss.
+
+Merges the mask-token flow from `train_masked_model_learnmask.py` with the
+Kronecker + marker covariance GP loss from `train_masked_model_gp.py`.
+"""
+
+import logging
+import math
+import os
+import sys
+from typing import Any
+
+import comet_ml  # noqa: F401
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.optim as optim
+from ruamel.yaml import YAML
+from torch.amp import GradScaler, autocast
+from torch.nn.functional import normalize
+from torch.utils.data import DataLoader
+from torchvision.transforms import (
+    Compose,
+    RandomCrop,
+    RandomHorizontalFlip,
+    RandomRotation,
+)
+from torchvision.transforms.functional import InterpolationMode
+from tqdm import tqdm
+
+from multiplex_model.data import DatasetFromTIFF, PanelBatchSampler, TestCrop
+from multiplex_model.losses import (
+    HybridKroneckerMarkerGPNLLLoss,
+    RankMe,
+    beta_nll_loss,
+    nll_loss,
+)
+from multiplex_model.modules import MultiplexAutoencoder
+from multiplex_model.modules.gp_covariance import KroneckerMarkerCovariance
+from multiplex_model.utils import (
+    ClampWithGrad,
+    TrainingConfig,
+    apply_channel_masking,
+    finish_experiment,
+    get_pixel_mask,
+    get_run_name,
+    get_scheduler_with_warmup,
+    init_experiment,
+    log_training_metrics,
+    log_validation_batch_metrics,
+    log_validation_images,
+    log_validation_metrics,
+    plot_reconstructs_with_masks,
+    plot_reconstructs_with_uncertainty,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def train_masked_learnmask_gp(
+    model,
+    optimizer,
+    scheduler,
+    train_dataloader,
+    val_dataloader,
+    device,
+    marker_names_map,
+    gp_covariance_module,
+    gp_loss_fn,
+    total_steps,
+    use_gp_loss=True,
+    epochs=10,
+    gradient_accumulation_steps=1,
+    beta=1.0,
+    min_channels_frac=0.75,
+    fully_masked_channels_max_frac=0.5,
+    spatial_masking_ratio=0.6,
+    mask_patch_size=8,
+    start_epoch=0,
+    save_checkpoint_every=5,
+    checkpoints_path="checkpoints",
+):
+    model.train()
+    scaler = GradScaler()
+    run_name = get_run_name()
+
+    if not os.path.exists(checkpoints_path):
+        os.makedirs(checkpoints_path, exist_ok=True)
+        print(f"Created checkpoints directory at {checkpoints_path}")
+
+    step = start_epoch * (len(train_dataloader) // gradient_accumulation_steps)
+
+    for epoch in range(start_epoch, epochs):
+        model.train()
+        epoch_loss_components: dict[str, list[float]] = {
+            "standard_nll": [],
+            "gp_nll": [],
+            "total_loss": [],
+        }
+
+        for batch_idx, (img, channel_ids, panel_idx, img_path) in enumerate(
+            tqdm(train_dataloader, desc=f"Epoch {epoch}")
+        ):
+            img = img.to(device, dtype=torch.float32)
+            channel_ids = channel_ids.to(device, dtype=torch.long)
+
+            img, channel_ids, masked_img, active_channel_ids = apply_channel_masking(
+                img,
+                channel_ids,
+                min_channels_frac,
+                fully_masked_channels_max_frac,
+                apply_channel_subset_sampling=True,
+            )
+
+            pixel_mask = get_pixel_mask(masked_img, spatial_masking_ratio, mask_patch_size)
+
+            with autocast(device_type="cuda", dtype=torch.bfloat16):
+                output = model(
+                    masked_img, active_channel_ids, channel_ids, spatial_mask=pixel_mask
+                )["output"]
+                mi, logvar = output.unbind(dim=-1)
+                mi = torch.sigmoid(mi)
+                logvar = ClampWithGrad.apply(logvar, -15.0, 15.0)
+
+            # GP loss runs in float32: linalg.solve in Woodbury (gp_covariance.py)
+            # rejects mixed bfloat16/float32 dtypes used by the precomputed Kronecker eigs.
+            if use_gp_loss and gp_loss_fn is not None:
+                marker_emb = model.encoder.hyperkernel.hyperkernel_weights(channel_ids)
+                loss, loss_dict = gp_loss_fn(
+                    img.float(), mi.float(), logvar.float(), marker_emb.float()
+                )
+                for key in loss_dict:
+                    epoch_loss_components[key].append(loss_dict[key])
+            else:
+                loss = beta_nll_loss(img, mi, logvar, beta=beta)
+                epoch_loss_components["total_loss"].append(loss.item())
+
+            if not loss.isfinite():
+                logger.warning("Non-finite loss at step %d epoch %d, skipping batch", batch_idx, epoch)
+                optimizer.zero_grad()
+                continue
+
+            scaler.scale(loss / gradient_accumulation_steps).backward()
+
+            if (batch_idx + 1) % gradient_accumulation_steps == 0:
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad()
+                scheduler.step()
+
+                mask_token_value = model.encoder.mask_token.item() if model.encoder.mask_token is not None else None
+
+                metrics: dict[str, Any] = {
+                    "loss": loss.item(),
+                    "lr": scheduler.get_last_lr()[0],
+                    "mu": mi.mean().item(),
+                    "logvar": logvar.mean().item(),
+                    "mae": torch.abs(img - mi).mean().item(),
+                    "mse": torch.square(img - mi).mean().item(),
+                    "step": step,
+                    "mask_token": mask_token_value,
+                }
+                if use_gp_loss and epoch_loss_components["gp_nll"]:
+                    metrics["standard_nll"] = epoch_loss_components["standard_nll"][-1]
+                    metrics["gp_nll"] = epoch_loss_components["gp_nll"][-1]
+
+                log_training_metrics(**metrics)
+                step += 1
+
+        if use_gp_loss and epoch_loss_components["gp_nll"]:
+            avg_standard_nll = float(np.mean(epoch_loss_components["standard_nll"]))
+            avg_gp_nll = float(np.mean(epoch_loss_components["gp_nll"]))
+            avg_total = float(np.mean(epoch_loss_components["total_loss"]))
+            print(f"\nEpoch {epoch} Loss Components:")
+            print(f"  Standard NLL: {avg_standard_nll:.4f}")
+            print(f"  GP NLL:       {avg_gp_nll:.4f}")
+            print(f"  Total Loss:   {avg_total:.4f}")
+
+        test_masked_learnmask_gp(
+            model,
+            val_dataloader,
+            device,
+            epoch,
+            gp_covariance_module=gp_covariance_module,
+            gp_loss_fn=gp_loss_fn,
+            spatial_masking_ratio=spatial_masking_ratio,
+            fully_masked_channels_max_frac=fully_masked_channels_max_frac,
+            mask_patch_size=mask_patch_size,
+            marker_names_map=marker_names_map,
+            use_gp_loss=use_gp_loss,
+        )
+
+        checkpoint = {
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "scheduler_state_dict": scheduler.state_dict(),
+            "epoch": epoch,
+            "total_steps": total_steps,
+        }
+        if gp_covariance_module is not None:
+            checkpoint["gp_covariance_state_dict"] = gp_covariance_module.state_dict()
+        if hasattr(model, "get_architecture_config"):
+            checkpoint["model_config"] = model.get_architecture_config()
+
+        if (epoch + 1) % save_checkpoint_every == 0:
+            torch.save(checkpoint, f"{checkpoints_path}/checkpoint-{run_name}-epoch_{epoch}.pth")
+        torch.save(checkpoint, f"{checkpoints_path}/last_checkpoint-{run_name}.pth")
+
+    final_model_path = f"{checkpoints_path}/final_model-{run_name}.pth"
+    print(f"Training completed. Saving final model at {final_model_path}...")
+    final_checkpoint: dict[str, Any] = {"model_state_dict": model.state_dict()}
+    if hasattr(model, "get_architecture_config"):
+        final_checkpoint["model_config"] = model.get_architecture_config()
+    if gp_covariance_module is not None:
+        final_checkpoint["gp_covariance_state_dict"] = gp_covariance_module.state_dict()
+    torch.save(final_checkpoint, final_model_path)
+
+
+def test_masked_learnmask_gp(
+    model,
+    test_dataloader,
+    device,
+    epoch,
+    gp_covariance_module,
+    gp_loss_fn,
+    marker_names_map,
+    num_plots=4,
+    spatial_masking_ratio=0.6,
+    fully_masked_channels_max_frac=0.5,
+    mask_patch_size=8,
+    use_gp_loss=True,
+):
+    model.eval()
+    running_loss = 0.0
+    running_mae = 0.0
+    running_mse = 0.0
+    running_standard_nll = 0.0
+    running_gp_nll = 0.0
+
+    plot_indices = np.random.choice(
+        np.arange(len(test_dataloader)), size=num_plots, replace=False
+    )
+    plot_indices = set(plot_indices)
+
+    all_latents: list[torch.Tensor] = []
+    all_channel_variances: list[torch.Tensor] = []
+    all_channel_maes: list[torch.Tensor] = []
+    all_channel_mses: list[torch.Tensor] = []
+
+    with torch.no_grad():
+        for idx, (img, channel_ids, panel_idx, img_path) in enumerate(
+            tqdm(test_dataloader, desc=f"Testing epoch {epoch}")
+        ):
+            img = img.to(device, dtype=torch.float32)
+            channel_ids = channel_ids.to(device, dtype=torch.long)
+
+            _, _, masked_img, active_channel_ids = apply_channel_masking(
+                img,
+                channel_ids,
+                fully_masked_channels_max_frac=fully_masked_channels_max_frac,
+                apply_channel_subset_sampling=False,
+            )
+
+            pixel_mask = get_pixel_mask(masked_img, spatial_masking_ratio, mask_patch_size)
+
+            latent = model.encode(masked_img, active_channel_ids, spatial_mask=pixel_mask)["output"]
+            output = model.decode(latent, channel_ids)
+            mi, logvar = output.unbind(dim=-1)
+            mi = torch.sigmoid(mi)
+            logvar = torch.clamp(logvar, -15.0, 15.0)
+
+            latent = normalize(latent.mean(dim=(2, 3)), p=2, dim=1)
+            all_latents.append(latent.cpu())
+
+            variance_per_channel = torch.exp(logvar).mean(dim=(0, 2, 3))
+            mae_per_channel = torch.abs(img - mi).mean(dim=(0, 2, 3))
+            mse_per_channel = torch.square(img - mi).mean(dim=(0, 2, 3))
+            all_channel_variances.append(variance_per_channel.cpu())
+            all_channel_maes.append(mae_per_channel.cpu())
+            all_channel_mses.append(mse_per_channel.cpu())
+
+            batch_var_mse_corr = torch.corrcoef(
+                torch.stack([variance_per_channel.cpu(), mse_per_channel.cpu()])
+            )[0, 1].item()
+            if math.isfinite(batch_var_mse_corr):
+                log_validation_batch_metrics(
+                    variance_mse_correlation_per_batch=batch_var_mse_corr,
+                    step=epoch * len(test_dataloader) + idx,
+                )
+
+            if use_gp_loss and gp_loss_fn is not None:
+                marker_emb = model.encoder.hyperkernel.hyperkernel_weights(channel_ids)
+                loss, loss_dict = gp_loss_fn(img, mi, logvar, marker_emb)
+                running_standard_nll += loss_dict["standard_nll"]
+                running_gp_nll += loss_dict["gp_nll"]
+                if idx == 0:
+                    _, _, K_C = gp_covariance_module._compute_marker_eigen(marker_emb[0])
+                    eigvals = torch.linalg.eigvalsh(K_C)
+                    print(
+                        f"  Marker cov diagnostics — "
+                        f"min_eigval: {eigvals.min().item():.4f}, "
+                        f"condition_number: {(eigvals.max() / eigvals.min()).item():.2f}"
+                    )
+            else:
+                loss = nll_loss(img, mi, logvar)
+
+            running_loss += loss.item()
+            running_mae += torch.abs(img - mi).mean().item()
+            running_mse += torch.square(img - mi).mean().item()
+
+            if idx in plot_indices:
+                unactive_channels = [
+                    i for i in channel_ids[0] if i not in active_channel_ids[0]
+                ]
+                masked_channels_names = " | ".join(
+                    [marker_names_map[i.item()] for i in unactive_channels]
+                )
+
+                reconstr_img = plot_reconstructs_with_masks(
+                    img,
+                    mi,
+                    pixel_mask,
+                    channel_ids,
+                    unactive_channels,
+                    markers_names_map=marker_names_map,
+                    ncols=9,
+                )
+                log_validation_images(
+                    fig=reconstr_img,
+                    panel_idx=panel_idx[0],
+                    img_path=img_path[0],
+                    epoch=epoch,
+                    masked_channels_names=masked_channels_names,
+                    img_idx=idx,
+                )
+
+                sigma = torch.exp(0.5 * logvar)
+                uncertainty_img = plot_reconstructs_with_uncertainty(
+                    img,
+                    mi,
+                    sigma,
+                    channel_ids,
+                    unactive_channels,
+                    markers_names_map=marker_names_map,
+                    ncols=9,
+                )
+                log_validation_images(
+                    fig=uncertainty_img,
+                    panel_idx=panel_idx[0],
+                    img_path=img_path[0],
+                    epoch=epoch,
+                    masked_channels_names=masked_channels_names,
+                    img_idx=idx,
+                    name_suffix="_sigma",
+                )
+                plt.close("all")
+
+    val_loss = running_loss / len(test_dataloader)
+    val_mae = running_mae / len(test_dataloader)
+    val_mse = running_mse / len(test_dataloader)
+
+    latents_cat = torch.cat(all_latents)
+    rankme = RankMe(latents_cat)
+
+    all_variances = torch.cat(all_channel_variances)
+    all_maes = torch.cat(all_channel_maes)
+    all_mses = torch.cat(all_channel_mses)
+    variance_mae_corr = torch.corrcoef(
+        torch.stack([all_variances.flatten(), all_maes.flatten()])
+    )[0, 1].item()
+    variance_mse_corr = torch.corrcoef(
+        torch.stack([all_variances.flatten(), all_mses.flatten()])
+    )[0, 1].item()
+
+    val_metrics: dict[str, Any] = {
+        "val_loss": val_loss,
+        "val_mae": val_mae,
+        "val_mse": val_mse,
+        "latent_rankme": rankme,
+        "variance_mae_correlation": variance_mae_corr,
+        "variance_mse_correlation": variance_mse_corr,
+        "epoch": epoch,
+    }
+    if use_gp_loss and gp_loss_fn is not None:
+        val_metrics["val_standard_nll"] = running_standard_nll / len(test_dataloader)
+        val_metrics["val_gp_nll"] = running_gp_nll / len(test_dataloader)
+
+    log_validation_metrics(**val_metrics)
+
+    print(f"{'=' * 40} EPOCH {epoch + 1} {'=' * 40}")
+    print(f"Total Loss: {val_loss:.4f}")
+    if use_gp_loss and gp_loss_fn is not None:
+        print(f"Standard NLL: {val_metrics['val_standard_nll']:.4f}")
+        print(f"GP NLL:       {val_metrics['val_gp_nll']:.4f}")
+    print(f"MAE: {val_mae:.6f}")
+    print(f"MSE: {val_mse:.6f}")
+    print(f"Pearson MAE vs Var: {variance_mae_corr:.4f}")
+    print(f"Pearson MSE vs Var: {variance_mse_corr:.4f}")
+    print("=" * 90)
+    print()
+
+    return val_metrics
+
+
+if __name__ == "__main__":
+    config_path = sys.argv[1]
+    yaml = YAML(typ="safe")
+    with open(config_path, "r") as f:
+        raw_config = yaml.load(f)
+
+    config = TrainingConfig(**raw_config)
+
+    device = config.device
+    print(f"Using device: {device}")
+
+    SIZE = config.input_image_size
+    BATCH_SIZE = config.batch_size
+    NUM_WORKERS = config.num_workers
+
+    PANEL_CONFIG = YAML().load(open(config.panel_config))
+    TOKENIZER = YAML().load(open(config.tokenizer_config))
+    INV_TOKENIZER = {v: k for k, v in TOKENIZER.items()}
+
+    train_transform = Compose(
+        [
+            RandomRotation(180, interpolation=InterpolationMode.BILINEAR),
+            RandomCrop(SIZE),
+            RandomHorizontalFlip(),
+        ]
+    )
+    test_transform = TestCrop(SIZE[0])
+
+    train_dataset = DatasetFromTIFF(
+        panels_config=PANEL_CONFIG,
+        split="train",
+        marker_tokenizer=TOKENIZER,
+        transform=train_transform,
+        use_preprocessing=False,
+        use_median_denoising=False,
+        use_butterworth_filter=True,
+        use_minmax_normalization=False,
+        use_clip_normalization=True,
+        file_extension="npy",
+    )
+    test_dataset = DatasetFromTIFF(
+        panels_config=PANEL_CONFIG,
+        split="test",
+        marker_tokenizer=TOKENIZER,
+        transform=test_transform,
+        use_preprocessing=False,
+        use_median_denoising=False,
+        use_butterworth_filter=True,
+        use_minmax_normalization=False,
+        use_clip_normalization=True,
+        file_extension="npy",
+    )
+
+    train_batch_sampler = PanelBatchSampler(train_dataset, BATCH_SIZE)
+    test_batch_sampler = PanelBatchSampler(test_dataset, BATCH_SIZE, shuffle=False)
+
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_sampler=train_batch_sampler,
+        num_workers=NUM_WORKERS,
+        pin_memory=True,
+        persistent_workers=True,
+        prefetch_factor=4,
+    )
+    test_dataloader = DataLoader(
+        test_dataset,
+        batch_sampler=test_batch_sampler,
+        num_workers=NUM_WORKERS,
+        pin_memory=True,
+        persistent_workers=True,
+        prefetch_factor=4,
+    )
+
+    num_channels = len(TOKENIZER)
+    model_config: dict[str, Any] = {
+        "num_channels": num_channels,
+        "encoder_config": config.encoder_config.model_dump(),
+        "decoder_config": config.decoder_config.model_dump(),
+    }
+
+    use_gp_loss = getattr(config, "use_gp_loss", False)
+    use_kronecker_gp = getattr(config, "use_kronecker_gp", False)
+    use_marker_covariance = getattr(config, "use_marker_covariance", False)
+    lambda_gp = getattr(config, "lambda_gp", 0.1)
+    gp_kernel_jitter = getattr(config, "gp_kernel_jitter", 1e-2)
+    gp_lengthscale = getattr(config, "gp_lengthscale", 5.0)
+    gp_downscale_factor = getattr(config, "gp_downscale_factor", 1)
+    marker_embed_dim = getattr(config, "marker_embed_dim", 32)
+    marker_jitter = getattr(config, "marker_jitter", 1e-2)
+
+    assert use_gp_loss and use_kronecker_gp and use_marker_covariance, (
+        "This script combines learnmask with Kronecker marker GP loss. "
+        "Set use_gp_loss=true, use_kronecker_gp=true, use_marker_covariance=true."
+    )
+
+    print("\nGP Loss Configuration:")
+    print(f"  Lambda GP:        {lambda_gp}")
+    print(f"  Kernel Jitter:    {gp_kernel_jitter}")
+    print(f"  Lengthscale:      {gp_lengthscale}")
+    print(f"  Downscale Factor: {gp_downscale_factor}")
+    print(f"  Marker Embed Dim: {marker_embed_dim}")
+    print(f"  Marker Jitter:    {marker_jitter}\n")
+
+    H, W = SIZE
+    H_gp = H // gp_downscale_factor
+    W_gp = W // gp_downscale_factor
+    assert H_gp == W_gp, (
+        f"Kronecker GP requires square spatial grid, got {H_gp}x{W_gp}."
+    )
+
+    hk_cfg = config.encoder_config
+    if len(hk_cfg.ma_layers_blocks) == 0:
+        hk_input_dim = 1
+    else:
+        hk_input_dim = hk_cfg.ma_embedding_dims[-1]
+    hk_embed_dim = hk_cfg.pm_embedding_dims[0]
+    hk_kernel_size = hk_cfg.hyperkernel_config.kernel_size
+    hyperkernel_model_dim = hk_embed_dim * (hk_kernel_size ** 2) * hk_input_dim
+
+    gp_covariance_module = KroneckerMarkerCovariance(
+        grid_size=H_gp,
+        marker_embed_dim=marker_embed_dim,
+        hyperkernel_model_dim=hyperkernel_model_dim,
+        kernel_jitter=gp_kernel_jitter,
+        marker_jitter=marker_jitter,
+        spatial_matern_kernel_length_scale=gp_lengthscale,
+        device=device,
+    ).to(device)
+
+    gp_loss_fn = HybridKroneckerMarkerGPNLLLoss(
+        covariance_module=gp_covariance_module,
+        lambda_gp=lambda_gp,
+        downscale_factor=gp_downscale_factor,
+        device=device,
+    )
+    print(f"Using Kronecker Marker GP loss with lambda_gp={lambda_gp}")
+
+    start_epoch = 0
+    checkpoint = None
+    if config.resolve_checkpoint():
+        assert config.from_checkpoint is not None
+        print(f"Loading model from checkpoint: {config.from_checkpoint}")
+        checkpoint = torch.load(config.from_checkpoint, map_location=device)
+        model = MultiplexAutoencoder.load_from_checkpoint(
+            checkpoint, model_config=model_config
+        ).to(device)
+        if "gp_covariance_state_dict" in checkpoint:
+            gp_covariance_module.load_state_dict(checkpoint["gp_covariance_state_dict"])
+        else:
+            logger.warning(
+                "Checkpoint missing 'gp_covariance_state_dict' — "
+                "KroneckerMarkerCovariance starts from random init"
+            )
+        start_epoch = checkpoint.get("epoch", -1) + 1
+    else:
+        model = MultiplexAutoencoder(**model_config).to(device)
+
+    # When extending training (bumping config.epochs), use reset_lr_schedule: true to get
+    # a fresh cosine cycle. Without it, total_steps is reused from the checkpoint, and if
+    # config.epochs > original epochs the scheduler may be past its annealing boundary.
+    if checkpoint is not None and "total_steps" in checkpoint and not config.reset_lr_schedule:
+        total_steps = checkpoint["total_steps"]
+    else:
+        remaining_epochs = config.epochs - start_epoch
+        total_steps = len(train_dataloader) * remaining_epochs // config.gradient_accumulation_steps
+    num_warmup_steps = int(total_steps * config.frac_warmup_steps)
+    num_annealing_steps = total_steps - num_warmup_steps
+
+    params_to_optimize = list(model.parameters()) + list(gp_covariance_module.parameters())
+    optimizer = optim.AdamW(
+        params_to_optimize, lr=config.peak_lr, weight_decay=config.weight_decay
+    )
+    scheduler = get_scheduler_with_warmup(
+        optimizer,
+        num_warmup_steps,
+        num_annealing_steps,
+        final_lr=config.final_lr,
+        peak_lr=config.peak_lr,
+        type="cosine",
+    )
+
+    if checkpoint is not None and not config.reset_lr_schedule:
+        if "optimizer_state_dict" in checkpoint:
+            optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+        if "scheduler_state_dict" in checkpoint:
+            scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
+
+    comet_config = config.model_dump()
+    comet_config.update(
+        {
+            "use_gp_loss": use_gp_loss,
+            "use_kronecker_gp": use_kronecker_gp,
+            "use_marker_covariance": use_marker_covariance,
+            "lambda_gp": lambda_gp,
+            "gp_kernel_jitter": gp_kernel_jitter,
+            "gp_lengthscale": gp_lengthscale,
+            "gp_downscale_factor": gp_downscale_factor,
+            "marker_embed_dim": marker_embed_dim,
+            "marker_jitter": marker_jitter,
+        }
+    )
+    init_experiment(comet_config)
+
+    train_masked_learnmask_gp(
+        model,
+        optimizer,
+        scheduler,
+        train_dataloader,
+        test_dataloader,
+        device,
+        marker_names_map=INV_TOKENIZER,
+        gp_covariance_module=gp_covariance_module,
+        gp_loss_fn=gp_loss_fn,
+        total_steps=total_steps,
+        use_gp_loss=use_gp_loss,
+        epochs=config.epochs,
+        start_epoch=start_epoch,
+        gradient_accumulation_steps=config.gradient_accumulation_steps,
+        min_channels_frac=config.min_channels_frac,
+        spatial_masking_ratio=config.spatial_masking_ratio,
+        fully_masked_channels_max_frac=config.fully_masked_channels_max_frac,
+        mask_patch_size=config.mask_patch_size,
+        save_checkpoint_every=config.save_checkpoint_freq,
+        checkpoints_path=config.checkpoints_dir,
+        beta=config.beta,
+    )
+
+    finish_experiment()