diff --git a/.gitignore b/.gitignore
index ed95289..83d8d98 100644
--- a/.gitignore
+++ b/.gitignore
@@ -224,44 +224,7 @@ vscode/
 *.svg
 docs/research_proposal.tex
 .vscode/settings.json
-hest_data/.gitattributes
-hest_data/HEST_v1_3_0.csv
-hest_data/README.md
-hest_data/cellvit_seg/INT1_cellvit_seg.geojson.zip
-hest_data/cellvit_seg/INT1_cellvit_seg.parquet
-hest_data/cellvit_seg/INT10_cellvit_seg.geojson.zip
-hest_data/cellvit_seg/INT10_cellvit_seg.parquet
-hest_data/cellvit_seg/INT11_cellvit_seg.geojson.zip
-hest_data/cellvit_seg/INT11_cellvit_seg.parquet
-hest_data/cellvit_seg/INT12_cellvit_seg.geojson.zip
-hest_data/cellvit_seg/INT12_cellvit_seg.parquet
-hest_data/cellvit_seg/INT13_cellvit_seg.geojson.zip
-hest_data/cellvit_seg/INT13_cellvit_seg.parquet
-hest_data/cellvit_seg/INT16_cellvit_seg.geojson.zip
-hest_data/cellvit_seg/INT16_cellvit_seg.parquet
-hest_data/cellvit_seg/INT19_cellvit_seg.geojson.zip
-hest_data/cellvit_seg/INT19_cellvit_seg.parquet
-hest_data/cellvit_seg/INT20_cellvit_seg.geojson.zip
-hest_data/cellvit_seg/INT20_cellvit_seg.parquet
-hest_data/cellvit_seg/INT21_cellvit_seg.geojson.zip
-hest_data/cellvit_seg/INT21_cellvit_seg.parquet
-hest_data/cellvit_seg/INT1_cellvit_seg.geojson
-hest_data/cellvit_seg/INT10_cellvit_seg.geojson
-hest_data/cellvit_seg/INT11_cellvit_seg.geojson
-hest_data/cellvit_seg/INT12_cellvit_seg.geojson
-hest_data/cellvit_seg/INT13_cellvit_seg.geojson
-hest_data/cellvit_seg/INT16_cellvit_seg.geojson
-hest_data/cellvit_seg/INT19_cellvit_seg.geojson
-hest_data/cellvit_seg/INT20_cellvit_seg.geojson
-hest_data/cellvit_seg/INT21_cellvit_seg.geojson
-hest_data/cellvit_seg/TENX175_cellvit_seg.geojson
-hest_data/cellvit_seg/TENX175_cellvit_seg.geojson.zip
-hest_data/cellvit_seg/TENX175_cellvit_seg.parquet
-hest_data/metadata/TENX175.json
-hest_data/patches/TENX175.h5
-hest_data/st/TENX175.h5ad
-hest_data/tissue_seg/TENX175_contours.geojson
-hest_data/wsis/TENX175.tif
+hest_data/*
 .idea/
 .gemini/settings.json
 .gemini/agents/literature-search.md
diff --git a/docs/TRAINING_GUIDE.md b/docs/TRAINING_GUIDE.md
index c6a71a9..2f47738 100644
--- a/docs/TRAINING_GUIDE.md
+++ b/docs/TRAINING_GUIDE.md
@@ -182,31 +182,43 @@ Available interaction tokens: `p2p`, `p2h`, `h2p`, `h2h`. Default is all four (F
 
 ---
 
-## 4. HPC Batch Experiments
+## 4. HPC Batch Experiments (Model Comparison)
 
-The `hpc/array_train.slurm` script runs all three whole-slide experiments as a SLURM array job:
+The `scripts/compare_models.py` script provides a unified mechanism to train all 5 models (SpatialTranscriptFormer, AttentionMIL, TransMIL, HE2RNA, and ViT-ST) under identical conditions (using precomputed features, identical seeds/splits, and the `mse_ccc` loss) and aggregate their results into a comparison table and chart.
 
-| Index | Model | Supervision | Key Flags |
-| :--- | :--- | :--- | :--- |
-| 0 | SpatialTranscriptFormer | Dense | `--whole-slide` |
-| 1 | AttentionMIL | Weak | `--whole-slide --weak-supervision` |
-| 2 | TransMIL | Weak | `--whole-slide --weak-supervision` |
-
-Submit with:
+### Running Locally
+To run all experiments sequentially on your local machine:
+```bash
+python scripts/compare_models.py --data-dir hest_data --output-dir runs/comparison
+```
 
+### Running on Slurm HPC
+To generate Slurm job submission scripts for the cluster:
 ```bash
-sbatch hpc/array_train.slurm
+python scripts/compare_models.py --data-dir /path/to/linux/hest_data --output-dir runs/comparison --slurm
 ```
 
-### Collecting Results (Currently broken!)
+> [!IMPORTANT]
+> **Windows Path Check**: If `--slurm` is active, the script will validation-check the `--data-dir` and throw an error if a Windows drive letter or backslashes are present, preventing job failures at discovery.
+> **Invoking Directory**: Note that relative log and output paths rely on `sbatch` being run from the project root directory.
+
+To automatically submit these scripts to the Slurm scheduler:
+```bash
+python scripts/compare_models.py --data-dir /path/to/linux/hest_data --output-dir runs/comparison --slurm --submit
+```
 
-After experiments complete, aggregate all `results_summary.json` files into a comparison table:
+This generates and submits 5 training jobs, plus a collection job. The collection job is configured with a Slurm dependency (`afterany`) on all 5 training runs, so it automatically aggregates results as soon as they complete.
 
+### Collecting Results Manually
+If you ran jobs individually or need to rebuild the report/plots from existing runs, run the script in collection-only mode:
 ```bash
-python hpc/collect_results.py --results-dir runs/ws_experiments
+python scripts/compare_models.py --collect-only --output-dir runs/comparison
 ```
 
-This produces a sorted comparison table and `comparison.csv`.
+This parses the `training_logs.sqlite` and `results_summary.json` inside each subdirectory and writes:
+1. `comparison_report.md` (Markdown summary table)
+2. `comparison_results.csv` (CSV statistics)
+3. `comparison_chart.png` (PCC vs. CCC and MAE comparison plots)
 
 ---
 
diff --git a/scripts/compare_models.py b/scripts/compare_models.py
index 251d4a3..7d357ea 100644
--- a/scripts/compare_models.py
+++ b/scripts/compare_models.py
@@ -244,21 +244,32 @@ def main():
     parser.add_argument(
         "--epochs-patch",
         type=int,
-        default=5,
+        default=30,
         help="Number of training epochs for patch baselines",
     )
     parser.add_argument(
         "--epochs-mil",
         type=int,
-        default=20,
+        default=30,
         help="Number of training epochs for MIL/STF models",
     )
+    parser.add_argument(
+        "--early-stopping-patience",
+        type=int,
+        default=10,
+        help="Patience epochs of no val_ccc improvement before stopping",
+    )
 
     # Slurm & Execution Options
+    parser.add_argument(
+        "--no-mil",
+        action="store_true",
+        help="Exclude Multiple-Instance Learning (MIL) models from the comparison",
+    )
     parser.add_argument(
         "--slurm",
         action="store_true",
-        help="Generate Slurm scripts instead of running locally",
+        help="Generate Slurm scripts instead of running locally. Note: relative paths assume sbatch is run from the project root.",
     )
     parser.add_argument(
         "--submit",
@@ -322,6 +333,15 @@ def main():
     )
     args = parser.parse_args()
 
+    if args.slurm:
+        import re
+
+        if re.match(r"^[a-zA-Z]:", args.data_dir) or "\\" in args.data_dir:
+            parser.error(
+                f"Data directory '{args.data_dir}' appears to be a Windows path (starts with drive letter or contains backslashes). "
+                "For Slurm execution, please provide a valid Linux/POSIX path to the data directory using --data-dir."
+            )
+
     # Determine hyperparameters based on --quick-test
     max_samples = 5 if args.quick_test else None
     epochs_patch = 1 if args.quick_test else args.epochs_patch
@@ -333,7 +353,10 @@ def main():
             "--model",
             "he2rna",
             "--backbone",
-            "resnet50",
+            "ctranspath",
+            "--precomputed",
+            "--loss",
+            "mse_ccc",
             "--batch-size",
             "64",
             "--epochs",
@@ -345,7 +368,10 @@ def main():
             "--model",
             "vit_st",
             "--backbone",
-            "vit_b_16",
+            "ctranspath",
+            "--precomputed",
+            "--loss",
+            "mse_ccc",
             "--batch-size",
             "32",
             "--epochs",
@@ -361,6 +387,8 @@ def main():
             "--whole-slide",
             "--precomputed",
             "--weak-supervision",
+            "--loss",
+            "mse_ccc",
             "--use-amp",
             "--batch-size",
             "1",
@@ -377,6 +405,8 @@ def main():
             "--whole-slide",
             "--precomputed",
             "--weak-supervision",
+            "--loss",
+            "mse_ccc",
             "--use-amp",
             "--batch-size",
             "1",
@@ -410,6 +440,10 @@ def main():
         ],
     }
 
+    if args.no_mil:
+        configs.pop("AttentionMIL", None)
+        configs.pop("TransMIL", None)
+
     # Add shared flags to all configs
     for name, cmd in configs.items():
         cmd.extend(
@@ -429,6 +463,8 @@ def main():
 
         if max_samples:
             cmd.extend(["--max-samples", str(max_samples)])
+        if args.early_stopping_patience is not None:
+            cmd.extend(["--early-stopping-patience", str(args.early_stopping_patience)])
 
     results = {}
 
@@ -454,7 +490,6 @@ def main():
 
         job_ids = []
         slurm_script_paths = {}
-        working_dir = pathlib.Path(os.getcwd()).as_posix()
 
         print(f"\nGenerating Slurm scripts in: {slurm_scripts_dir}")
 
@@ -491,6 +526,8 @@ def main():
 #SBATCH --cpus-per-task={args.slurm_cpus}
 #SBATCH --mem={args.slurm_mem}
 
+set -e
+
 {setup_cmds}
 
 # Change to submit directory (project root)
@@ -521,6 +558,8 @@ def main():
 #SBATCH --cpus-per-task=1
 #SBATCH --mem=8G
 
+set -e
+
 {setup_cmds}
 
 # Change to submit directory (project root)
diff --git a/src/spatial_transcript_former/models/__init__.py b/src/spatial_transcript_former/models/__init__.py
index f34a4b3..1ffec6b 100644
--- a/src/spatial_transcript_former/models/__init__.py
+++ b/src/spatial_transcript_former/models/__init__.py
@@ -1,4 +1,4 @@
 # Models directory
 from .mil import AttentionMIL, TransMIL
-from .regression import HE2RNA, ViT_ST
+from .regression import HE2RNA, ViT_ST, LinearProbe, MLPProbe
 from .interaction import SpatialTranscriptFormer
diff --git a/src/spatial_transcript_former/models/interaction.py b/src/spatial_transcript_former/models/interaction.py
index 03ccbd9..d4323cb 100644
--- a/src/spatial_transcript_former/models/interaction.py
+++ b/src/spatial_transcript_former/models/interaction.py
@@ -78,7 +78,7 @@ def __init__(
             n_heads (int): Number of attention heads.
             n_layers (int): Number of transformer/interaction layers.
             dropout (float): Dropout probability.
-            use_spatial_pe (bool): Incorporate relative gradients into attention.
+            use_spatial_pe (bool): Incorporate slide-stationary spatial coordinates into attention.
             interactions (list[str], optional): Which attention interactions to
                 enable.  Valid keys are ``p2p``, ``p2h``, ``h2p``, ``h2h``.
                 Defaults to all four (full self-attention).
@@ -120,6 +120,10 @@ def __init__(
         # 3. Learnable pathway tokens (one per pathway, shared across batch)
         self.pathway_tokens = nn.Parameter(torch.randn(1, num_pathways, token_dim))
 
+        # Learnable scale and bias for dot-product Softplus head (Option 1)
+        self.scale = nn.Parameter(torch.ones(1, 1, num_pathways))
+        self.bias = nn.Parameter(torch.zeros(1, 1, num_pathways))
+
         # 4. Spatial Positional Encoder (optional)
         self.spatial_encoder = None
         if use_spatial_pe:
@@ -207,7 +211,7 @@ def forward(
             x (torch.Tensor): Image data or pre-computed features.
                 - (B, 3, H, W): Single image patch.
                 - (B, S, D): Pre-computed features.
-            rel_coords (torch.Tensor, optional): Spatial relative coordinates.
+            rel_coords (torch.Tensor, optional): Slide-stationary spatial coordinates.
             mask (torch.Tensor, optional): Boolean padding mask for patches (B, S) where True = Padding.
             return_dense (bool): If True, returns per-patch pathway predictions instead of global predictions.
             return_attention (bool): If True, returns attention maps from all layers.
@@ -300,20 +304,35 @@ def forward(
 
         # 5. Compute pathway scores via cosine similarity
         # L2-normalize both sets of tokens to produce cosine similarities in [-1, 1]
-        norm_pathway = F.normalize(processed_pathway_tokens, dim=-1)  # (B, P, D)
+        # 5. Compute pathway scores via scaled dot product + softplus with learnable scale/bias
+        d_dim = processed_patch_tokens.shape[-1]
 
         if return_dense:
-            # Dense prediction: per-patch cosine similarity with pathway tokens
-            norm_patch = F.normalize(processed_patch_tokens, dim=-1)  # (B, S, D)
-            # (B, S, D) @ (B, D, P) -> (B, S, P)
-            pathway_scores = torch.matmul(norm_patch, norm_pathway.transpose(1, 2))
+            # Dense prediction: scaled dot-product between patch and pathway tokens
+            raw_scores = torch.matmul(
+                processed_patch_tokens, processed_pathway_tokens.transpose(1, 2)
+            ) / (
+                d_dim**0.5
+            )  # (B, S, P)
+            pathway_scores = F.softplus(self.scale * raw_scores + self.bias)
         else:
-            # Global prediction: pool patches first, then compute scores
-            global_patch_token = processed_patch_tokens.mean(
-                dim=1, keepdim=True
-            )  # (B, 1, D)
-            norm_global = F.normalize(global_patch_token, dim=-1)  # (B, 1, D)
-            pathway_scores = torch.matmul(norm_global, norm_pathway.transpose(1, 2))
+            # Global prediction: pool patches first (using mask if provided), then compute scores
+            if mask is not None:
+                valid_mask = (~mask).unsqueeze(-1).float()  # (B, S, 1)
+                global_patch_token = (processed_patch_tokens * valid_mask).sum(
+                    dim=1, keepdim=True
+                ) / valid_mask.sum(dim=1, keepdim=True).clamp(min=1.0)
+            else:
+                global_patch_token = processed_patch_tokens.mean(
+                    dim=1, keepdim=True
+                )  # (B, 1, D)
+
+            raw_scores = torch.matmul(
+                global_patch_token, processed_pathway_tokens.transpose(1, 2)
+            ) / (
+                d_dim**0.5
+            )  # (B, 1, P)
+            pathway_scores = F.softplus(self.scale * raw_scores + self.bias)
             pathway_scores = pathway_scores.squeeze(1)  # (B, P)
 
         results = [pathway_scores]
diff --git a/src/spatial_transcript_former/models/regression.py b/src/spatial_transcript_former/models/regression.py
index 696ff7f..892cbe5 100644
--- a/src/spatial_transcript_former/models/regression.py
+++ b/src/spatial_transcript_former/models/regression.py
@@ -54,3 +54,33 @@ def __init__(self, num_pathways, model_name="vit_b_16", pretrained=True):
 
     def forward(self, x):
         return self.backbone(x)
+
+
+class LinearProbe(nn.Module):
+    """Linear probe baseline on precomputed features."""
+
+    def __init__(self, input_dim=1024, num_pathways=50):
+        super().__init__()
+        self.classifier = nn.Linear(input_dim, num_pathways)
+
+    def forward(self, x):
+        # Supports input shape (B, L) or (B, N, L)
+        return self.classifier(x)
+
+
+class MLPProbe(nn.Module):
+    """MLP baseline on precomputed features."""
+
+    def __init__(self, input_dim=1024, num_pathways=50, hidden_dim=512, dropout=0.1):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, num_pathways),
+        )
+
+    def forward(self, x):
+        # Supports input shape (B, L) or (B, N, L)
+        return self.mlp(x)
diff --git a/src/spatial_transcript_former/predict.py b/src/spatial_transcript_former/predict.py
index 0cbf1db..5189444 100644
--- a/src/spatial_transcript_former/predict.py
+++ b/src/spatial_transcript_former/predict.py
@@ -369,6 +369,9 @@ def inject_predictions(
 # ═══════════════════════════════════════════════════════════════════════
 # Kept here for backwards compatibility with training scripts.
 
+import matplotlib
+
+matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
 
@@ -407,8 +410,18 @@ def plot_training_summary(
     if plot_pathways_list is not None and len(plot_pathways_list) > 0:
         target_pathways = [p.replace("HALLMARK_", "") for p in plot_pathways_list]
     else:
-        # Default to the first 6 pathways available if none specified
-        target_pathways = short_names[:6]
+        # Spatially and biologically prominent Hallmark pathways to monitor by default
+        default_interesting = [
+            "WNT_BETA_CATENIN_SIGNALING",
+            "E2F_TARGETS",
+            "MYC_TARGETS_V1",
+            "EPITHELIAL_MESENCHYMAL_TRANSITION",
+            "GLYCOLYSIS",
+            "OXIDATIVE_PHOSPHORYLATION",
+        ]
+        target_pathways = [p for p in default_interesting if p in short_names]
+        if not target_pathways:
+            target_pathways = short_names[:6]
 
     for target_pw in target_pathways:
         if target_pw in short_names:
diff --git a/src/spatial_transcript_former/recipes/hest/compute_pathway_activities.py b/src/spatial_transcript_former/recipes/hest/compute_pathway_activities.py
index e464784..4ff85f7 100644
--- a/src/spatial_transcript_former/recipes/hest/compute_pathway_activities.py
+++ b/src/spatial_transcript_former/recipes/hest/compute_pathway_activities.py
@@ -186,7 +186,15 @@ def _score_pathways(expr_matrix, gene_names, pathway_dict, min_genes=5):
     n_scored : int
         Number of pathways that met the min_genes threshold.
     """
-    gene_to_idx = {g: i for i, g in enumerate(gene_names)}
+
+    def clean_gene_name(g):
+        g_upper = g.upper()
+        for prefix in ["GRCH38_", "MM10_", "GRCM38_", "HS_", "MOUSE_"]:
+            if g_upper.startswith(prefix):
+                return g_upper[len(prefix) :]
+        return g_upper
+
+    gene_to_idx = {clean_gene_name(g): i for i, g in enumerate(gene_names)}
     n_spots = expr_matrix.shape[0]
 
     all_pathways = list(pathway_dict.keys())
@@ -194,7 +202,11 @@ def _score_pathways(expr_matrix, gene_names, pathway_dict, min_genes=5):
     n_scored = 0
 
     for i, (pw_name, pw_genes) in enumerate(pathway_dict.items()):
-        col_indices = [gene_to_idx[g] for g in pw_genes if g in gene_to_idx]
+        col_indices = [
+            gene_to_idx[clean_gene_name(g)]
+            for g in pw_genes
+            if clean_gene_name(g) in gene_to_idx
+        ]
         if len(col_indices) < min_genes:
             continue
         activities[:, i] = expr_matrix[:, col_indices].mean(axis=1)
diff --git a/src/spatial_transcript_former/recipes/hest/dataset.py b/src/spatial_transcript_former/recipes/hest/dataset.py
index 7a03faa..10f7851 100644
--- a/src/spatial_transcript_former/recipes/hest/dataset.py
+++ b/src/spatial_transcript_former/recipes/hest/dataset.py
@@ -40,6 +40,17 @@
 from scipy.spatial import KDTree
 import torch.nn.functional as F
 
+
+def worker_init_fn(worker_id):
+    """Set random seeds for dataloader worker processes to prevent duplicate data streams."""
+    worker_seed = torch.initial_seed() % 2**32
+    import numpy as np
+    import random
+
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+
+
 # Augmentation helpers and normalize_coordinates are now in data.base
 # and imported above. Kept here for backward compatibility:
 # from spatial_transcript_former.recipes.hest.dataset import apply_dihedral_augmentation
@@ -162,13 +173,21 @@ def __getitem__(self, idx):
             data = patch  # (3, H, W)
             rel_coords = torch.zeros((1, 2))
 
+        # mask construction for raw patches neighborhood
+        if self.neighborhood_indices is not None:
+            mask = torch.zeros(data.shape[0], dtype=torch.bool)
+            if len(self.coords) < data.shape[0]:
+                mask[len(self.coords) :] = True
+        else:
+            mask = torch.zeros(1, dtype=torch.bool)
+
         # gene_counts removed (pathway-only)
         pathway_act = (
             self.pathway_activities[idx]
             if self.pathway_activities is not None
             else None
         )
-        return data, None, pathway_act, rel_coords
+        return data, None, pathway_act, rel_coords, mask
 
     def __del__(self):
         if self.h5_file:
@@ -334,13 +353,15 @@ def get_hest_dataloader(
                             if name in pw_names:
                                 p_indices.append(pw_names.index(name))
                             else:
-                                p_indices.append(-1)
+                                raise ValueError(
+                                    f"Pathway '{name}' not found in pathway target file {pw_h5_path}. "
+                                    f"Available pathway names: {pw_names}"
+                                )
 
                         p = len(pathway_names)
                         subset_acts = np.zeros((acts.shape[0], p), dtype=np.float32)
                         for i, idx in enumerate(p_indices):
-                            if idx != -1:
-                                subset_acts[:, i] = acts[:, idx]
+                            subset_acts[:, i] = acts[:, idx]
 
                         pathway_activities = torch.tensor(
                             subset_acts[mask_bool], dtype=torch.float32
@@ -378,6 +399,7 @@ def get_hest_dataloader(
         collate_fn=collate_fn_patch,
         pin_memory=True,
         persistent_workers=(num_workers > 0),
+        worker_init_fn=worker_init_fn,
     )
 
 
@@ -487,14 +509,15 @@ def _load_data(self):
                         if name in pw_names:
                             indices.append(pw_names.index(name))
                         else:
-                            # If a requested pathway is missing, we'll use a zero column
-                            indices.append(-1)
+                            raise ValueError(
+                                f"Pathway '{name}' not found in pathway target file {h5_path}. "
+                                f"Available pathway names: {pw_names}"
+                            )
 
                     p = len(self.target_pathway_names)
                     subset_acts = np.zeros((acts.shape[0], p), dtype=np.float32)
                     for i, idx in enumerate(indices):
-                        if idx != -1:
-                            subset_acts[:, i] = acts[:, idx]
+                        subset_acts[:, i] = acts[:, idx]
 
                     self.pathway_activities = torch.tensor(
                         subset_acts[mask_bool], dtype=torch.float32
@@ -564,7 +587,12 @@ def _getitem_patch(self, idx):
         else:
             combined_idxs = neighbor_idxs
 
-        feats = self.features[combined_idxs]
+        feats = self.features[combined_idxs].clone()
+
+        # --- Mask construction for padding and dropped neighbours ---
+        mask = torch.zeros(len(combined_idxs), dtype=torch.bool)
+        if len(self.coords) < self.n_neighbors + 1:
+            mask[len(self.coords) : self.n_neighbors + 1] = True
 
         # --- Neighbourhood dropout augmentation ---
         # Randomly zero out 0–2 neighbour feature vectors (never the centre).
@@ -575,10 +603,11 @@ def _getitem_patch(self, idx):
                     range(1, self.n_neighbors + 1), size=n_to_drop, replace=False
                 )
                 feats[drop_idxs] = 0.0
+                mask[drop_idxs] = True
 
-        # --- Relative coordinates ---
-        center_coord = self.coords[idx]
-        rel_coords = self.coords[combined_idxs] - center_coord  # (S, 2)
+        # --- Slide-stationary coordinates ---
+        # Aligned with whole-slide mode coordinates to ensure consistency
+        rel_coords = self.coords[combined_idxs].clone()
 
         if self.augment:
             # Dihedral rotation / flip
@@ -590,7 +619,7 @@ def _getitem_patch(self, idx):
             if self.pathway_activities is not None
             else None
         )
-        return feats, target_genes, pathway_acts, rel_coords
+        return feats, target_genes, pathway_acts, rel_coords, mask
 
 
 # ---------------------------------------------------------------------------
@@ -599,16 +628,16 @@ def _getitem_patch(self, idx):
 
 
 def collate_fn_patch(batch):
-    """Collate ``(feats, genes, pathway_acts, coords)`` tuples.
+    """Collate ``(feats, genes, pathway_acts, coords, mask)`` tuples.
 
     Handles ``pathway_acts=None`` (when no pathway targets dir is configured)
     by passing ``None`` through.
 
     Args:
-        batch: List of ``(feats, genes, pathway_acts, coords)`` tuples.
+        batch: List of ``(feats, genes, pathway_acts, coords, mask)`` tuples.
 
     Returns:
-        tuple: ``(feats, genes, pathway_acts, coords)`` where ``pathway_acts``
+        tuple: ``(feats, genes, pathway_acts, coords, mask)`` where ``pathway_acts``
         is a stacked tensor or ``None``.
     """
     feats = torch.stack([item[0] for item in batch])
@@ -618,7 +647,9 @@ def collate_fn_patch(batch):
     has_pathways = batch[0][2] is not None
     pathways = torch.stack([item[2] for item in batch]) if has_pathways else None
     coords = torch.stack([item[3] for item in batch])
-    return feats, genes, pathways, coords
+    has_mask = len(batch[0]) > 4 and batch[0][4] is not None
+    mask = torch.stack([item[4] for item in batch]) if has_mask else None
+    return feats, genes, pathways, coords, mask
 
 
 # ---------------------------------------------------------------------------
@@ -768,6 +799,7 @@ def collate_fn_ws(batch):
             collate_fn=collate_fn_ws,
             pin_memory=True,
             persistent_workers=(num_workers > 0),
+            worker_init_fn=worker_init_fn,
         )
 
     return DataLoader(
@@ -778,4 +810,5 @@ def collate_fn_ws(batch):
         collate_fn=collate_fn_patch,
         pin_memory=True,
         persistent_workers=(num_workers > 0),
+        worker_init_fn=worker_init_fn,
     )
diff --git a/src/spatial_transcript_former/train.py b/src/spatial_transcript_former/train.py
index e8a9857..d07bef6 100644
--- a/src/spatial_transcript_former/train.py
+++ b/src/spatial_transcript_former/train.py
@@ -132,7 +132,7 @@ def main():
     print(
         f"Targets: pathway_format_version={PATHWAY_FILE_VERSION} "
         "(mean log1p CP10k of pathway members). "
-        "Validation MAE/loss are in those units; best-model selection uses CCC."
+        "Validation MAE/loss are in those units; best-model selection uses CCC (concordance)."
     )
 
     # 3. Output & Logger
@@ -142,8 +142,8 @@ def main():
 
     # 4. Resume
     # ``best_val_metric`` tracks the highest val_ccc seen so far (CCC is
-    # higher-is-better and pathway-scale-invariant; preferable to MSE-based
-    # selection now that targets live in raw log1p CP10k units).
+    # higher-is-better and offset-sensitive (concordance-measuring); preferable to
+    # MSE-based selection now that targets live in raw log1p CP10k units).
     start_epoch, best_val_metric = 0, -float("inf")
     schedulers = {"main": main_scheduler}
     if args.resume:
@@ -159,6 +159,7 @@ def main():
             for _ in range(start_epoch):
                 main_scheduler.step()
 
+    epochs_no_improve = 0
     # 5. Training Loop
     for epoch in range(start_epoch, args.epochs):
         print(f"\nEpoch {epoch + 1}/{args.epochs}")
@@ -230,6 +231,7 @@ def main():
         val_ccc = val_metrics.get("val_ccc")
         if val_ccc is not None and val_ccc > best_val_metric:
             best_val_metric = val_ccc
+            epochs_no_improve = 0
 
             # Legacy state_dict path (kept for tools that still load .pth directly)
             best_path = os.path.join(args.output_dir, f"best_model_{args.model}.pth")
@@ -248,6 +250,9 @@ def main():
             except Exception as e:
                 print(f"  (skipped save_pretrained bundle: {e})")
             print(f"Saved best model (val_ccc={val_ccc:.4f}) -> {best_path}")
+        else:
+            if val_ccc is not None:
+                epochs_no_improve += 1
 
         # Save latest
         save_checkpoint(
@@ -261,11 +266,21 @@ def main():
             args.model,
         )
 
+        if (
+            args.early_stopping_patience is not None
+            and epochs_no_improve >= args.early_stopping_patience
+        ):
+            print(
+                f"Early stopping triggered: val_ccc has not improved for {args.early_stopping_patience} epochs."
+            )
+            break
+
         # Periodic visualization
         if val_ids and (epoch + 1) % args.vis_interval == 0:
-            vis_id = args.vis_sample if args.vis_sample else val_ids[0]
-            print(f"Generating visualization for sample {vis_id}...")
-            run_inference_plot(model, args, vis_id, epoch + 1, device)
+            if not getattr(model, "weak_supervision", False):
+                vis_id = args.vis_sample if args.vis_sample else val_ids[0]
+                print(f"Generating visualization for sample {vis_id}...")
+                run_inference_plot(model, args, vis_id, epoch + 1, device)
 
     # 6. Finalize
     logger.finalize(best_val_metric)
diff --git a/src/spatial_transcript_former/training/arguments.py b/src/spatial_transcript_former/training/arguments.py
index e2ff8a4..409ad57 100644
--- a/src/spatial_transcript_former/training/arguments.py
+++ b/src/spatial_transcript_former/training/arguments.py
@@ -131,6 +131,12 @@ def parse_args():
     )
     g.add_argument("--compile", action="store_true")
     g.add_argument("--resume", action="store_true")
+    g.add_argument(
+        "--early-stopping-patience",
+        type=int,
+        default=None,
+        help="Number of epochs to wait for val_ccc improvement before early stopping (default: None/disabled)",
+    )
     g.add_argument(
         "--run-name",
         type=str,
diff --git a/src/spatial_transcript_former/training/builder.py b/src/spatial_transcript_former/training/builder.py
index 8329a43..d997b39 100644
--- a/src/spatial_transcript_former/training/builder.py
+++ b/src/spatial_transcript_former/training/builder.py
@@ -1,7 +1,13 @@
 import os
 import torch
 import torch.nn as nn
-from spatial_transcript_former.models import HE2RNA, ViT_ST, SpatialTranscriptFormer
+from spatial_transcript_former.models import (
+    HE2RNA,
+    ViT_ST,
+    SpatialTranscriptFormer,
+    LinearProbe,
+    MLPProbe,
+)
 from spatial_transcript_former.training.losses import (
     CCCLoss,
     CLIPAlignmentLoss,
@@ -26,17 +32,35 @@ def setup_model(args, device):
     args.num_pathways = _resolve_num_pathways(args)
 
     if args.model == "he2rna":
-        model = HE2RNA(
-            num_pathways=args.num_pathways,
-            backbone=args.backbone,
-            pretrained=args.pretrained,
-        )
+        if getattr(args, "precomputed", False):
+            from spatial_transcript_former.models.backbones import get_backbone
+
+            _, feature_dim = get_backbone(args.backbone, pretrained=False)
+            model = LinearProbe(
+                input_dim=feature_dim,
+                num_pathways=args.num_pathways,
+            )
+        else:
+            model = HE2RNA(
+                num_pathways=args.num_pathways,
+                backbone=args.backbone,
+                pretrained=args.pretrained,
+            )
     elif args.model == "vit_st":
-        model = ViT_ST(
-            num_pathways=args.num_pathways,
-            model_name=args.backbone if "vit_" in args.backbone else "vit_b_16",
-            pretrained=args.pretrained,
-        )
+        if getattr(args, "precomputed", False):
+            from spatial_transcript_former.models.backbones import get_backbone
+
+            _, feature_dim = get_backbone(args.backbone, pretrained=False)
+            model = MLPProbe(
+                input_dim=feature_dim,
+                num_pathways=args.num_pathways,
+            )
+        else:
+            model = ViT_ST(
+                num_pathways=args.num_pathways,
+                model_name=args.backbone if "vit_" in args.backbone else "vit_b_16",
+                pretrained=args.pretrained,
+            )
     elif args.model == "interaction":
         print(
             f"Initializing SpatialTranscriptFormer ({args.backbone}, pretrained={args.pretrained}, num_pathways={args.num_pathways})"
diff --git a/src/spatial_transcript_former/training/engine.py b/src/spatial_transcript_former/training/engine.py
index c2221f3..98a458a 100644
--- a/src/spatial_transcript_former/training/engine.py
+++ b/src/spatial_transcript_former/training/engine.py
@@ -14,11 +14,17 @@
 from spatial_transcript_former.data.spatial_stats import spatial_coherence_score
 
 
-def _criterion_call(criterion, preds, targets, mask=None):
-    """Call the criterion, passing the mask only when provided."""
+def _criterion_call(criterion, preds, targets, mask=None, pathway_mask=None):
+    """Call the criterion, passing the mask and pathway_mask only when provided."""
+    kwargs = {}
     if mask is not None:
-        return criterion(preds, targets, mask=mask)
-    return criterion(preds, targets)
+        kwargs["mask"] = mask
+    if pathway_mask is not None:
+        kwargs["pathway_mask"] = pathway_mask
+    try:
+        return criterion(preds, targets, **kwargs)
+    except TypeError:
+        return criterion(preds, targets)
 
 
 def _optimizer_step(
@@ -86,21 +92,34 @@ def train_one_epoch(
                 )
             pathway_targets = pathway_targets.to(device, non_blocking=True)
 
+            # Calculate pathway validity mask dynamically (std > 1e-6 over valid spots)
+            B, N, G = pathway_targets.shape
+            pathway_mask = torch.zeros(B, G, dtype=torch.bool, device=device)
+            for b_idx in range(B):
+                valid_idx = ~mask[b_idx]
+                t_slide = pathway_targets[b_idx, valid_idx]
+                if t_slide.shape[0] >= 2:
+                    pathway_mask[b_idx] = torch.std(t_slide, dim=0) > 1e-6
+                else:
+                    pathway_mask[b_idx] = True
+
             with torch.amp.autocast("cuda", enabled=scaler is not None):
-                if isinstance(model, SpatialTranscriptFormer) and not getattr(
-                    model, "weak_supervision", False
-                ):
-                    preds = model(
-                        feats,
-                        return_dense=True,
-                        mask=mask,
-                        rel_coords=coords,
-                    )
+                if not getattr(model, "weak_supervision", False):
+                    if isinstance(model, SpatialTranscriptFormer):
+                        preds = model(
+                            feats,
+                            return_dense=True,
+                            mask=mask,
+                            rel_coords=coords,
+                        )
+                    else:
+                        preds = model(feats)
                     loss = _criterion_call(
                         criterion,
                         preds,
                         pathway_targets,
                         mask=mask,
+                        pathway_mask=pathway_mask,
                     )
                 else:
                     preds = model(feats)
@@ -122,19 +141,29 @@ def train_one_epoch(
             pbar.set_postfix({"loss": f"{current_loss:.4f}"})
     else:
         for batch_idx, batch in enumerate(pbar):
-            # Unpack: (images, None, pathway_targets, rel_coords)
-            images, _, pathway_targets, rel_coords = batch
+            # Unpack: (images, None, pathway_targets, rel_coords, mask)
+            images, _, pathway_targets, rel_coords, mask = batch
             images = images.to(device, non_blocking=True)
             rel_coords = rel_coords.to(device, non_blocking=True)
+            if mask is not None:
+                mask = mask.to(device, non_blocking=True)
             if pathway_targets is None:
                 raise ValueError(
                     "pathway_targets is None, but training now requires pathway targets."
                 )
             pathway_targets = pathway_targets.to(device, non_blocking=True)
 
+            # Calculate pathway validity mask dynamically (std > 1e-6 over batch)
+            if pathway_targets.shape[0] >= 2:
+                pathway_mask = torch.std(pathway_targets, dim=0) > 1e-6
+            else:
+                pathway_mask = torch.ones(
+                    pathway_targets.shape[1], dtype=torch.bool, device=device
+                )
+
             with torch.amp.autocast("cuda", enabled=scaler is not None):
                 if isinstance(model, SpatialTranscriptFormer):
-                    outputs = model(images, rel_coords=rel_coords)
+                    outputs = model(images, rel_coords=rel_coords, mask=mask)
                 else:
                     outputs = model(images)
 
@@ -142,6 +171,8 @@ def train_one_epoch(
                     criterion,
                     outputs,
                     pathway_targets,
+                    mask=mask,
+                    pathway_mask=pathway_mask,
                 )
                 loss = loss / grad_accum_steps
 
@@ -181,6 +212,8 @@ def validate(model, loader, criterion, device, whole_slide=False, use_amp=False)
     pred_var_list = []
     attn_correlations = []
     spatial_coherence_list = []
+    all_mil_preds = []
+    all_mil_targets = []
 
     with torch.no_grad():
         for batch in tqdm(
@@ -195,28 +228,48 @@ def validate(model, loader, criterion, device, whole_slide=False, use_amp=False)
                 if pathway_targets is None:
                     raise ValueError("pathway_targets is None in validation.")
                 pathway_targets = pathway_targets.to(device, non_blocking=True)
+
+                B, N, G = pathway_targets.shape
+                pathway_mask = torch.zeros(B, G, dtype=torch.bool, device=device)
+                for b_idx in range(B):
+                    valid_idx = ~mask[b_idx]
+                    t_slide = pathway_targets[b_idx, valid_idx]
+                    if t_slide.shape[0] >= 2:
+                        pathway_mask[b_idx] = torch.std(t_slide, dim=0) > 1e-6
+                    else:
+                        pathway_mask[b_idx] = True
             else:
-                # Unpack: (images, None, pathway_targets, rel_coords)
-                images, _, pathway_targets, rel_coords = batch
+                # Unpack: (images, None, pathway_targets, rel_coords, mask)
+                images, _, pathway_targets, rel_coords, mask = batch
                 images = images.to(device, non_blocking=True)
                 rel_coords = rel_coords.to(device, non_blocking=True)
+                if mask is not None:
+                    mask = mask.to(device, non_blocking=True)
                 if pathway_targets is None:
                     raise ValueError("pathway_targets is None in validation.")
                 pathway_targets = pathway_targets.to(device, non_blocking=True)
 
+                if pathway_targets.shape[0] >= 2:
+                    pathway_mask = torch.std(pathway_targets, dim=0) > 1e-6
+                else:
+                    pathway_mask = torch.ones(
+                        pathway_targets.shape[1], dtype=torch.bool, device=device
+                    )
+
             with torch.amp.autocast("cuda", enabled=use_amp):
                 attn = None
 
                 if whole_slide:
-                    if isinstance(model, SpatialTranscriptFormer) and not getattr(
-                        model, "weak_supervision", False
-                    ):
-                        outputs = model(
-                            feats,
-                            return_dense=True,
-                            mask=mask,
-                            rel_coords=coords,
-                        )
+                    if not getattr(model, "weak_supervision", False):
+                        if isinstance(model, SpatialTranscriptFormer):
+                            outputs = model(
+                                feats,
+                                return_dense=True,
+                                mask=mask,
+                                rel_coords=coords,
+                            )
+                        else:
+                            outputs = model(feats)
                         targets = pathway_targets
                     else:
                         # MIL models: extract attention if supported
@@ -231,35 +284,24 @@ def validate(model, loader, criterion, device, whole_slide=False, use_amp=False)
                 else:
                     targets = pathway_targets
                     if isinstance(model, SpatialTranscriptFormer):
-                        outputs = model(images, rel_coords=rel_coords)
+                        outputs = model(images, rel_coords=rel_coords, mask=mask)
                     else:
                         outputs = model(images)
 
                 # Compute loss
-                if (
-                    whole_slide
-                    and isinstance(model, SpatialTranscriptFormer)
-                    and not getattr(model, "weak_supervision", False)
-                ):
-                    loss = _criterion_call(
-                        criterion,
-                        outputs,
-                        targets,
-                        mask=mask,
-                    )
-                else:
-                    loss = _criterion_call(
-                        criterion,
-                        outputs,
-                        targets,
-                    )
+                loss = _criterion_call(
+                    criterion,
+                    outputs,
+                    targets,
+                    mask=mask,
+                    pathway_mask=pathway_mask,
+                )
 
                 # --- Interpretability Metrics (MAE, PCC, CCC) ---
                 eval_preds = outputs
                 mae_diff = torch.abs(eval_preds - targets)
                 if (
                     whole_slide
-                    and isinstance(model, SpatialTranscriptFormer)
                     and not getattr(model, "weak_supervision", False)
                     and mask is not None
                 ):
@@ -287,68 +329,73 @@ def validate(model, loader, criterion, device, whole_slide=False, use_amp=False)
                 # PCC and CCC — computed for all model/mode combinations
                 if torch.isfinite(eval_preds).all() and torch.isfinite(targets).all():
                     if whole_slide:
-                        for b_idx in range(eval_preds.shape[0]):
-                            p_slide = eval_preds[b_idx]  # (N, G)
-                            t_slide = targets[b_idx]  # (N, G)
-                            valid_idx = ~mask[b_idx]
-                            p_slide = p_slide[valid_idx]  # (V, G)
-                            t_slide = t_slide[valid_idx]  # (V, G)
-
-                            if p_slide.shape[0] >= 2:
-                                pred_mean = p_slide.mean(dim=0)  # (G,)
-                                tgt_mean = t_slide.mean(dim=0)  # (G,)
-                                vx = p_slide - pred_mean
-                                vy = t_slide - tgt_mean
-                                cov = (vx * vy).sum(dim=0)  # (G,)
-                                var_x = (vx**2).sum(dim=0)  # (G,)
-                                var_y = (vy**2).sum(dim=0)  # (G,)
-
-                                corr = cov / (
-                                    torch.sqrt(var_x + 1e-8) * torch.sqrt(var_y + 1e-8)
-                                )
-                                N_v = p_slide.shape[0]
-                                mean_diff_sq = (pred_mean - tgt_mean) ** 2
-                                ccc_vals = (2 * cov) / (
-                                    var_x + var_y + N_v * mean_diff_sq + 1e-8
-                                )
-
-                                valid_genes = torch.std(t_slide, dim=0) > 1e-6
-                                if valid_genes.any():
-                                    vc = corr[valid_genes]
-                                    vc = vc[torch.isfinite(vc)]
-                                    if len(vc) > 0:
-                                        pcc_list.append(vc.mean().item())
-                                    vk = ccc_vals[valid_genes]
-                                    vk = vk[torch.isfinite(vk)]
-                                    if len(vk) > 0:
-                                        ccc_list.append(vk.mean().item())
-
-                                # Per-pathway accumulation (slide-level entries)
-                                valid_idx_genes = torch.where(valid_genes)[0].tolist()
-                                for g in valid_idx_genes:
-                                    c = corr[g].item()
-                                    k = ccc_vals[g].item()
-                                    if np.isfinite(c):
-                                        per_pathway_pcc.setdefault(g, []).append(c)
-                                    if np.isfinite(k):
-                                        per_pathway_ccc.setdefault(g, []).append(k)
+                        if not getattr(model, "weak_supervision", False):
+                            for b_idx in range(eval_preds.shape[0]):
+                                p_slide = eval_preds[b_idx]  # (N, G)
+                                t_slide = targets[b_idx]  # (N, G)
+                                valid_idx = ~mask[b_idx]
+                                p_slide = p_slide[valid_idx]  # (V, G)
+                                t_slide = t_slide[valid_idx]  # (V, G)
+
+                                if p_slide.shape[0] >= 2:
+                                    pred_mean = p_slide.mean(dim=0)  # (G,)
+                                    tgt_mean = t_slide.mean(dim=0)  # (G,)
+                                    vx = p_slide - pred_mean
+                                    vy = t_slide - tgt_mean
+                                    N_v = p_slide.shape[0]
+                                    cov = (vx * vy).sum(dim=0) / N_v  # (G,)
+                                    var_x = (vx**2).sum(dim=0) / N_v  # (G,)
+                                    var_y = (vy**2).sum(dim=0) / N_v  # (G,)
+
+                                    corr = cov / (
+                                        torch.sqrt(var_x + 1e-8)
+                                        * torch.sqrt(var_y + 1e-8)
+                                    )
+                                    mean_diff_sq = (pred_mean - tgt_mean) ** 2
+                                    ccc_vals = (2 * cov) / (
+                                        var_x + var_y + mean_diff_sq + 1e-8
+                                    )
+
+                                    valid_genes = torch.std(t_slide, dim=0) > 1e-6
+                                    if valid_genes.any():
+                                        vc = corr[valid_genes]
+                                        vc = vc[torch.isfinite(vc)]
+                                        if len(vc) > 0:
+                                            pcc_list.append(vc.mean().item())
+                                        vk = ccc_vals[valid_genes]
+                                        vk = vk[torch.isfinite(vk)]
+                                        if len(vk) > 0:
+                                            ccc_list.append(vk.mean().item())
+
+                                    # Per-pathway accumulation (slide-level entries)
+                                    valid_idx_genes = torch.where(valid_genes)[
+                                        0
+                                    ].tolist()
+                                    for g in valid_idx_genes:
+                                        c = corr[g].item()
+                                        k = ccc_vals[g].item()
+                                        if np.isfinite(c):
+                                            per_pathway_pcc.setdefault(g, []).append(c)
+                                        if np.isfinite(k):
+                                            per_pathway_ccc.setdefault(g, []).append(k)
+                        else:
+                            all_mil_preds.append(eval_preds.detach().cpu())
+                            all_mil_targets.append(targets.detach().cpu())
                     else:
                         pred_mean = eval_preds.mean(dim=0)  # (G,)
                         tgt_mean = targets.mean(dim=0)  # (G,)
                         vx = eval_preds - pred_mean
                         vy = targets - tgt_mean
-                        cov = (vx * vy).sum(dim=0)
-                        var_x = (vx**2).sum(dim=0)
-                        var_y = (vy**2).sum(dim=0)
+                        B_size = eval_preds.shape[0]
+                        cov = (vx * vy).sum(dim=0) / B_size
+                        var_x = (vx**2).sum(dim=0) / B_size
+                        var_y = (vy**2).sum(dim=0) / B_size
 
                         corr = cov / (
                             torch.sqrt(var_x + 1e-8) * torch.sqrt(var_y + 1e-8)
                         )
-                        B_size = eval_preds.shape[0]
                         mean_diff_sq = (pred_mean - tgt_mean) ** 2
-                        ccc_vals = (2 * cov) / (
-                            var_x + var_y + B_size * mean_diff_sq + 1e-8
-                        )
+                        ccc_vals = (2 * cov) / (var_x + var_y + mean_diff_sq + 1e-8)
 
                         valid_genes = torch.std(targets, dim=0) > 1e-6
                         if valid_genes.any():
@@ -393,7 +440,6 @@ def validate(model, loader, criterion, device, whole_slide=False, use_amp=False)
                 if (
                     whole_slide
                     and mask is not None
-                    and isinstance(model, SpatialTranscriptFormer)
                     and not getattr(model, "weak_supervision", False)
                 ):
                     for b in range(eval_preds.shape[0]):
@@ -418,6 +464,45 @@ def validate(model, loader, criterion, device, whole_slide=False, use_amp=False)
                 else:
                     pred_var_list.append(eval_preds.var(dim=0).mean().item())
 
+    if all_mil_preds:
+        mil_preds = torch.cat(all_mil_preds, dim=0).to(device)  # (M, G)
+        mil_targets = torch.cat(all_mil_targets, dim=0).to(device)  # (M, G)
+
+        # Compute PCC and CCC across slides (cohort-level)
+        pred_mean = mil_preds.mean(dim=0)  # (G,)
+        tgt_mean = mil_targets.mean(dim=0)  # (G,)
+        vx = mil_preds - pred_mean
+        vy = mil_targets - tgt_mean
+        M_size = mil_preds.shape[0]
+        cov = (vx * vy).sum(dim=0) / M_size
+        var_x = (vx**2).sum(dim=0) / M_size
+        var_y = (vy**2).sum(dim=0) / M_size
+
+        corr = cov / (torch.sqrt(var_x + 1e-8) * torch.sqrt(var_y + 1e-8))
+        mean_diff_sq = (pred_mean - tgt_mean) ** 2
+        ccc_vals = (2 * cov) / (var_x + var_y + mean_diff_sq + 1e-8)
+
+        valid_genes = torch.std(mil_targets, dim=0) > 1e-6
+        if valid_genes.any():
+            vc = corr[valid_genes]
+            vc = vc[torch.isfinite(vc)]
+            if len(vc) > 0:
+                pcc_list.append(vc.mean().item())
+            vk = ccc_vals[valid_genes]
+            vk = vk[torch.isfinite(vk)]
+            if len(vk) > 0:
+                ccc_list.append(vk.mean().item())
+
+        # Also populate per-pathway dictionary for reporting
+        valid_idx_genes = torch.where(valid_genes)[0].tolist()
+        for g in valid_idx_genes:
+            c = corr[g].item()
+            k = ccc_vals[g].item()
+            if np.isfinite(c):
+                per_pathway_pcc.setdefault(g, []).append(c)
+            if np.isfinite(k):
+                per_pathway_ccc.setdefault(g, []).append(k)
+
     avg_loss = running_loss / len(loader)
     avg_mae = running_mae / len(loader)
     avg_baseline_mae = (
diff --git a/src/spatial_transcript_former/training/losses.py b/src/spatial_transcript_former/training/losses.py
index d95e4ae..2c72f87 100644
--- a/src/spatial_transcript_former/training/losses.py
+++ b/src/spatial_transcript_former/training/losses.py
@@ -11,6 +11,45 @@
 import torch.nn.functional as F
 
 
+def get_combined_mask(preds, mask, pathway_mask):
+    """Combine spatial padding mask and pathway validity mask.
+
+    Args:
+        preds: (B, N, G) or (B, G)
+        mask (spatial padding): (B, N) boolean where True = padding/ignore.
+        pathway_mask (pathway validity): (B, G) or (G,) boolean where True = valid.
+
+    Returns:
+        torch.Tensor (boolean): Same shape as preds, True = valid.
+    """
+    if preds.dim() == 3:
+        B, N, G = preds.shape
+        if mask is not None:
+            valid_spatial = ~mask.unsqueeze(-1)  # (B, N, 1)
+        else:
+            valid_spatial = torch.ones(B, N, 1, dtype=torch.bool, device=preds.device)
+
+        if pathway_mask is not None:
+            if pathway_mask.dim() == 1:
+                valid_pathway = pathway_mask.view(1, 1, G).expand(B, N, G)
+            else:
+                valid_pathway = pathway_mask.unsqueeze(1)  # (B, 1, G)
+        else:
+            valid_pathway = torch.ones(B, 1, G, dtype=torch.bool, device=preds.device)
+
+        return valid_spatial & valid_pathway
+    else:
+        B, G = preds.shape
+        if pathway_mask is not None:
+            if pathway_mask.dim() == 1:
+                valid_pathway = pathway_mask.view(1, G).expand(B, G)
+            else:
+                valid_pathway = pathway_mask
+        else:
+            valid_pathway = torch.ones(B, G, dtype=torch.bool, device=preds.device)
+        return valid_pathway
+
+
 class PCCLoss(nn.Module):
     """
     Pearson Correlation Coefficient Loss.
@@ -28,12 +67,13 @@ def __init__(self, eps=1e-8):
         super().__init__()
         self.eps = eps
 
-    def forward(self, preds, target, mask=None):
+    def forward(self, preds, target, mask=None, pathway_mask=None):
         """
         Args:
             preds:  (B, G) or (B, N, G)
             target: (B, G) or (B, N, G)
             mask:   (B, N) boolean, True = padded (ignore). Optional.
+            pathway_mask: (B, G) or (G,) boolean, True = valid. Optional.
 
         Returns:
             Scalar loss = 1 - mean(PCC).
@@ -53,7 +93,7 @@ def forward(self, preds, target, mask=None):
             preds = preds.squeeze(1)  # (B, G)
             target = target.squeeze(1)  # (B, G)
             if mask is not None:
-                valid = ~mask.squeeze(1)  # (B)
+                valid = ~mask.view(-1)  # (B)
                 preds = preds[valid]
                 target = target[valid]
 
@@ -66,6 +106,15 @@ def forward(self, preds, target, mask=None):
                 torch.sqrt(torch.sum(vx**2, dim=0) + self.eps)
                 * torch.sqrt(torch.sum(vy**2, dim=0) + self.eps)
             )
+            if pathway_mask is not None:
+                if pathway_mask.dim() == 2:
+                    p_mask = pathway_mask.any(dim=0)
+                else:
+                    p_mask = pathway_mask
+                cost = cost[p_mask]
+                if cost.numel() > 0:
+                    return 1 - cost.mean()
+                return torch.tensor(0.0, device=preds.device, requires_grad=True)
             return 1 - cost.mean()
 
         # 1. Masking: Zero out padded positions so they don't contribute to sums
@@ -101,6 +150,13 @@ def forward(self, preds, target, mask=None):
         )  # (B, G)
 
         # Average across genes, then average across batch
+        if pathway_mask is not None:
+            if pathway_mask.dim() == 1:
+                pathway_mask = pathway_mask.unsqueeze(0).expand(B, -1)
+            valid_cost = cost[pathway_mask]
+            if valid_cost.numel() > 0:
+                return 1 - valid_cost.mean()
+            return torch.tensor(0.0, device=preds.device, requires_grad=True)
         return 1 - cost.mean()
 
 
@@ -115,12 +171,13 @@ class CCCLoss(PCCLoss):
     but systematically offset (wrong mean or variance) will have CCC < PCC.
     """
 
-    def forward(self, preds, target, mask=None):
+    def forward(self, preds, target, mask=None, pathway_mask=None):
         """
         Args:
             preds:  (B, G) or (B, N, G)
             target: (B, G) or (B, N, G)
             mask:   (B, N) boolean, True = padded (ignore). Optional.
+            pathway_mask: (B, G) or (G,) boolean, True = valid. Optional.
 
         Returns:
             Scalar loss = 1 - mean(CCC).
@@ -137,22 +194,33 @@ def forward(self, preds, target, mask=None):
             preds = preds.squeeze(1)
             target = target.squeeze(1)
             if mask is not None:
-                valid = ~mask.squeeze(1)
+                valid = ~mask.view(-1)
                 preds = preds[valid]
                 target = target[valid]
 
-            if preds.shape[0] < 2:
+            B_size = preds.shape[0]
+            if B_size < 2:
                 return torch.tensor(0.0, device=preds.device, requires_grad=True)
 
             pred_mean = preds.mean(dim=0)
             target_mean = target.mean(dim=0)
             vx = preds - pred_mean
             vy = target - target_mean
-            cov = (vx * vy).sum(dim=0)
-            var_x = (vx**2).sum(dim=0)
-            var_y = (vy**2).sum(dim=0)
+            cov = (vx * vy).sum(dim=0) / B_size
+            var_x = (vx**2).sum(dim=0) / B_size
+            var_y = (vy**2).sum(dim=0) / B_size
             mean_diff_sq = (pred_mean - target_mean) ** 2
             ccc = (2 * cov) / (var_x + var_y + mean_diff_sq + self.eps)
+
+            if pathway_mask is not None:
+                if pathway_mask.dim() == 2:
+                    p_mask = pathway_mask.any(dim=0)
+                else:
+                    p_mask = pathway_mask
+                ccc = ccc[p_mask]
+                if ccc.numel() > 0:
+                    return 1 - ccc.mean()
+                return torch.tensor(0.0, device=preds.device, requires_grad=True)
             return 1 - ccc.mean()
 
         if mask is not None:
@@ -173,45 +241,53 @@ def forward(self, preds, target, mask=None):
             vx = vx * valid.float()
             vy = vy * valid.float()
 
-        cov = (vx * vy).sum(dim=1)  # (B, G)
-        var_x = (vx**2).sum(dim=1)  # (B, G)
-        var_y = (vy**2).sum(dim=1)  # (B, G)
+        # Mean-based covariance and variance calculations (stable under varying N)
+        counts = valid_counts.squeeze(-1)
+        cov = (vx * vy).sum(dim=1) / counts  # (B, G)
+        var_x = (vx**2).sum(dim=1) / counts  # (B, G)
+        var_y = (vy**2).sum(dim=1) / counts  # (B, G)
 
         mean_diff_sq = (pred_means.squeeze(1) - target_means.squeeze(1)) ** 2  # (B, G)
         ccc = (2 * cov) / (var_x + var_y + mean_diff_sq + self.eps)  # (B, G)
 
+        if pathway_mask is not None:
+            if pathway_mask.dim() == 1:
+                pathway_mask = pathway_mask.unsqueeze(0).expand(B, -1)
+            valid_ccc = ccc[pathway_mask]
+            if valid_ccc.numel() > 0:
+                return 1 - valid_ccc.mean()
+            return torch.tensor(0.0, device=preds.device, requires_grad=True)
         return 1 - ccc.mean()
 
 
 class MaskedMSELoss(nn.Module):
     """
-    MSE loss with optional masking for padded positions.
+    MSE loss with optional masking for padded positions and invalid pathways.
 
     When no mask is provided, behaves identically to nn.MSELoss().
     """
 
-    def forward(self, preds, target, mask=None):
+    def forward(self, preds, target, mask=None, pathway_mask=None):
         """
         Args:
             preds:  (B, G) or (B, N, G)
             target: (B, G) or (B, N, G)
             mask:   (B, N) boolean, True = padded (ignore). Optional.
+            pathway_mask: (B, G) or (G,) boolean, True = valid. Optional.
 
         Returns:
             Scalar MSE loss over valid positions.
         """
         diff = (preds - target) ** 2
-
-        if mask is not None and preds.dim() == 3:
-            valid = ~mask.unsqueeze(-1).expand_as(diff)
+        valid = get_combined_mask(preds, mask, pathway_mask)
+        if valid.any():
             return (diff * valid.float()).sum() / valid.sum()
-
-        return diff.mean()
+        return torch.tensor(0.0, device=preds.device, requires_grad=True)
 
 
 class MaskedHuberLoss(nn.Module):
     """
-    Huber (smooth L1) loss with optional masking for padded positions.
+    Huber (smooth L1) loss with optional masking for padded positions and invalid pathways.
 
     More robust to outlier pathway activity values than MSE: quadratic near
     zero, linear beyond delta.
@@ -224,23 +300,22 @@ def __init__(self, delta=1.0):
         super().__init__()
         self.delta = delta
 
-    def forward(self, preds, target, mask=None):
+    def forward(self, preds, target, mask=None, pathway_mask=None):
         """
         Args:
             preds:  (B, G) or (B, N, G)
             target: (B, G) or (B, N, G)
             mask:   (B, N) boolean, True = padded (ignore). Optional.
+            pathway_mask: (B, G) or (G,) boolean, True = valid. Optional.
 
         Returns:
             Scalar Huber loss over valid positions.
         """
         diff = F.huber_loss(preds, target, reduction="none", delta=self.delta)
-
-        if mask is not None and preds.dim() == 3:
-            valid = ~mask.unsqueeze(-1).expand_as(diff)
+        valid = get_combined_mask(preds, mask, pathway_mask)
+        if valid.any():
             return (diff * valid.float()).sum() / valid.sum()
-
-        return diff.mean()
+        return torch.tensor(0.0, device=preds.device, requires_grad=True)
 
 
 class CLIPAlignmentLoss(nn.Module):
@@ -339,18 +414,19 @@ def __init__(
         self.clip = CLIPAlignmentLoss(clip_temperature) if clip_weight > 0 else None
         self.clip_weight = clip_weight
 
-    def forward(self, preds, target, mask=None):
+    def forward(self, preds, target, mask=None, pathway_mask=None):
         """
         Args:
             preds:  (B, G) or (B, N, G)
             target: (B, G) or (B, N, G)
             mask:   (B, N) boolean, True = padded (ignore). Optional.
+            pathway_mask: (B, G) or (G,) boolean, True = valid. Optional.
 
         Returns:
             Scalar composite loss.
         """
-        loss = self.mse(preds, target, mask)
-        loss = loss + self.alpha * self.pcc(preds, target, mask)
+        loss = self.mse(preds, target, mask, pathway_mask)
+        loss = loss + self.alpha * self.pcc(preds, target, mask, pathway_mask)
         if self.clip is not None:
             loss = loss + self.clip_weight * self.clip(preds, target, mask)
         return loss
diff --git a/src/spatial_transcript_former/visualization.py b/src/spatial_transcript_former/visualization.py
index 803cc68..eee57b8 100644
--- a/src/spatial_transcript_former/visualization.py
+++ b/src/spatial_transcript_former/visualization.py
@@ -83,7 +83,7 @@ def run_inference_plot(model, args, sample_id, epoch, device):
     from spatial_transcript_former.predict import plot_training_summary
 
     # 1. Setup Data
-    _, val_loader, _ = setup_dataloaders(args, [], [sample_id])
+    _, val_loader, val_whole_slide = setup_dataloaders(args, [], [sample_id])
     if val_loader is None:
         return
 
@@ -96,49 +96,59 @@ def run_inference_plot(model, args, sample_id, epoch, device):
     # 2. Run Inference
     with torch.no_grad():
         for batch in val_loader:
-            if args.whole_slide:
+            if val_whole_slide:
                 image_features, _, target, coords, mask = batch
                 image_features = image_features.to(device)
                 coords = coords.to(device)
                 mask = mask.to(device)
                 target = target.to(device)
+                mask_to_store = mask
             else:
-                image_features, _, target, coords = batch
+                image_features, _, target, coords, mask = batch
                 image_features = image_features.to(device)
                 coords = coords.to(device)
-                mask = torch.ones(target.shape[0], device=device)
                 target = target.to(device)
+                if mask is not None:
+                    mask = mask.to(device)
+                    mask_to_store = mask[:, 0]
+                else:
+                    mask_to_store = torch.zeros(
+                        target.shape[0], dtype=torch.bool, device=device
+                    )
 
             # Forward pass
-            if args.whole_slide:
-                outputs = model(
-                    image_features,
-                    rel_coords=coords,
-                    mask=mask,
-                    return_dense=True,
-                )
+            if val_whole_slide:
+                if isinstance(model, SpatialTranscriptFormer):
+                    outputs = model(
+                        image_features,
+                        rel_coords=coords,
+                        mask=mask,
+                        return_dense=True,
+                    )
+                else:
+                    outputs = model(image_features)
             else:
                 if isinstance(model, SpatialTranscriptFormer):
-                    outputs = model(image_features, rel_coords=coords)
+                    outputs = model(image_features, rel_coords=coords, mask=mask)
                 else:
                     outputs = model(image_features)
 
             preds_list.append(outputs.cpu())
             targets_list.append(target.cpu())
             coords_list.append(coords.cpu())
-            masks_list.append(mask.cpu())
+            masks_list.append(mask_to_store.cpu())
 
-            if args.whole_slide:
+            if val_whole_slide:
                 break  # Whole slide is one batch
 
     # Concatenate results (for patch-based)
-    all_preds = torch.cat(preds_list, dim=1 if args.whole_slide else 0)
-    all_targets = torch.cat(targets_list, dim=1 if args.whole_slide else 0)
-    all_coords = torch.cat(coords_list, dim=1 if args.whole_slide else 0)
-    all_masks = torch.cat(masks_list, dim=1 if args.whole_slide else 0)
+    all_preds = torch.cat(preds_list, dim=1 if val_whole_slide else 0)
+    all_targets = torch.cat(targets_list, dim=1 if val_whole_slide else 0)
+    all_coords = torch.cat(coords_list, dim=1 if val_whole_slide else 0)
+    all_masks = torch.cat(masks_list, dim=1 if val_whole_slide else 0)
 
     # Squeeze batch dim for processing
-    if args.whole_slide:
+    if val_whole_slide:
         pathway_preds = all_preds.numpy()[0]
         pathway_truth = all_targets.numpy()[0]
         coords = all_coords.numpy()[0]
@@ -148,31 +158,44 @@ def run_inference_plot(model, args, sample_id, epoch, device):
         pathway_preds = all_preds.numpy()
         pathway_truth = all_targets.numpy()
         if all_coords.ndim == 3:
-            coords = all_coords.squeeze(1).numpy()
+            coords = all_coords[:, 0].numpy()
         else:
             coords = all_coords.numpy()
         mask = all_masks.numpy()
-        valid_idx = mask.astype(bool)
+        valid_idx = ~mask.astype(bool)
 
     coords = coords[valid_idx]
     pathway_preds = pathway_preds[valid_idx]
     pathway_truth = pathway_truth[valid_idx]
 
     # The dataloader applies normalize_coordinates which breaks alignment with the full-res histology
-    # image. Let's fetch the raw coordinates directly from the .pt file.
+    # image. Let's fetch the raw coordinates directly from the .pt or .h5 file.
     try:
         from spatial_transcript_former.data.paths import resolve_feature_dir
         from spatial_transcript_former.recipes.hest.dataset import get_h5ad_valid_mask
 
-        feat_dir = resolve_feature_dir(
-            args.data_dir,
-            getattr(args, "backbone", "resnet50"),
-            getattr(args, "feature_dir", None),
-        )
-        pt_path = os.path.join(feat_dir, f"{sample_id}.pt")
-        saved_data = torch.load(pt_path, map_location="cpu", weights_only=True)
-        raw_coords = saved_data["coords"].numpy()
-        barcodes = saved_data["barcodes"]
+        # Try to load from precomputed features .pt file
+        try:
+            feat_dir = resolve_feature_dir(
+                args.data_dir,
+                getattr(args, "backbone", "resnet50"),
+                getattr(args, "feature_dir", None),
+            )
+            pt_path = os.path.join(feat_dir, f"{sample_id}.pt")
+            saved_data = torch.load(pt_path, map_location="cpu", weights_only=True)
+            raw_coords = saved_data["coords"].numpy()
+            barcodes = saved_data["barcodes"]
+        except Exception:
+            # Fall back to raw patches .h5 file if precomputed is not available
+            h5_path = os.path.join(args.data_dir, "patches", f"{sample_id}.h5")
+            if os.path.exists(h5_path):
+                with h5py.File(h5_path, "r") as f:
+                    raw_coords = f["coords"][:]
+                    barcodes = f["barcode"][:].flatten()
+            else:
+                raise FileNotFoundError(
+                    f"Neither .pt features nor patches .h5 found for {sample_id}"
+                )
 
         st_dir = os.path.join(args.data_dir, "st")
         h5ad_path = os.path.join(st_dir, f"{sample_id}.h5ad")
diff --git a/tests/recipes/hest/test_dataset.py b/tests/recipes/hest/test_dataset.py
index 1097fc5..58967b2 100644
--- a/tests/recipes/hest/test_dataset.py
+++ b/tests/recipes/hest/test_dataset.py
@@ -211,8 +211,8 @@ def test_hest_feature_dataset_neighborhood_dropout():
         # Run multiple times to trigger the stochastic dropout
         dropout_occurred = False
         for _ in range(100):
-            # Batch: (feats, genes, pathways, coords)
-            f, g, _, _ = ds[0]
+            # Batch: (feats, genes, pathways, coords, mask)
+            f, g, _, _, mask = ds[0]
             assert g is None, "Genes should be None in pathway-only mode"
             # Center (index 0) should NEVER be zero
             assert not torch.all(f[0] == 0)