MT-PCR: Hybrid Mamba-Transformer Network with Spatial Serialization for Point Cloud Registration

Titel

MT-PCR: Hybride Mamba-Transformer Netwerk met Ruimtelijke Serialisatie voor Registratie van Puntwolken

1. Het Probleem

Puntwolkregistratie (Point Cloud Registration - PCR) is een fundamentele taak in 3D-computervisie en robotica, waarbij het doel is om twee gedeeltelijk overlappende 3D-puntwolken uit te lijnen door een optimale starre transformatie (rotatie en translatie) te schatten.

Bestaande oplossingen lijden onder twee belangrijke beperkingen:

• Transformer-architecturen: De meeste huidige state-of-the-art (SOTA) methoden gebruiken Transformer-architecturen. Hoewel deze uitstekend zijn in het modelleren van globale ruimtelijke relaties, hebben ze een kwadratische computationele complexiteit ($O(N^2)$) ten opzichte van de sequentielengte. Dit beperkt de resolutie van verwerkbare puntwolken, wat vaak leidt tot het noodzakelijke downsamplen van data en het verlies van belangrijke geometrische details.
• Mamba-modellen: Mamba, een nieuw model gebaseerd op State Space Models (SSM), biedt lineaire complexiteit ($O(N)$) en sterke lange-afstandscontextmodelleringscapaciteiten. Echter, het direct toepassen van Mamba op puntwolken levert suboptimale resultaten op omdat puntwolken ongestructureerd en ongeordend zijn, terwijl Mamba ontworpen is voor causale, gestructureerde sequenties (zoals tekst).

2. Methodologie

De auteurs stellen MT-PCR voor, het eerste hiërarchische registratiekader dat de efficiëntie van Mamba combineert met de ruimtelijke bewustwording van Transformers. De kern van de methode bestaat uit vier fasen:

A. Ruimtelijke Serialisatie (Z-Order)

Om Mamba toe te passen op ongeordende 3D-data, moeten de punten worden omgezet in een betekenisvolle 1D-sequentie die de ruimtelijke nabijheid behoudt.

• De auteurs gebruiken Z-Order ruimte-vullende krommen (Morton codes).
• 3D-coördinaten worden gekwantiseerd naar een rooster en de bits van de $x, y, z$-coördinaten worden verstrengeld (interleaved) om een index te genereren.
• Dit zorgt ervoor dat punten die dicht bij elkaar liggen in de 3D-ruimte ook dicht bij elkaar liggen in de gegenereerde 1D-sequentie, wat essentieel is voor de effectiviteit van Mamba.

B. Hybride Mamba-Transformer Architectuur

Het netwerk volgt een "coarse-to-fine" (grof-naar-fijn) aanpak:

1. Multi-scale Feature Extractie: KPConv wordt gebruikt om features op verschillende schalen te extraheren.
2. Coarse Matching (Mamba Encoder): Op het niveau van "superpoints" worden features verwerkt door een Mamba-encoder. Dankzij de serialisatie kan Mamba hier efficiënt globale geometrische context modelleren met lineaire complexiteit.
3. Feature Refinement (Transformer): De gesequenteerde features worden vervolgens verwerkt door een Transformer-blok (met cross-attention) om de overeenkomsten (correspondences) verder te verfijnen.
4. Fine Registration: Een lokaal attention-mechanisme en gewogen SVD worden gebruikt voor de uiteindelijke uitlijning.

C. Afschaffing van de "Order-Indicator"

Een cruciale observatie in dit paper is dat de Order-Indicator tokens (die normaal gesproken in Mamba worden gebruikt om de volgorde in sequentiemodellen aan te geven) niet nodig zijn en zelfs schadelijk zijn voor PCR.

• In plaats van de volgorde expliciet te markeren, vertrouwt het model op de inherente ruimtelijke continuïteit van de Z-Order serialisatie.
• Het verwijderen van deze tokens verbetert de prestaties omdat het de geometrische samenhang van de puntwolk behoudt zonder kunstmatige onderbrekingen.

3. Belangrijkste Bijdragen

1. Eerste Hybride Kader: MT-PCR is het eerste framework dat Mamba en Transformers combineert voor hiërarchische puntwolkregistratie.
2. Ruimtelijke Serialisatie: Introductie van een Z-Order gebaseerde serialisatiestrategie die het mogelijk maakt om ongeordende 3D-data efficiënt te verwerken met lineaire complexiteit.
3. Design Insight: Het bewijs dat Order-Indicator tokens overbodig zijn voor 3D-geometrische taken en dat hun verwijdering de registratieprestaties verbetert.
4. Efficiëntie en Schaalbaarheid: Het framework reduceert de GPU-geheugengebruik en FLOPs aanzienlijk ten opzichte van pure Transformer-methoden, terwijl het de nauwkeurigheid behoudt of verbetert.

4. Resultaten

De methode is uitgebreid getest op meerdere benchmarks: 3DMatch, 3DLoMatch (indoor), KITTI (outdoor LiDAR) en ETH-Challenging.

• Nauwkeurigheid:
  • Op 3DMatch bereikt MT-PCR een Registration Recall (RR) van 95,5%, wat een nieuw state-of-the-art resultaat is.
  • Op het uitdagende 3DLoMatch (lage overlap) wordt 75,4% RR behaald.
  • Op KITTI wordt een RR van 100% bereikt met een Relative Rotation Error (RRE) van slechts 0,16°, een verbetering ten opzichte van eerdere beste modellen.
• Efficiëntie:
  • MT-PCR is aanzienlijk sneller en verbruikt minder geheugen dan Transformer-baselines.
  • Bij een sequentielengte van 1536 tokens reduceert MT-PCR de rekencost (FLOPs) met 85,4% (een factor 6,9 lager) en het geheugengebruik met 58,3% vergeleken met GeoTransformer.
  • Waar Transformer-modellen vaak "Out of Memory" (OOM) fouten krijgen bij hoge resoluties, blijft MT-PCR stabiel en efficiënt.
• Generalisatie: Het model toont sterke generalisatievermogen bij overdracht van KITTI naar ETH-datasets en kan zelfs met ongesuperviseerde domeinadaptatie (UDA) binnen één epoch worden aangepast.

5. Betekenis en Conclusie

MT-PCR overbrugt de kloof tussen de efficiëntie van State Space Models (Mamba) en de complexiteit van 3D-puntwolkregistratie. Door de kwadratische complexiteit van Transformers te omzeilen via lineaire Mamba-architecturen en slimme ruimtelijke serialisatie, biedt deze methode een schaalbare oplossing voor hoge-resolutie 3D-scènes.

De bevindingen tonen aan dat:

1. Mamba zeer geschikt is voor 3D-geometrie mits de data correct wordt geserialiseerd.
2. Traditionele sequentiemodulaties (zoals order-indicators) niet altijd nodig zijn voor ruimtelijke data.
3. Een hybride aanpak (Mamba voor globale context + Transformer voor lokale verfijning) de beste balans biedt tussen snelheid en nauwkeurigheid.

Dit werk opent nieuwe wegen voor efficiënte 3D-perceptie in toepassingen zoals autonoom rijden, SLAM en augmented reality, waar rekenkracht en geheugen vaak beperkt zijn.