💻 computer science

RI-Mamba: Rotation-Invariant Mamba for Robust Text-to-Shape Retrieval

The paper proposes RI-Mamba, the first rotation-invariant state-space model for point clouds that disentangles pose from geometry using global and local reference frames and Hilbert sorting, enabling robust text-to-shape retrieval across diverse object categories and arbitrary orientations without requiring canonical poses.

Original authors: Khanh Nguyen, Dasith de Silva Edirimuni, Ghulam Mubashar Hassan, Ajmal Mian

Published 2026-02-13

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Khanh Nguyen, Dasith de Silva Edirimuni, Ghulam Mubashar Hassan, Ajmal Mian

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.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 have a massive digital library filled with millions of 3D objects: chairs, mugs, trains, and backpacks. You want to find a specific item, say, "a red backpack," just by typing that into a search bar. This is the dream of Text-to-Shape Retrieval.

However, there's a huge problem with current search engines for 3D objects. They are incredibly picky. If you upload a photo of a backpack that is lying on its side, or if the database has a backpack standing upright, the search engine gets confused. It's like trying to find a specific book in a library where the librarian only recognizes books if they are standing perfectly straight on the shelf. If you tilt the book, they can't find it.

Furthermore, most of these systems only know how to find a tiny handful of objects (like just chairs and tables) and require humans to painstakingly label every single part of every object before the computer can learn. This is slow, expensive, and doesn't work for the messy, real world.

Enter RI-Mamba, the new hero of this story.

The Problem: The "Spinning Top" Dilemma

Think of a 3D object as a spinning top. Current AI models are like people trying to describe a spinning top. If the top is spinning, they can't tell if it's a red top or a blue top, or if it's a top at all. They get dizzy. They also struggle if the object is a "weird" shape they haven't seen before.

The Solution: RI-Mamba

The researchers at the University of Western Australia built RI-Mamba (Rotation-Invariant Mamba). Here is how it works, using some simple analogies:

1. The "Magic Compass" (Rotation Invariance)

Imagine you are looking at a puzzle piece. If you rotate the piece, the shape of the puzzle piece doesn't change, but its position on the table does.

Old AI: Tries to memorize the exact position of the piece on the table. If you rotate it, it thinks it's a completely different piece.
RI-Mamba: It has a "Magic Compass" built into its brain. No matter how you spin the object, the compass instantly re-orientates the object in its mind so it looks "upright" again. It separates what the object is (its geometry) from how it is sitting (its pose). This allows it to recognize a chair whether it's upside down, sideways, or standing tall.

2. The "Hilbert Snake" (Organizing the Chaos)

3D objects are made of thousands of tiny dots (points). To process them, the AI needs to turn this messy cloud of dots into a neat list (a sequence), like turning a pile of LEGOs into a single long train of bricks.

The Challenge: If you just pick dots randomly, the order changes every time you rotate the object. The AI gets confused.
The Fix: RI-Mamba uses a mathematical trick called a Hilbert Curve. Imagine a snake that slithers through the 3D space, visiting every dot in a specific, winding pattern. Because this snake follows the shape of the object itself, no matter how you rotate the object, the snake visits the dots in the exact same order. This gives the AI a consistent "story" to read, even when the object is spinning.

3. The "Memory Lane" (State-Space Models)

Most advanced AI uses "Transformers" (like the technology behind ChatGPT) to understand these lists. But Transformers are like trying to read a book by looking at every single page at once—it gets very heavy and slow as the book gets longer.

RI-Mamba uses Mamba, which is like a super-efficient reader. It reads the story one word at a time, remembering the important parts as it goes. This makes it incredibly fast and allows it to handle huge, complex 3D objects without crashing the computer's memory.

4. The "Contextual Glasses" (Reintroducing Orientation)

Here is the tricky part: When the "Magic Compass" straightens the object, it accidentally throws away the information about how the object was originally sitting. If you just look at a straightened chair, you don't know if it was originally lying on its back.

The Fix: RI-Mamba wears "Contextual Glasses." It calculates a special "orientation tag" for every part of the object. It remembers, "Ah, this part of the chair was originally facing North." It then slips this information back into the AI's brain using a technique called FiLM (Feature-wise Linear Modulation). This lets the AI understand both the shape and the original pose, making it much smarter.

5. The "Self-Taught Student" (No Manual Labels)

Usually, teaching AI requires humans to draw boxes around objects and write descriptions like "This is a red chair."

RI-Mamba is a self-taught student. The researchers used a clever trick: they took 3D models, took pictures of them from different angles, and used another AI to write captions for those pictures automatically. Then, they taught RI-Mamba to match the 3D shape with the picture and the caption.
Because it didn't need humans to label everything, they could train it on 200+ categories of objects (from trains to mugs to bunk beds) instead of just 13. It learned to find a "blue planter" even if it had never seen one before.

The Result

In tests, RI-Mamba was a champion.

Robustness: It could find a "wooden bench" even if the bench was rotated 90 degrees, while other systems failed completely.
Speed: It was much faster and used less computer memory than the previous best methods.
Versatility: It works on a massive variety of objects, not just the boring ones.

In a Nutshell

RI-Mamba is like a super-smart librarian who can find any book in the library, no matter how it's tilted, spun, or stacked. It doesn't need a human to tell it what every book is called; it learns by looking at pictures and reading descriptions on its own. And it does all this incredibly fast, making it ready for real-world applications like virtual reality shopping or finding lost items in a 3D warehouse.

1. Problem Statement

The rapid expansion of 3D assets in virtual reality, gaming, and digital design has created a need for efficient Text-to-Shape retrieval. However, existing methods face two critical limitations that hinder real-world applicability:

