💻 computer science

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

The paper proposes SLAM-AGS, a slide-label-aware multi-task pretraining framework that utilizes Adaptive Gradient Surgery to resolve conflicting task gradients and stabilize learning, thereby significantly improving downstream bag-level prediction and abnormal cell retrieval in computational cytology, especially under conditions of low witness rates and unreliable instance-level labels.

Original authors: Marco Acerbis, Swarnadip Chatterjee, Christophe Avenel, Joakim Lindblad

Published 2026-01-15

📖 5 min read🧠 Deep dive

CC0 1.0

Original authors: Marco Acerbis, Swarnadip Chatterjee, Christophe Avenel, Joakim Lindblad

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to find a single, tiny, rare criminal hiding in a massive city of 1,000 innocent people. This is essentially the challenge doctors face when using AI to analyze computational cytology (looking at cells under a microscope to find cancer).

Here is a simple breakdown of the paper's solution, SLAM-AGS, using everyday analogies.

The Problem: The "Needle in a Haystack" Dilemma

In this medical world, a "slide" is like a whole city (a massive image with millions of pixels). A "patch" is a single neighborhood or block within that city.

The Label Issue: Doctors usually only know if the entire city has a criminal (the slide is "positive" or "negative"). They don't have the time or money to point out exactly which specific neighborhood contains the criminal.
The Witness Rate Issue: Even in a "criminal" city, the bad guy might be hiding in just 1 out of 1,000 blocks. In some cases, it's as rare as 1 out of 2,000.
The AI Struggle: If you train an AI to find the criminal using only the "city is guilty" label, it gets confused. It sees 999 innocent blocks and thinks, "Maybe the whole city is innocent," or it gets so confused by the conflicting clues that it stops learning entirely.

The Solution: SLAM-AGS (The Smart Detective)

The authors created a new training method called SLAM-AGS. Think of it as a two-part training camp for the AI detective, designed to handle the "needle in a haystack" problem without getting confused.

1. Two Different Training Drills (Multi-Task Learning)

Instead of giving the AI just one job, they give it two simultaneous drills based on what kind of "city" it is looking at:

Drill A: The "Innocent City" Drill (Slide-Negative)
If the AI looks at a city known to be 100% innocent, every single block is innocent. The AI is told: "Group all these innocent blocks together tightly." This helps the AI learn what "normal" looks like very clearly.
- Analogy: It's like teaching a student to recognize a perfect, clean room. Every object in the room belongs there.
Drill B: The "Suspicious City" Drill (Slide-Positive)
If the AI looks at a city known to have a criminal, it knows some blocks are bad, but it doesn't know which ones. The AI is told: "Look at these blocks and learn to tell them apart from each other, even if you don't know exactly which one is the criminal."
- Analogy: It's like showing a student a bag of mixed fruits where you know there is a rotten apple inside, but you don't know which one. The student learns to spot differences between the fruits so they can eventually find the bad one.

2. The "Gradient Surgery" (Fixing the Confusion)

Here is the tricky part. When you try to teach the AI both drills at the same time, the instructions can clash.

The Conflict: Drill A says, "Make all these blocks look the same!" while Drill B says, "Make these blocks look different!"
The Result: Without help, the AI gets a headache (mathematically called "gradient conflict") and stops learning or learns the wrong thing.
The Fix (Adaptive Gradient Surgery): The authors added a "surgeon" to the training process. This surgeon looks at the conflicting instructions, cuts out the part of the instruction that causes the fight, and then rescales the remaining instructions so the AI doesn't move too slowly.
- Analogy: Imagine two coaches shouting at an athlete. Coach A says "Run North!" and Coach B says "Run East!" The athlete spins in circles. The "Gradient Surgeon" steps in, says, "Okay, we'll run Northeast, but we'll run just as fast as if we were running North," ensuring the athlete keeps moving forward without spinning out.

The Results: Finding the Needle

The researchers tested this on a dataset of bone marrow cells (simulating the "needle in a haystack" scenario with very few bad cells).

The Competition: They compared their method against standard AI training, pure self-learning, and other hybrid methods.
The Outcome:
- At easy levels (10% bad cells): Everyone did okay, but SLAM-AGS was the best.
- At hard levels (0.5% bad cells): The other methods collapsed. They basically started guessing randomly because the signal was too weak.
- SLAM-AGS: It kept performing strongly. It could still identify the "bad" neighborhoods (instances) and correctly label the "city" (slide), even when the bad cells were incredibly rare.

Summary

SLAM-AGS is a smarter way to train AI to find rare diseases in cell images. It teaches the AI two different lessons at once (recognizing normal groups and spotting differences in suspicious groups) and uses a "surgical" technique to stop those lessons from fighting each other. This allows the AI to remain sharp and accurate even when the disease is hiding in just a tiny fraction of the sample.

The authors have made their code open-source so others can try this "surgical" training method themselves.

Technical Summary: SLAM-AGS

1. Problem Statement

Computational cytology faces two critical challenges that hinder the application of deep learning for cancer detection:

Unreliable and Costly Instance Labels: Obtaining exhaustive patch-level annotations for gigapixel Whole Slide Images (WSIs) is prohibitively expensive. Consequently, fully supervised training is often infeasible.
Extremely Low Witness Rates: In cytology WSIs, the ratio of positive patches (containing abnormal cells) to total patches is often less than 1%. At such low witness rates, standard weakly supervised methods (relying solely on slide-level labels) are dominated by negative instances. This leads to overfitting, poor generalization to rare positives, and model collapse where the model fails to distinguish positive bags from negative ones.

