SAM+D: Parameter-Efficient Dimensional Lifting of SAM-Family Models via Depth-Routed LoRA and Depth Shifting

Technische Samenvatting: SAM+D

Probleemstelling
Bestaande fundamentele modellen voor visie, zoals het Segment Anything Model (SAM) en de video-extensie SAM2, zijn van nature ontworpen voor 2D-beelden of 2D+T (video) sequenties. Hoewel deze modellen sterke zero-shot capaciteiten vertonen in 2D-domeinen, missen ze het werkelijke volumetrische begrip dat vereist is voor 3D medische beeldvorming (bijv. CT-scans) en 4D spatiotemporele taken (bijv. 3D time-lapse microscopie). Huidige benaderingen om deze modellen aan te passen aan volumetrische gegevens kampen met significante beperkingen:

1. Slice-by-slice verwerking: Het behandelen van 3D-volumes als onafhankelijke 2D-slices negeert de context tussen de slices, waardoor het niet mogelijk is om overlappende structuren te onderscheiden.
2. Zware hertraining: Volledige 3D-architecturen of methoden die SAM/SAM2 herbouwen met 3D-componenten vereisen vaak enorme geannoteerde datasets en gooien waardevolle 2D pre-trained priors weg.
3. Inefficiënte adaptatie: Bestaande adapter-methoden passen vaak uniforme transformaties toe over het volume, waarbij voorbij wordt gegaan aan het feit dat verschillende anatomische regio's (bijv. grenzen versus interieurs) verschillende adaptaties vereisen. Bovendien introduceert het modelleren van cross-slice afhankelijkheden doorgaans 3D-convoluties of attention-mechanismen die aanzienlijke parameters en computationele kosten met zich meebrengen.
4. Gebrek aan unificatie: De meeste methoden richten zich op een enkel basismodel (ofwel SAM of SAM2) en een enkele dimensionaliteit, en missen een verenigde architectuur om beide respectievelijk naar 3D en 4D te tillen.

Methodologie: SAM+D
De auteurs stellen SAM+D voor, een parameter-efficiënt framework dat bevroren 2D fundamentele modellen met één ruimtelijke dimensie "verhoogt" (lift) om 3D- en 4D-gegevens te verwerken. De kerninnovatie ligt in het invoegen van twee lichtgewicht, model-agnostische modules in elke bevroren transformer block van de pre-trained encoder, waarbij het overgrote deel van de parameters bevroren blijft.

1. Depth-Routed LoRA (DRLoRA):

   • Mechanisme: In plaats van een enkele LoRA-adapter die uniform wordt toegepast, behoudt DRLoRA meerdere parallelle LoRA-experts ($E$) per transformer block.
   • Routing: Een lichtgewicht router brengt de genormaliseerde dieptepositie ($z = d/D_s$) van een slice in kaart naar een softmax-distributie over deze experts.
   • Functie: Dit stelt het model in staat om diepte-afhankelijke adaptaties te leren. Bijvoorbeeld: kenmerken nabij orgaanranden kunnen andere low-rank updates ontvangen dan homogene interieurregio's. De router is uitsluitend geconditioneerd op de dieptepositie, waardoor de noodzaak voor hulp-load-balancing losses die vereist zijn bij content-gebaseerde Mixture-of-Experts (MoE) methoden wordt geëlimineerd.
   • Injectie: DRLoRA wordt toegepast op de Query ($Q$) en Value ($V$) projecties van de self-attention lagen, terwijl de Keys ($K$) bevroren blijven om de pre-trained feature space te behouden.

2. Depth Shift Module (DSM):

   • Mechanisme: Geïnspireerd door de Temporal Shift Module (TSM) voor video, is DSM hergebruikt om feature channels langs de ruimtelijke diepte-as te verschuiven.
   • Operatie: Vóór de self-attention wordt een fractie ($\alpha$) van de channels in slice $i$ naar voren geschoven vanuit slice $i-1$ en naar achteren vanuit slice $i+1$. De resterende channels blijven ongewijzigd.
   • Efficiëntie: Deze operatie bevat nul leerbare parameters en nul multiply-accumulate (MAC) operaties. Het veroorzaakt slechts een verwaarloosbare overhead in geheugentrafiek (gemeten op ~1,7% van de encoder block latentie).
   • Doel: Het maakt uitwisseling van cross-slice features mogelijk, wat zorgt voor volume-niveau context zonder de kosten van 3D-convoluties of cross-slice attention.