Canonical Pose Dependency: Most state-of-the-art (SOTA) models assume objects are in a predefined "canonical" orientation. They fail when objects appear in random rotations (arbitrary SO(3) orientations), which is common in real-world datasets.
Scalability and Annotation: Current approaches rely heavily on human-annotated part segmentation labels and curated datasets (e.g., only chairs and tables). This limits their ability to scale to diverse object categories (200+) and prevents training on large-scale, unlabelled 3D repositories.

Existing Rotation-Invariant (RI) models often suffer from high computational costs (quadratic complexity in Transformers) or limited expressiveness due to the loss of pose information during invariance enforcement.

2. Methodology: RI-Mamba

The authors propose RI-Mamba, the first rotation-invariant architecture based on State-Space Models (SSMs), specifically the Mamba architecture. The pipeline consists of four key stages:

A. Reference Frame Computation (RFC) & Local Alignment

To disentangle geometry from pose, the model processes the point cloud in local patches:

Local Reference Frame (LRF): For each patch of points, a local reference frame is computed using Principal Component Analysis (PCA).
Disambiguation: To ensure deterministic results, the sign of PCA axes is resolved by projecting points onto the axes and selecting the direction with the majority of points.
Alignment: Points within a patch are projected into their LRF, creating a rotation-invariant geometric representation ( $\hat{p}_i$ ).

B. Rotation-Invariant Serialization

Mamba requires a 1D sequence input. To maintain invariance while preserving spatial structure:

Global Reference Frame (GRF): A global frame is computed for the entire object.
Hilbert Sorting: Patch centers are projected onto the GRF and sorted using a Hilbert space-filling curve. This maps 3D spatial proximity to 1D sequence order, ensuring the token order remains consistent regardless of the object's global rotation.

C. Embedding Strategy (The Core Innovation)

Standard RI methods often discard orientation information, reducing model expressiveness. RI-Mamba introduces a novel Patchwise Orientational Embedding:

Geometric & Positional Embeddings: Extracted from the aligned local coordinates (LRF).
Orientational Embedding: Instead of expensive pairwise comparisons (quadratic complexity), the model computes the relative orientation of each patch's LRF with respect to the Global Reference Frame ( $F_i^T F_g$ ). This is processed via an MLP.
Feature-wise Linear Modulation (FiLM): The orientational embedding is reintegrated into the hidden states using FiLM. This allows the model to dynamically scale and shift features based on spatial context, recovering the pose information lost during alignment without breaking rotation invariance.

D. Bidirectional Scanning & Contrastive Learning

Bidirectional Scanning: Since Mamba is inherently causal (unidirectional), the authors employ a Reverse Operator that alternates the scanning direction (left-to-right, then right-to-left) across layers. This captures long-range dependencies from all directions.
Training Strategy: The model is trained using Cross-Modal Contrastive Learning (InfoNCE loss) aligned with CLIP's image and text embeddings. Crucially, they use automated triplet generation (rendering + image captioning) to create training data, eliminating the need for manual annotations and enabling training on massive datasets (e.g., OmniObject3D with 200+ categories).

3. Key Contributions

First RI-SSM Architecture: Introduces RI-Mamba, the first rotation-invariant model based on State-Space Models, achieving linear-time complexity ( $O(N)$ ) compared to the quadratic complexity ( $O(N^2)$ ) of RI-Transformers.
Novel Orientational Embedding: Proposes a linear-time, patchwise orientational embedding strategy compatible with SSMs. It recovers pose information via FiLM, balancing rotation invariance with high model expressiveness.
Annotation-Free Scalability: Demonstrates a pipeline for text-to-shape retrieval on 200+ object categories using automated data generation and contrastive learning, removing the bottleneck of manual part segmentation.
New Benchmark: Establishes the first comprehensive benchmark for text-to-shape retrieval on diverse objects with arbitrary orientations (OmniObject3D).

4. Experimental Results

The authors evaluated RI-Mamba on Text2Shape (supervised) and OmniObject3D (zero-shot) benchmarks.

Supervised Retrieval (Text2Shape): RI-Mamba achieved SOTA performance (RR@1: 13.87) and significantly outperformed all other methods under random rotation (SO(3)), where competitors dropped to near-zero performance.
Zero-Shot Retrieval (OmniObject3D):
- Under random rotation, RI-Mamba achieved 19.34% RR@1, significantly outperforming non-RI models (e.g., DuoMamba at 3.28%) and other RI models (RI-Transformer at 3.76%).
- It demonstrated superior generalization across 214 categories.
3D-to-3D Retrieval & Classification: RI-Mamba maintained high Mean Average Precision (mAP) and classification accuracy across all rotation settings (I/I, I/SO(3), SO(3)/I, SO(3)/SO(3)), proving its robustness.
Efficiency:
- Memory: At 2048 tokens, RI-Mamba uses ~2GB GPU memory, whereas RI-Transformer exceeds 20GB.
- Speed: RI-Mamba scales linearly with token count, while RI-Transformer scales quadratically.

5. Significance

RI-Mamba represents a paradigm shift in 3D retrieval by solving the rotation invariance vs. expressiveness vs. efficiency trilemma.

Practicality: It enables robust retrieval in real-world scenarios where object orientation is unpredictable, a capability previously lacking in scalable models.
Scalability: By removing the dependency on manual annotations and leveraging linear-time SSMs, it opens the door to training on massive, diverse 3D datasets.
Architecture: It successfully adapts the efficient Mamba architecture to 3D geometry, proving that SSMs can handle complex spatial relationships better than Transformers in terms of scalability and memory efficiency.

The code is publicly available, and the work sets a new standard for rotation-robust 3D understanding.