LoC-Path: Learning to Compress for Pathology Multimodal Large Language Models

Probleemstelling

De integratie van Multimodale Large Language Models (MLLMs) in de digitale pathologie stuit op fundamentele uitdagingen door de aard van Whole Slide Images (WSI's):

1. Extreme Sequentielengte: WSI's zijn gigapixel-afbeeldingen die worden opgedeeld in duizenden kleine patches (tiles). Dit resulteert in visuele token-sequenties van duizenden lengte, wat de computationele kosten voor attention-mechanismen kwadratisch doet stijgen.
2. Hoge Redundantie: De meeste tokens in een WSI zijn diagnostisch irrelevant. Het weefsel vertoont sterke lokale en globale redundantie (bijvoorbeeld grote gebieden van normaal weefsel).
3. Schaars Signaal: Alleen een klein, specifiek subset van de patches bevat de diagnostisch relevante informatie die nodig is voor het beantwoorden van vragen of het genereren van rapporten.
4. Resource-inefficiëntie: Bestaande methoden (zoals LLaVA-achtige late-fusion modellen) voegen alle visuele tokens toe aan de tekstsequenties. Dit leidt tot hoge geheugenvereisten, lange inferentielatentie en moeilijke end-to-end training, vooral onder beperkte rekenkracht.

Methodologie: LoC-Path

LoC-Path is een efficiënt framework dat "compressie vóór fusie" toepast en een "routing"-mechanisme gebruikt tijdens de fusie. Het doel is om de visuele input te comprimeren tot een compacte set latente representaties en vervolgens alleen de meest relevante tokens naar de LLM te routeren.

Het framework bestaat uit vier kerncomponenten:

1. Sparse Token Merger (STM) voor Lokale Redundantie:

   • Voordat de tokens worden geresampled, worden aangrenzende tokens op het 2D-rooster samengevoegd in niet-overlappende vensters ($s \times s$).
   • Dit vermindert de lokale redundantie (bijvoorbeeld identieke weefselstructuur) en verkort de sequentielengte aanzienlijk, terwijl de grove ruimtelijke indeling behouden blijft.

2. MAE-voorgevormde Resampler voor Globale Redundantie:

   • In plaats van een zware slide-level encoder te gebruiken, gebruikt LoC-Path een lichtgewicht resampler die geïnspireerd is op hoe pathologen een slide scannen (breed kijken en dan comprimeren).
   • MAE Pre-training: De resampler wordt vooraf getraind met een Masked Autoencoder (MAE) doel. Hierbij worden 75% van de tokens gemaskeerd en moet de resampler deze reconstrueren. Dit dwingt het model om een breed overzicht te nemen en diverse, complementaire latente representaties te leren in plaats van een triviaal gemiddelde te berekenen.
   • De resampler comprimeert de lange reeks van samengevoegde tokens naar een compacte set van $L$ latente queries (bijv. 256 tokens).

3. Token Importance Scorer (TIS) voor Query-afhankelijke Selectie:

   • Niet alle comprimerende latente tokens zijn even relevant voor een specifieke tekstvraag (bijv. een VQA-taak).
   • TIS gebruikt de tekst-query om een score te berekenen voor elke visuele latent.
   • Alleen de top-$M$ meest relevante latente tokens worden geselecteerd (bijv. 96 tokens) voor fusie. Dit vermindert de visuele context die de LLM moet verwerken drastisch.

4. Cross-Attention Routing Adapter (CARA) voor Efficiënte Fusie:

   • In plaats van de geselecteerde visuele tokens te concateneren met de tekst (wat leidt tot $O((T+L)^2)$ complexiteit), gebruikt LoC-Path cross-attention.
   • De geselecteerde visuele latents worden via een gated cross-attention mechanisme geïnjecteerd in meerdere decoder-lagen van de LLM.
   • Dit verlaagt de complexiteit van $O((T+L)^2)$ naar $O(T^2 + TM)$, waarbij $T$ de tekstlengte is en $M$ het kleine aantal geselecteerde visuele tokens.

Kernbijdragen

• Characterisering van het WSI-regime: Het paper identificeert dat WSI's een uniek regime hebben met extreme redundantie en schaarse, query-afhankelijke relevantie, wat de standaard-aannames van bestaande MLLMs doorbreekt.
• Resource-efficiënt Framework: LoC-Path combineert lokale token-merging, globale latente compressie en query-afhankelijke routing. Dit maakt end-to-end ontwikkeling en deploy mogelijk op beperkte hardware.
• MAE Pre-training voor Pathologie: Het introduceert een specifieke pre-training strategie voor resamplers in de pathologie, waar grote datasets voor image-text pairs schaars zijn.
• Efficiëntie zonder Prestatieverlies: Het framework reduceert zowel de trainingskosten als de inferentielatentie en -geheugengebruik aanzienlijk, terwijl het concurrerende prestaties behoudt ten opzichte van state-of-the-art modellen.

Resultaten

De auteurs evalueren LoC-Path op diverse benchmarks, waaronder WSI-Bench, Slide-Bench en Path-Bench, en vergelijken het met modellen zoals WSI-LLaVA, SlideChat en Quilt-LLaVA.

• Prestaties: LoC-Path behaalt vergelijkbare of zelfs betere resultaten op taken zoals morfologische analyse, diagnose en rapportgeneratie vergeleken met bestaande slide-level MLLMs. Op de WSI-Bench testset behaalde het een gemiddelde score van 0.739, vergelijkbaar met WSI-LLaVA (0.754), maar met aanzienlijk minder resources.
• Efficiëntie:
  • Inferentie: Het model reduceert de piek-geheugengebruik met 38,9% en de TFLOPs (rekenkracht) met 81,91% ten opzichte van de gemiddelde "General" baselines.
  • Training: De trainingstijd en het geheugengebruik tijdens de tweede fusiestap zijn aanzienlijk lager dan bij LLaVA-style concatenatie.
  • Scalabiliteit: In tegenstelling tot modellen die lineair of kwadratisch groeien met de visuele context, blijft de inferentielatentie van LoC-Path bijna constant, zelfs bij toenemende inputgrootte.
• Feature Kwaliteit: De MAE-voorgevormde resampler toont een hoge "attention coverage" (het model kijkt naar het grootste deel van de slide) en hoge diversiteit in de latente codes, wat aantoont dat het model niet simpelweg een gemiddelde berekent maar relevante gebieden selecteert.

Betekenis

LoC-Path is een doorbraak voor de praktische toepassing van MLLMs in de digitale pathologie. Het lost het fundamentele probleem op van de "gigapixel" aard van WSI's die onverenigbaar is met de beperkte contextvensters en rekenkracht van huidige LLMs.

Door te bewijzen dat compressie en slimme routing (in plaats van brute kracht concatenatie) leiden tot efficiëntere en even goede modellen, maakt LoC-Path het mogelijk om geavanceerde diagnostische AI-systemen te ontwikkelen en te deployen op standaard hardware. Dit opent de deur voor bredere adoptie van AI-assistenten in klinische omgevingen waar rekenkracht vaak een beperkende factor is.