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

1. Problema Identificado

El crecimiento explosivo de activos 3D en realidad virtual y videojuegos ha hecho necesaria la búsqueda eficiente mediante texto (text-to-shape). Sin embargo, los métodos existentes presentan dos limitaciones críticas que restringen su aplicabilidad en el mundo real:

• Restricción de Categorías: La mayoría de los modelos dependen de descripciones humanas y anotaciones de segmentación de partes, lo que limita su entrenamiento a conjuntos de datos pequeños y curados con un número reducido de categorías (ej. solo sillas y mesas).
• Falta de Invarianza Rotacional: Los modelos actuales asumen que los objetos 3D están en una "pose canónica" (alineados). En repositorios reales, los objetos aparecen en orientaciones arbitrarias (rotaciones SO(3)). Los métodos actuales sufren caídas drásticas de rendimiento cuando los objetos están rotados, ya que no son invariantes a la rotación.

2. Metodología Propuesta: RI-Mamba

Los autores proponen RI-Mamba, el primer modelo basado en Modelos de Espacio de Estados (State-Space Models - SSM), específicamente la arquitectura Mamba, diseñado para ser invariante a la rotación en nubes de puntos.

Componentes Clave del Pipeline:

1. Desacoplamiento de Pose y Geometría (RFC):

   • Se dividen los puntos en parches locales.
   • Para cada parche, se calcula un Marco de Referencia Local (LRF) utilizando Análisis de Componentes Principales (PCA) con una estrategia de desambiguación de signos (proyectar puntos para determinar la dirección positiva).
   • Los puntos se alinean a este LRF, eliminando la información de pose local y capturando solo la geometría intrínseca.

2. Serialización Invariante a la Rotación (RI-Serialization):

   • Para definir el orden de los tokens en la secuencia (necesario para Mamba), se calcula un Marco de Referencia Global (GRF) para todo el objeto.
   • Los centros de los parches se proyectan al GRF y se ordenan utilizando una curva de relleno de espacio de Hilbert. Esto asegura que el orden de los tokens sea consistente independientemente de la rotación global del objeto, preservando la localidad espacial.

3. Embeddings Orientacionales Lineales:

   • Un desafío de la invarianza a la rotación es perder la información de orientación, lo que reduce la expresividad del modelo.
   • RI-Mamba introduce un embedding orientacional que codifica la orientación de cada parche relativa al GRF global ($F_i^T F_g$).
   • A diferencia de métodos anteriores que calculan orientaciones relativas entre todos los pares de parches (complejidad cuadrática), este enfoque es lineal y compatible con la naturaleza secuencial de Mamba.

4. Bloques RI-Mamba y FiLM:

   • Los embeddings geométricos, posicionales y orientacionales se alimentan a bloques Mamba.
   • Se utiliza Modulación Lineal Específica de Características (FiLM) para reintegrar la información de contexto espacial (posición y orientación) en las representaciones latentes, permitiendo que el modelo adapte las características de los puntos según su contexto.
   • Se emplea un operador de reversión para escanear la secuencia en ambas direcciones (izquierda-derecha y derecha-izquierda) en capas alternas, capturando dependencias bidireccionales sin costo computacional adicional significativo.

5. Aprendizaje Contrastivo Cross-Modal:

   • Para escalar el entrenamiento sin anotaciones manuales, se utiliza un marco de aprendizaje contrastivo cruzado (similar a CLIP/TAMM).
   • Se generan triplets (texto, imagen, forma 3D) automáticamente mediante renderizado y descripción de imágenes (image captioning) sobre grandes conjuntos de datos (ShapeNet, Objaverse, etc.).
   • El modelo se entrena alineando las representaciones 3D con las de CLIP (texto e imagen).

3. Contribuciones Clave

1. Primera Arquitectura Invariante a Rotación basada en SSM: RI-Mamba es el primer modelo Mamba invariante a rotaciones arbitrarias SO(3), logrando un rendimiento competitivo con una complejidad computacional lineal, superando a los Transformers invariables a la rotación (RI-Transformer) que tienen complejidad cuadrática.
2. Embeddings Orientacionales Eficientes: Diseño de embeddings orientacionales por parche que se calculan en tiempo lineal y se integran perfectamente con modelos de espacio de estados, evitando el costo cuadrático de los métodos de pares anteriores.
3. Entrenamiento Escalable sin Anotación Manual: Introducción de un pipeline de generación de datos automatizada y aprendizaje contrastivo que permite entrenar en más de 200 categorías de objetos sin necesidad de etiquetas humanas costosas.
4. Nuevo Benchmark: Establecimiento del primer benchmark para la recuperación texto-forma en objetos diversos con orientaciones arbitrarias, evaluando la robustez en escenarios realistas.

4. Resultados Experimentales

Los experimentos se realizaron en conjuntos de datos como Text2Shape, OmniObject3D y ModelNet40, evaluando tanto en modo supervisado como zero-shot.

• Recuperación Texto-Forma (Text-to-Shape):

  • En Text2Shape (sillas y mesas), RI-Mamba supera a todos los métodos que no usan anotaciones de partes y es competitivo con los SOTA que sí las usan.
  • Bajo rotación aleatoria (SO(3)), RI-Mamba mantiene un rendimiento robusto, mientras que los métodos existentes (como SCA3D) colapsan, fallando en recuperar el objeto correcto.
  • En OmniObject3D (214 categorías, zero-shot), RI-Mamba logra el mejor rendimiento en todas las métricas (RR@1, NDCG), superando a modelos no invariantes (PointBERT, DuoMamba) y a otros modelos invariantes (RI-Transformer, LocoTrans).

• Eficiencia Computacional:

  • RI-Mamba muestra una escalabilidad superior. Mientras que RI-Transformer excede los 20 GB de memoria GPU con 2048 tokens, RI-Mamba se mantiene alrededor de 2 GB, gracias a su complejidad lineal.

• Robustez:

  • El modelo demuestra alta robustez ante el intercambio de ejes (ej. cambiar el eje de gravedad de Y a Z) y en tareas de clasificación 3D zero-shot, manteniendo alta precisión tanto en poses canónicas como rotadas.

5. Significado e Impacto

RI-Mamba representa un avance significativo en la visión por computadora 3D al resolver el dilema entre expresividad (capacidad de entender la orientación) y eficiencia/robustez (invarianza a la rotación y bajo costo computacional).

• Aplicabilidad Real: Al eliminar la necesidad de poses canónicas y anotaciones manuales, habilita la búsqueda de objetos 3D en repositorios masivos y desordenados, un requisito fundamental para aplicaciones de realidad aumentada, diseño de interiores y motores de búsqueda 3D.
• Eficiencia: La adopción de Mamba permite procesar nubes de puntos grandes de manera eficiente, superando las limitaciones de escalabilidad de los Transformers en tareas 3D.
• Generalización: La capacidad de generalizar a más de 200 categorías sin fine-tuning específico demuestra el potencial de los enfoques zero-shot basados en aprendizaje contrastivo para el futuro de la interacción humano-máquina con activos 3D.

En resumen, el trabajo propone una solución integral que combina una arquitectura de espacio de estados eficiente con técnicas geométricas avanzadas para lograr una recuperación de formas 3D robusta, escalable y lista para el mundo real.