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

Resumo Técnico: RI-Mamba

1. O Problema

A recuperação de formas 3D a partir de texto (text-to-shape retrieval) é crucial para aplicações em realidade virtual, jogos e design, onde os usuários buscam objetos em grandes repositórios usando linguagem natural. No entanto, os métodos existentes enfrentam duas limitações fundamentais que impedem sua aplicação no mundo real:

1. Restrição de Categorias: A maioria dos modelos é treinada em conjuntos de dados pequenos e curados (como Text2Shape), limitando-se a poucas categorias de objetos (ex: apenas cadeiras e mesas).
2. Dependência de Pose Canônica: Os modelos atuais assumem que todos os objetos 3D estão em uma pose padrão (alinhada). No mundo real, os objetos aparecem em orientações aleatórias. Quando os objetos são rotacionados arbitrariamente (rotações SO(3)), o desempenho desses modelos cai drasticamente.
3. Custo Computacional: As abordagens mais robustas para invariância rotacional (como o RI-Transformer) utilizam mecanismos de atenção que têm complexidade quadrática, tornando o treinamento em larga escala e a inferência em grandes conjuntos de dados computacionalmente proibitivos.

2. Metodologia (RI-Mamba)

Os autores propõem o RI-Mamba, o primeiro modelo baseado em State-Space Models (SSM) que é invariante à rotação para nuvens de pontos. A arquitetura combina eficiência linear com robustez geométrica através dos seguintes componentes:

• Desacoplamento de Pose e Geometria (Reference Frame Computation - RFC):

  • Para cada "patch" (aglomerado local) de pontos, o modelo calcula um Referencial Local (LRF) usando Análise de Componentes Principais (PCA).
  • Os pontos são projetados neste referencial local, alinhando-os e removendo a informação de pose global, garantindo que a representação geométrica local seja invariante à rotação.
  • Uma etapa de "desambiguação" é aplicada para garantir que os eixos do referencial tenham direção consistente (evitando ambiguidades de sinal da PCA).

• Serialização Invariante à Rotação (RI-Serialization):

  • Para alimentar o modelo Mamba (que processa sequências), os patches desordenados devem ser convertidos em uma sequência 1D.
  • O modelo define um Referencial Global (GRF) para o objeto inteiro.
  • Os centros dos patches são projetados no GRF e ordenados usando uma Curva de Preenchimento de Espaço de Hilbert. Isso preserva a proximidade espacial e garante que a ordem dos tokens permaneça consistente independentemente da rotação do objeto.

• Embutimentos Orientacionais Lineares (Linear-time Orientational Embeddings):

  • Um desafio crítico em modelos invariantes é que a projeção para o referencial local descarta a informação de orientação global do objeto, reduzindo a expressividade do modelo.
  • O RI-Mamba introduz um embedding orientacional que codifica a orientação de cada patch relativa ao referencial global.
  • Diferente de métodos anteriores que calculam pares relativos (complexidade quadrática), esta abordagem calcula a orientação de cada patch em relação ao GRF de forma independente, mantendo a complexidade linear e sendo compatível com o processamento unidirecional do Mamba.

• Recuperação de Contexto Espacial (FiLM e Scanning Bidirecional):

  • FiLM (Feature-wise Linear Modulation): Os embeddings posicionais e orientacionais são reintegrados aos estados ocultos do modelo via modulação linear, permitindo que o modelo adapte as características dos pontos ao contexto espacial.
  • Scanning Bidirecional: Como o Mamba é causal (processa em uma direção), o modelo alterna a direção de varredura (esquerda-direita e direita-esquerda) entre as camadas para capturar dependências de longo alcance em todas as direções.

• Aprendizado por Contraste Multimodal:

  • O modelo é treinado usando aprendizado por contraste entre modalidades (texto, imagem e forma 3D), alinhando os embeddings do RI-Mamba com os do CLIP.
  • Geração Automatizada de Dados: Para escalar o treinamento para mais de 200 categorias sem anotação manual, o sistema gera automaticamente trios (texto-imagem-forma) usando renderização e legendagem de imagens de modelos 3D brutos.

3. Principais Contribuições

1. Primeira Arquitetura Invariante Baseada em Mamba: Introduz o RI-Mamba, que alcança desempenho competitivo com complexidade computacional linear, superando a ineficiência dos Transformers baseados em atenção para nuvens de pontos.
2. Novo Mecanismo de Embedding Orientacional: Desenvolve uma estratégia de embedding orientacional baseada em patches que opera em tempo linear, resolvendo o problema de perda de informação de pose sem o custo quadrático de métodos anteriores.
3. Escalabilidade sem Anotação Manual: Estabelece um pipeline de treinamento que elimina a necessidade de anotações humanas detalhadas (como segmentação de partes), permitindo o treinamento em conjuntos de dados massivos e diversos.
4. Novo Benchmark: Cria o primeiro benchmark para recuperação texto-forma em objetos com orientações arbitrárias, testando robustez em mais de 200 categorias.

4. Resultados Experimentais

Os experimentos foram conduzidos nos conjuntos de dados Text2Shape (supervisionado) e OmniObject3D (zero-shot), além de ModelNet40.

• Recuperação Texto-Forma (Text-to-Shape):

  • No cenário de rotação aleatória (SO(3)), o RI-Mamba superou significativamente todos os métodos anteriores (incluindo RI-Transformer e LocoTrans). Enquanto outros modelos sofriam quedas drásticas de desempenho com objetos rotacionados, o RI-Mamba manteve alta precisão.
  • No benchmark OmniObject3D (214 categorias), o RI-Mamba alcançou o estado da arte (SOTA) tanto em poses canônicas quanto rotacionadas, demonstrando excelente capacidade de generalização.

• Recuperação 3D-a-3D e Classificação:

  • O modelo demonstrou robustez superior em cenários onde tanto a consulta quanto o banco de dados estão rotacionados, mantendo métricas de mAP (Mean Average Precision) estáveis, ao contrário de modelos não invariantes que falharam nessas condições.

• Eficiência Computacional:

  • Comparado ao RI-Transformer, o RI-Mamba reduziu drasticamente o uso de memória GPU e o tempo de execução. Para sequências de 2048 tokens, o RI-Transformer exigiu mais de 20 GB de memória, enquanto o RI-Mamba manteve-se em cerca de 2 GB, permitindo a escalabilidade para objetos complexos.

5. Significância e Impacto

O trabalho do RI-Mamba representa um avanço significativo na interseção entre visão computacional 3D e processamento de linguagem natural.

• Viabilidade Prática: Ao resolver o problema da invariância rotacional com eficiência linear, o modelo torna viável a implementação de sistemas de busca 3D em ambientes do mundo real, onde os objetos não estão perfeitamente alinhados.
• Escalabilidade: A abordagem de treinamento sem anotação manual abre caminho para o uso de repositórios 3D massivos e não curados, superando a limitação de conjuntos de dados pequenos e específicos.
• Eficiência: A adoção de State-Space Models (Mamba) para tarefas 3D complexas demonstra que é possível obter alta expressividade geométrica sem o custo proibitivo dos Transformers, tornando a tecnologia acessível para aplicações em tempo real.

Em resumo, o RI-Mamba estabelece um novo padrão para a recuperação de formas 3D, combinando robustez geométrica, eficiência computacional e capacidade de generalização em larga escala.