SLAM-AGS: Slide-Label Aware Multi-Task Pretraining Using Adaptive Gradient Surgery in Computational Cytology

AI summary

Overview

Research area: Computational cytology and computational pathology — specifically weakly supervised and self-supervised representation learning for whole-slide image (WSI) analysis, combined with Multiple Instance Learning (MIL).

Technical level: Advanced. The paper assumes familiarity with contrastive learning (SimCLR, SupCon), multi-task gradient conflict methods (PCGrad, MGDA), and attention-based MIL aggregators.

Scope: The paper introduces SLAM-AGS, a slide-label-aware multi-task pretraining framework with adaptive gradient surgery, and evaluates it on a bone-marrow cytology dataset under simulated witness rates from 10% down to 0.5%.

What This Paper Is About

In computational cytology, exhaustive cell-level labels are unreliable and expensive, so models must learn from slide-level labels alone — yet abnormal cells are extremely rare within a slide (witness rates often below 1%), which causes weakly supervised models to overfit to negatives. The authors build on the WeakSupCon multi-task pretraining scheme, which combines a supervised-style similarity objective on slide-negative patches with a SimCLR contrastive objective on slide-positive patches, but show that naive combination creates conflicting gradients that destabilize training. Their goal is to resolve that gradient interference so that both bag-level classification and instance-level abnormal-cell retrieval remain strong even at sub-percent witness rates.

Key Contributions

1. Adaptation and evaluation of multi-task pretraining in cytology. The authors adapt existing multi-task pretraining techniques (WeakSupCon) to the cytology setting and evaluate them under small dataset sizes and witness rates of 10%, 5%, 1%, and 0.5%.
2. Adaptive Gradient Surgery during pretraining. They apply PCGrad to project conflicting task gradients onto each other's normal planes, then add a gradient rescaling step that restores the pre-projection gradient norm when the projected norm is smaller, preserving update scale while keeping non-conflicting directions.
3. A renewed MIL pipeline for cytology. The pretrained ResNet-18 encoder is integrated into a PAMIL (Prototype Attention-Based MIL) aggregator for bag-level prediction plus attention-guided retrieval of the most abnormal instances in a bag.
4. Open-source release. The complete implementation and evaluation framework is shared at https://github.com/Ace95/SLAM-AGS.

Main Findings

• Highest bag-level F1 at every witness rate. SLAM-AGS reaches 97.1 ± 6.4 at 10% WR, 96.0 ± 8.9 at 5% WR, 56.7 ± 14.9 at 1% WR, and 66.4 ± 6.3 at 0.5% WR. For comparison, WeakSupCon scores 60.0 ± 14.9, 84.6 ± 2.6, 55.6 ± 9.6, and 49.5 ± 9.7 respectively; SimCLR scores 80.0 ± 11.2, 84.3 ± 10.3, 51.0 ± 17.5, and 50.0 ± 9.7.
• Baseline collapse at low witness rates. The weakly supervised baseline starts high (92.0 ± 11.0 at 10% WR) but degrades to 53.4 ± 29.8, 53.3 ± 18.3, and 52.7 ± 19.2 as the witness rate falls to 5%, 1%, and 0.5%, with the paper describing this as collapsing to near-random performance.
• Best instance retrieval, with the largest gains at low witness rates. On Top 400 key-instance detection, SLAM-AGS achieves 64.3 ± 7.3 (10% WR), 63.9 ± 8.8 (5% WR), 62.0 ± 19.4 (1% WR), and 65.3 ± 11.9 (0.5% WR) — an average detection rate above 60%, while competing methods drop to roughly 40–63%.
• Self-supervised and hybrid baselines approach random retrieval. SimCLR and WeakSupCon tend toward a random distribution of instances within positive bags, with near-random Top-400 detection rates slightly over 40%, which the authors link to their poorer bag-level classification at lower witness rates.
• The rescaling step is essential. SLAM-AGS without rescaling scores 66.0 ± 10.6, 60.0 ± 9.1, 54.1 ± 16.9, and 46.0 ± 18.9 on F1 and 50.7 ± 13.9, 43.7 ± 14.9, 47.3 ± 13.2, and 46.7 ± 19.4 on Top 400 — consistently below the full method, confirming that projection alone is not sufficient.
• Instance detection is more stable than bag classification for the proposed method. SLAM-AGS Top 400 detection stays relatively stable across witness rates while its bag-level F1 still declines as the witness rate drops, suggesting the PAMIL aggregator struggles to combine correct instance signals under very sparse supervision.
• General downward trend across methods. The paper states that performance decreases for all methods as the witness rate diminishes, since fewer key instances provide a weaker positive signal for the aggregator.

Methodology in Plain English

The authors start from a public bone-marrow cytology dataset of May-Grünwald-Giemsa/Pappenheim-stained WSIs. They treat Lymphocytes (the majority class) as "normal" and merge seven minority hematologic cell classes into a single "abnormal" class. The dataset contains 18,369 normal and 910 abnormal cells for training, and 7,873 normal and 369 abnormal cells for testing. Because no cytology dataset provides reliable instance labels with controlled witness rates, the authors synthesize MIL bags: patches are grouped into bags of 1,000 images, and positive bags are built by injecting a controlled fraction of abnormal instances into otherwise normal bags, matching target witness rates of 10%, 5%, 1%, and 0.5%. The same matching is done for test bags.