While Multiple Instance Learning (MIL) is a standard solution for slide-level labeling, it struggles when instance-level discrimination is weakly supervised and positive signals are sparse. Furthermore, existing multi-task pretraining methods that combine self-supervised and weakly supervised objectives often suffer from gradient conflicts, destabilizing training and limiting the benefits of weak supervision.

2. Methodology: SLAM-AGS

The authors propose SLAM-AGS (Slide-Label Aware Multi-task pretraining using Adaptive Gradient Surgery), a framework designed to jointly optimize representation learning for both slide-negative and slide-positive patches while stabilizing the training process.

2.1. Multi-Task Pretraining Objectives

The framework adapts the WeakSupCon approach, treating patches based on their slide-level labels:

Slide-Negative Patches (Supervised Similarity): Patches from healthy (negative) slides are treated as having a known, shared negative label. The model optimizes a Similarity Loss derived from Supervised Contrastive (SupCon) loss to encourage a compact representation of these negative instances in the feature space.
Slide-Positive Patches (Self-Supervised Contrastive): Patches from positive slides (which may contain a mix of normal and abnormal cells) are treated as unlabeled. The model optimizes a SimCLR contrastive objective on these patches to learn robust, label-free visual features.

The total loss is the sum of these two objectives: $L = L_{Similarity} + L_{SimCLR}$ .

2.2. Adaptive Gradient Surgery

To address the instability caused by combining heterogeneous objectives, SLAM-AGS integrates Adaptive Gradient Surgery based on the PCGrad algorithm:

Conflict Resolution: When the cosine similarity between task gradients is negative (indicating conflict), the algorithm projects each task's gradient onto the normal plane of the other task's gradient. This removes destructive interference.
Gradient Rescaling: A critical innovation in this work is the gradient rescaling step. Standard projection can reduce the effective gradient magnitude, slowing convergence. SLAM-AGS computes the norm of the original combined gradient ( $\|g_{sum}\|$ ) and the projected gradient ( $\|g_{pc}\|$ ). If the projected norm is smaller, the gradient is rescaled to preserve the original update magnitude:
$g_{pc} \leftarrow g_{pc} \cdot \frac{\|g_{sum}\|}{\|g_{pc}\|}$
This ensures that while conflicting directions are removed, the overall scale of the learning signal is maintained.

2.3. Downstream Pipeline

The pretrained encoder is integrated into an Attention-based Multiple Instance Learning (MIL) aggregator (specifically PAMIL). This allows for:

Bag-level prediction: Classifying the entire slide (e.g., malignancy detection).
Instance-level retrieval: Using attention scores to identify and retrieve the most abnormal patches within a positive bag.

3. Experimental Setup

Dataset: The Bone Marrow Cytology dataset [14], featuring May-Grünwald-Giemsa/Pappenheim-stained WSIs.
Simulation: To simulate MIL settings with controlled witness rates, "Lymphocytes" (majority class) were labeled as "normal," and seven minority cell types were merged into an "abnormal" class. Abnormal instances were injected into normal bags to create target witness rates of 10%, 5%, 1%, and 0.5%.
Baselines: The method was compared against Weakly Supervised training, SimCLR, WeakSupCon, and a variant of SLAM-AGS without gradient rescaling.
Metrics:
- Bag-level: F1-Score.
- Instance-level: Recall@400 (the number of correctly identified key instances among the top 400 patches ranked by attention score).

4. Key Results

The experiments demonstrate that SLAM-AGS significantly outperforms existing methods, particularly as witness rates decrease:

Bag-Level Classification (F1-Score):
- At 10% and 5% witness rates, SLAM-AGS achieved the highest scores (97.1% and 96.0%, respectively), outperforming Weakly Supervised (53.4%) and WeakSupCon (84.6%).
- At 0.5% witness rates, the advantage was most pronounced. SLAM-AGS achieved 66.4%, whereas Weakly Supervised, SimCLR, and WeakSupCon all collapsed to near-random performance (~50%).
Instance-Level Detection (Recall@400):
- SLAM-AGS maintained a detection rate above 60% across all witness rates (e.g., 65.3% at 0.5% WR).
- Competing methods (SimCLR, WeakSupCon) dropped to near-random levels (~40-45%) at low witness rates, failing to effectively identify key abnormal instances.
Ablation Study: The variant "SLAM-AGS w/o rescaling" performed significantly worse (e.g., 46.0% F1 at 0.5% WR), confirming that the gradient rescaling component is essential for stabilizing optimization and preserving learning signals.

5. Significance and Claims

The paper claims that SLAM-AGS provides a robust solution for computational cytology where positive signals are extremely sparse. The core significance lies in the demonstration that:

Gradient Interference is a Primary Bottleneck: In multi-task pretraining with sparse positive labels, conflicting gradients between self-supervised and weakly supervised objectives can destabilize training.
Adaptive Surgery Enables Stability: By combining gradient projection with magnitude-preserving rescaling, the framework allows supervised and self-supervised objectives to reinforce rather than interfere with each other.
Performance at Extreme Sparsity: The method enables stable pretraining and superior downstream performance even at witness rates as low as 0.5%, a regime where standard weakly supervised and self-supervised methods fail.

The authors conclude that mitigating gradient conflicts and rescaling updates are critical for robust representation learning in scenarios with scarce positive supervision. They note that while instance-level retrieval remains stable, bag-level F1 scores still degrade slightly at the lowest witness rates, suggesting that future work should focus on improving the MIL aggregator to better leverage the identified instance-level signals.

Reproducibility: The authors have released their complete implementation and evaluation framework as open source at https://github.com/Ace95/SLAM-AGS.