docs: expand Notes from the Authors with case aggregator, benchmarking, and stage 2 hyperparameters

Add three new subsections under Notes from the Authors:

- On Case Aggregator: argue the second-stage [CASE] module is a
  task-distribution projection head, not just a pooling operation,
  using the 1.0 slide-per-patient RCB task as evidence.
- On Benchmarking: flag that held-out task cohorts partially overlap
  with pretraining corpora, and introduce two cleaner evaluation axes
  (cohorts absent from both stages; task families absent from Stage 2
  within seen cohorts).
- On Hyperparameters of Second Stage: acknowledge the limited
  compute budget for Stage 2 sweeps and call on the community to run
  controlled ablations under the new evaluation protocol.

Author: yousefkotp

README.md

@@ -33,6 +33,9 @@
   - [On the multi-task training dynamics](#on-the-multi-task-training-dynamics)
   - [On the effect of scaling](#on-the-effect-of-scaling)
   - [On the limitations of slide encoders](#on-the-limitations-of-slide-encoders)
+  - [On Hyperparameters of Second Stage](#on-hyperparameters-of-second-stage)
+  - [On Case Aggregator](#on-case-aggregator)
+  - [On Benchmarking](#on-benchmarking)
 - [Acknowledgment](#acknowledgment)
 - [Citation](#citation)
 - [Contact](#contact)
@@ -225,6 +228,20 @@ The issue is that a single vector is, by construction, an information bottleneck
 
 MOOZY itself is a single-vector model, and the results in this repository should be read accordingly. They are evidence of transferability on *scalar* endpoints like classification, survival, and retrieval, not a claim that whole-patient understanding has been solved. Our contribution is orthogonal to the multi-vector question. We argue for *patient-level* rather than slide-level aggregation, and the compression bottleneck is still there. The natural next step is to relax the single `[CASE]` token into a small bank of learned patient latents, and to find a training signal that pushes each latent to represent the slide or patient at a different scale.
 
+### On Benchmarking
+
+We want to flag a methodological caveat in how we evaluate. The cohorts behind our held-out tasks are not entirely out of distribution. Their slides overlap with the cohorts used in both training stages (stage 1 and 2), so the underlying slide distribution (morphology, staining, scanners, demographics) is plausibly familiar to the model, and there is a real possibility that some of our headline numbers are inflated by this overlap. They should be read as transfer within a partially overlapping universe of cohorts, not as pure out of distribution (OOD) generalization to genuinely unseen patient populations. To partially address this, we are currently structuring our evaluation along two axes, and we recommend future patient/slide-encoder work do the same: i) Cohorts absent from both stages, the task is evaluated on a dataset whose slides were never in Stage 1 SSL nor in Stage 2 supervision. This is the cleanest test of cohort-level generalization, and ii) Task families absent from Stage 2 within seen cohorts, the cohort was visible during pretraining, but the family of clinical readout being probed was not part of Stage 2 supervision on that cohort. This isolates how well the representation transfers to a new clinical readout while holding the slide distribution roughly fixed.
+
+The two axes test different things and both matter. Cohort generalization is the more clinically meaningful direction. Task-family generalization is a cleaner probe of representation quality, where with cohort shift held roughly fixed, it isolates whether the learned features carry signal beyond the specific clinical readouts that Stage 2 supervised against. Without reporting along both axes, headline numbers can look like generalization but be partly a function of evaluation-cohort overlap with pretraining.
+
+### On Hyperparameters of Second Stage
+
+Our compute budget did not allow a broad sweep over the second-stage configuration space, learning rate schedules, warmup, weight decay, loss weighting, which slide-encoder layers to fine-tune, freezing slide encoder, and so on. What we ended up reporting is the configuration that performed best on the eight held-out tasks, which is a perfectly normal thing to do, but it inherits the issue we raise in [On Benchmarking](#on-benchmarking). If the held-out tasks themselves draw from cohorts that were partially visible during pretraining, then selecting hyperparameters against those tasks can quietly tune the model to that same overlapping evaluation universe rather than to a genuinely held-out signal. Stage 2 configuration is not optimal in any global sense, it is the best we found within a constrained search, and there are almost certainly better settings we did not get to explore. We strongly encourage the community to run controlled Stage 2 hyperparameter ablations themselves, ideally with model selection done on the two evaluation axes from the previous subsection rather than on a task suite that may share cohort distribution with pretraining.
+
+### On Case Aggregator
+
+We call the second-stage module that turns slide tokens into a `[CASE]` embedding the "case aggregator", and that name is convenient but, we think, slightly misleading. The clearest piece of evidence comes from Residual Cancer Burden task, every patient in RCB contributes exactly one slide, so the slides-per-patient ratio is 1.0. There is nothing to aggregate, and yet, in our case-aggregator ablation, adding the case aggregator helped in that task. Our reading is that the case transformer is best thought of as a *task-distribution projection head*, not just a pooling operation. During Stage 2 it is trained jointly with the slide encoder over supervised tasks, and the `[CASE]` query learns to attend over slide tokens and re-project them into a subspace that is shaped by the distribution of those tasks. When the case has multiple slides, the projection happens to also pool across them, which is the role the name "aggregator" captures. When the case has a single slide, the same module still runs, and it still applies that learned re-projection. Mechanically, it is transforming a slide representation that was optimized largely by Stage 1 SSL, into a representation that lives in the joint geometry of all the case-level pathology tasks the model was aligned to. The downstream probe then sees a more linearly (or shallow-MLP) decodable signal. The implication which we want to flag explicitly, is that "case aggregator" undersells what this module does and over-constrains how the community might think about it. It is a learned, case-level, task-aware projection that *also* aggregates when there is more than one slide. We think the more accurate framing is closer to a case-level adapter that closes the gap between a generic slide representation and the distribution of clinical readouts the model was supervised against. 
+
 ## Acknowledgment
 
 This work was supported by NSERC-DG RGPIN-2022-05378 [M.S.H], Amazon Research Award [M.S.H], and Gina Cody RIF [M.S.H], FRQNT scholarship [Y.K]. Computational resources were provided in part by [Calcul Québec](https://www.calculquebec.ca) and the [Digital Research Alliance of Canada](https://www.alliancecan.ca).