For pretraining, a ResNet-18 encoder is shared across two objectives. Patches from slide-negative patients all carry the same known negative label, so they are trained with a similarity loss derived from SupCon that pulls their embeddings together. Patches from slide-positive patients may contain a mix of normal and abnormal cells, so they are trained with the SimCLR contrastive objective using two stochastic augmentations per patch. The total loss is the sum of the two.

The problem is that these two losses can push the shared encoder in conflicting directions. The authors apply PCGrad, which projects each task gradient onto the normal plane of the other when their cosine similarity is negative, removing destructive interference. Because that projection shrinks the combined gradient and slows convergence, they add a rescaling step: they compute the norm of the summed gradient before projection and the norm after, and if the projected norm is smaller, they scale the projected gradient back up to the original norm. This keeps the update size while retaining non-conflicting directions.

The pretrained encoder then feeds a PAMIL aggregator, which performs bag-level classification and assigns class-specific attention scores that let the authors rank patches and retrieve the Top 400 most abnormal instances in positive bags. Training runs for 150 epochs with cosine annealing, a 10-epoch warm-up, initial learning rate 0.001, and batch size 256 with an equal number of positive and negative patches per batch. Each training set has 18 bags (9 positive, 9 negative) and each test set has 6 bags (3 positive, 3 negative). Every configuration is run with five random seeds, and mean and standard deviation are reported. The code is released openly.

Why This Matters

Impact on research. The paper shows that gradient conflict in multi-task medical pretraining is not a minor implementation detail but a decisive factor at low witness rates — the same backbone and losses without rescaling lose roughly 30 F1 points at 10% WR (66.0 vs 97.1) and about 20 points at 0.5% WR (46.0 vs 66.4). It also gives the cytology community a controlled, reproducible benchmark for witness-rate sensitivity, since the authors explicitly note that no cytology dataset with reliable instance labels and controlled witness rates exists and that they therefore simulate one.

Real-world applications.

• Screening programs where only slide-level diagnostic labels are recorded, avoiding the cost of cell-level annotation.
• Hematological malignancy triage from bone-marrow aspirate slides, the dataset used here.
• Attention-guided review, where a cytologist or pathologist is shown the Top 400 ranked patches instead of an entire slide.
• Any WSI-based cancer screening workflow with extreme class imbalance among instances, where rare abnormal cells must be found among a majority of normal ones.

Industry relevance. The approach targets the dominant practical constraint in digital pathology and cytology pipelines: abundant slides but scarce fine-grained labels. A pretraining recipe that stays stable at 0.5% witness rates could reduce the annotation budget needed before a model is deployable, and the open-source release lowers the barrier to reproducing and extending the method. The reliance on a standard ResNet-18 encoder and a PAMIL aggregator means the technique is compatible with existing MIL infrastructure rather than requiring a bespoke architecture.

Future Directions

• Improving the aggregator. The authors observe that instance retrieval stays stable across witness rates while bag-level F1 declines, and explicitly call for future work to improve the aggregator so it better leverages instance-level information.
• Extending to extremely low witness rates. Performance is evaluated down to 0.5% WR; the paper frames further improvement in this regime as an open goal rather than a solved problem.
• Validation on real cytology datasets. The evaluation relies on simulated witness rates built by injecting abnormal cells into normal bags because suitable annotated cytology data with controlled witness rates is unavailable; testing on naturally occurring witness rates remains open.
• Generalization beyond bone marrow. The paper describes cytology as lagging behind histopathology in deep learning adoption, citing limited public datasets and large intra-class variability, which raises the question of whether SLAM-AGS transfers to other cytology domains.

Target Audience

Researchers and practitioners working on computational cytology, computational pathology, and whole-slide image analysis who deal with weak slide-level supervision, extremely low witness rates, or multi-task pretraining. It is also relevant to machine learning researchers interested in gradient conflict resolution methods such as PCGrad and in combining self-supervised contrastive learning with weak labels. Readers need a working understanding of contrastive learning and multiple instance learning to follow the method in detail. The paper is less suited to beginners, since the loss formulations and gradient surgery derivations are presented without introductory exposition.

Authors’ abstract

Computational cytology faces two major challenges: i) instance-level labels are unreliable and prohibitively costly to obtain, ii) witness rates are extremely low. We propose SLAM-AGS, a Slide-Label-Aware Multitask pretraining framework that jointly optimizes (i) a weakly supervised similarity objective on slide-negative patches and (ii) a self-supervised contrastive objective on slide-positive patches, yielding stronger performance on downstream tasks. To stabilize learning, we apply Adaptive Gradient Surgery to tackle conflicting task gradients and prevent model collapse. We integrate the pretrained encoder into an attention-based Multiple Instance Learning aggregator for bag-level prediction and attention-guided retrieval of the most abnormal instances in a bag. On a publicly available bone-marrow cytology dataset, with simulated witness rates from 10% down to 0.5%, SLAM-AGS improves bag-level F1-Score and Top 400 positive cell retrieval over other pretraining methods, with the largest gains at low witness rates, showing that resolving gradient interference enables stable pretraining and better performance on downstream tasks. To facilitate reproducibility, we share our complete implementation and evaluation framework as open source: https://github.com/Ace95/SLAM-AGS.