Linear probe merged labels

contrastive-pretraining/README.md

@@ -601,6 +601,105 @@ Outputs in `--results_dir`:
 
 The best epoch by val mean-AUROC is restored before test evaluation. Single-class columns (all 0 or all 1 in the test split) are gracefully reported as `NaN` and excluded from the macro mean.
 
+## Frozen-encoder MIL probing
+
+MIL keeps the visual encoder frozen and trains only a gated
+Classify-Then-Aggregate head. Unlike linear probing, which uses one pooled vector
+per study, MIL operates on the projected visual tokens from every valid series.
+
+Two execution modes are available:
+
+| Mode | Token embeddings | Best use |
+|------|------------------|----------|
+| Cached MIL | Written once as ragged per-study bags | Repeated experiments and faster training |
+| Online MIL | Recomputed each epoch and discarded | Limited storage or one-off experiments |
+
+The encoder configuration must match the contrastive checkpoint, including
+`--encoder`, `--fusion_mode`, `--pooling_strategy`, and `--dim_latent`. MIL
+requires `--fusion_mode late` so valid series remain separate before their tokens
+are concatenated into each study bag. Add `--extra_latent_projection` when that
+projection was used during contrastive pretraining.
+
+### Cached MIL
+
+Extract token-level features for every split:
+
+```bash
+WEIGHTS=/path/to/model_checkpoint.pt
+DATA=/path/to/mri
+REPORTS=/path/to/reports.jsonl
+LABELS=/path/to/labels.csv
+SPLITS=/path/to/splits.csv
+FEATURES=/path/to/mil_features
+
+for SPLIT in train val test; do
+  python scripts/extract_features.py \
+    --weights_path "$WEIGHTS" \
+    --encoder vjepa2 \
+    --fusion_mode late \
+    --pooling_strategy simple_attn \
+    --dim_latent 512 \
+    --data_folder "$DATA" \
+    --jsonl_file "$REPORTS" \
+    --labels_file "$LABELS" \
+    --splits_csv "$SPLITS" \
+    --split "$SPLIT" \
+    --normalizer zscore \
+    --feature_level tokens \
+    --cache_dtype float16 \
+    --out_dir "$FEATURES"
+done
+```
+
+Train and evaluate the MIL head:
+
+```bash
+python scripts/mil_probe.py \
+  --features_dir "$FEATURES" \
+  --results_dir /path/to/mil_results \
+  --epochs 50 \
+  --batch_size 4
+```
+
+Cached MIL stores all retained token embeddings plus study offsets and provenance.
+The pooled `features_<split>.npy` files used by the linear probe cannot be reused
+for MIL. To bound cache size, add `--max_tokens_per_study 2048` during extraction.
+Leaving it at `0` retains every token; any positive limit is deterministic
+subsampling.
+
+### Online MIL without a token cache
+
+Online MIL reconstructs the same frozen encoder and sends each study's tokens
+directly to the same MIL head:
+
+```bash
+python scripts/mil_probe_online.py \
+  --weights_path "$WEIGHTS" \
+  --encoder vjepa2 \
+  --fusion_mode late \
+  --data_folder "$DATA" \
+  --jsonl_file "$REPORTS" \
+  --labels_file "$LABELS" \
+  --splits_csv "$SPLITS" \
+  --results_dir /path/to/mil_online_results \
+  --epochs 50 \
+  --grad_accum_steps 4
+```
+
+The encoder remains in evaluation mode, runs under `torch.no_grad()`, and is
+excluded from the optimizer. Token embeddings are discarded after each study and
+are never written to disk. This avoids the token-cache storage cost but repeats
+encoder computation every epoch.
+
+Studies are encoded one at a time because their numbers of series differ.
+`--grad_accum_steps` provides a larger effective head-training batch size.
+`--max_tokens_per_study` can reduce MIL-head memory after encoding but does not
+reduce encoder computation. Preprocessed image inputs remain supported through
+`--use_preprocessed --preprocessed_dir /path/to/preprocessed`.
+
+Both modes select checkpoints and class thresholds using validation data only,
+then evaluate the fixed model and thresholds on the test split.
+
 ### Reusing features for a different label set (no re-extract)
 
 Because labels are not baked into the encoder features, you only ever run `extract_features.py` **once**. To probe a different label CSV — a different ground-truth rule, or a different grouping — relabel the cached features instead of re-encoding:

contrastive-pretraining/mr_rate/mr_rate/mr_rate.py