Implementatie-instellingen

• 3D Instelling (SAM $\to$ 3D): DRLoRA en DSM worden ingevoegd in de bevroren ViT-B encoder van SAM. Het model verwerkt een geresamplede volume van slices in een enkele forward pass. Een leerbare diepte-embedding wordt toegevoegd aan prompt-embeddings om voorspellingen te moduleren op basis van de nabijheid van de prompt tot de slice.
• 4D Instelling (SAM2 $\to$ 3D+T): Dezelfde modules worden ingevoegd in de bevroren Hiera encoder van SAM2. Temporele consistentie wordt afgehandeld door de native memory architectuur van SAM2, terwijl DRLoRA/DSM de ruimtelijke diepte afhandelen. LoRA wordt aanvullend geïnjecteerd in de memory attention module om cross-attention voor volumetrische data aan te passen.

Belangrijkste Bijdragen

1. Verenigd Framework: SAM+D is de eerste methode die zowel SAM (2D $\to$ 3D) als SAM2 (2D+T $\to$ 3D+T) verhoogt naar hogere dimensies met behulp van dezelfde lichtgewicht componenten zonder de basisarchitectuur te herontwerpen.
2. Nieuwe Modules: Introductie van DRLoRA voor ruimtelijk adaptieve, diepte-geconditioneerde low-rank updates en DSM voor zero-cost cross-slice feature uitwisseling.
3. Parameter Efficiëntie: De methode traint slechts ~2,8% van de parameters voor SAM en ~3,7% voor SAM2, aanzienlijk minder dan bestaande adapter-gebaseerde of volledig gefinetunede benaderingen.
4. Generalisatie: De aanpak demonstreert generalisatie over verschillende SAM-familie architecturen, doel-dimensionaliteiten (3D, 4D) en domeinen (medische CT, bio-scène celtracking).

Experimentele Resultaten
De auteurs hebben SAM+D geëvalueerd op vijf benchmarks onder een single-point prompt instelling:

• 3D Medische Segmentatie: Geëvalueerd op vier CT tumor segmentatie datasets (KiTS, Pancreas, LiTS, Colon). SAM+D behaalde competitieve of superieure resultaten vergeleken met volledig gesuperviseerde 3D methoden (bijv. nnU-Net) en bestaande SAM-gebaseerde adapters (bijv. 3DSAM-adapter, Med-SA), ondanks het gebruik van veel minder leerbare parameters.
• 4D Spatiotemporele Segmentatie: Geëvalueerd op de Fluo-N3DH-SIM+ dataset van de Cell Tracking Challenge. SAM+D presteerde beter dan slice-gebaseerde SAM2 benaderingen en andere tracking methoden, wat aantoont dat temporele propagatie de segmentatiekwaliteit kan verfijnen voorbij de initiële detector (Cellpose).
• Ablatie Studies: Bevestigden dat zowel DRLoRA als DSM complementaire winsten bieden. Diepte-gebaseerde routing presteerde beter dan content-gebaseerde MoE routers, en DSM bood betere cross-slice context dan parameter-rijke 3D attention of convolutie-adapters.

Betekenis en Claims
Het artikel claimt dat SAM+D een praktische en generaliseerbare methode biedt voor het uitbreiden van pre-trained 2D fundamentele modellen naar hogere-dimensie taken. Door de toegevoegde dimensie (diepte) te behandelen als een specifieke as voor routing en shifting, bereikt de methode een volumetrisch begrip zonder de 2D priors weg te gooien of zware computationele kosten te veroorzaken. De auteurs merken op dat hoewel het huidige werk zich richt op SAM en SAM2, het principe van "dimensionaal verhogen" via parameter-efficiënte adaptatie van nature kan worden uitgebreid naar andere dimensies (bijv. het verhogen van 2D beeldmodellen naar video door tijd als de toegevoegde as te behandelen). Het werk is bescheiden van omvang en erkent dat het nog geen SAM3 (tekst-geprompte) heeft geëvalueerd en dat volledige automatische 4D tracking nog steeds externe detectoren vereist voor celdeling gebeurtenissen.