@@ -340,6 +340,48 @@ def _encode_visual_tokens(self, image, real_volume_mask, vis_proj_layer,
 
         return merged, token_mask
 
+    def _encode_visual_instances(self, image, real_volume_mask, vis_proj_layer):
+        """Return every valid series token as one study-level MIL bag.
+
+        Contrastive late fusion averages corresponding tokens across series.
+        Classify-Then-Aggregate MIL must instead receive the complete set of
+        series-specific tokens and perform the first study-level aggregation.
+        """
+        if self.fusion_mode != "late":
+            raise ValueError(
+                "Series-preserving MIL extraction requires fusion_mode='late'. "
+                "Other fusion modes aggregate series before projected tokens are available."
+            )
+
+        b, r, _c, _d, _h, _w = image.shape
+        per_series = []
+        tokens_per_series = None
+        for series_index in range(r):
+            rank = dist.get_rank() if dist.is_initialized() else 0
+            if rank == 0:
+                print(
+                    f"[MRRATE rank=0] encoding MIL series {series_index + 1}/{r}...",
+                    flush=True,
+                )
+            encoded = self.run_checkpoint(self.visual_transformer, image[:, series_index])
+            projected = vis_proj_layer(encoded)
+            if projected.ndim != 3:
+                raise RuntimeError(
+                    f"Visual encoder must return [B,T,D] tokens, got {projected.shape}"
+                )
+            if tokens_per_series is None:
+                tokens_per_series = projected.shape[1]
+            elif projected.shape[1] != tokens_per_series:
+                raise RuntimeError("All padded series must produce the same token count")
+            per_series.append(projected)
+
+        stacked = torch.stack(per_series, dim=1)  # [B,R,T,D]
+        token_mask = real_volume_mask.unsqueeze(-1).expand(b, r, tokens_per_series)
+        return (
+            stacked.reshape(b, r * tokens_per_series, stacked.shape[-1]),
+            token_mask.reshape(b, r * tokens_per_series),
+        )
+
     def state_dict(self, *args, **kwargs):
         return super().state_dict(*args, **kwargs)
 
@@ -370,6 +412,7 @@ def forward(
             aug_text=None,
             mode='clip',
             debug=False,
+            return_visual_tokens=False,
             **kwargs
     ):
         b, r, c, d, h, w = image.shape
@@ -380,7 +423,10 @@ def forward(
         if real_volume_mask is None:
             real_volume_mask = torch.ones(b, r, dtype=torch.bool, device=device)
 
-        use_extra_proj = self.extra_latent_projection and return_loss
+        if return_visual_tokens and return_loss:
+            raise ValueError("return_visual_tokens and return_loss are mutually exclusive")
+
+        use_extra_proj = self.extra_latent_projection and (return_loss or return_visual_tokens)
         vis_proj_layer = self.to_visual_latent_extra if use_extra_proj else self.to_visual_latent
         text_proj_layer = self.to_text_latent_extra if use_extra_proj else self.to_text_latent
 
@@ -406,15 +452,26 @@ def forward(
             print(f"[DEBUG MRRATE rank={dist.get_rank()}] text done, encoding visual...", flush=True)
 
         # 2. Encode Visual Tokens (Returns Tokens + Mask)
-        visual_tokens, token_mask = self._encode_visual_tokens(
-            image, real_volume_mask, vis_proj_layer,
-            text_latents=text_latents,
-            num_sentences_per_image=num_sentences_per_image
-        )
+        if return_visual_tokens:
+            visual_tokens, token_mask = self._encode_visual_instances(
+                image, real_volume_mask, vis_proj_layer
+            )
+        else:
+            visual_tokens, token_mask = self._encode_visual_tokens(
+                image, real_volume_mask, vis_proj_layer,
+                text_latents=text_latents,
+                num_sentences_per_image=num_sentences_per_image,
+            )
 
         if debug and dist.is_initialized():
             print(f"[DEBUG MRRATE rank={dist.get_rank()}] visual done: tokens={visual_tokens.shape}, mask={token_mask.shape}", flush=True)
 
+        # Downstream MIL consumes the projected visual instances before the
+        # global masked mean used by standard MR-RATE inference. This opt-in
+        # branch leaves contrastive training and existing inference unchanged.
+        if return_visual_tokens:
+            return visual_tokens, token_mask
+
         # 3. Training Loss
         if return_loss:
             temp = self.logit_temperature.float().clamp(min=0.0, max=4.0).exp()
@@ -518,4 +575,4 @@ def forward(
             mask_bc = token_mask.unsqueeze(-1)
             sum_tokens = (visual_tokens * mask_bc).sum(dim=1)
             count_tokens = mask_bc.sum(dim=1).clamp(min=1.0)
-            return l2norm(sum_tokens / count_tokens)
+            return l2norm(sum_tokens / count_tokens)