Mask-aware inference with State-Space Models

Este artigo apresenta o Partial Vision Mamba (PVM), um novo componente arquitetural que adapta os princípios de convoluções parciais para o modelo State-Space Mamba, permitindo que redes neurais processem eficientemente dados de visão computacional com regiões inválidas ou faltantes em tarefas como completamento de profundidade, inpainting de imagens e classificação.

O essencial

Imagine que você é um pintor talentoso (o seu modelo de Inteligência Artificial) e alguém te entrega uma tela para pintar, mas a tela está cheia de buracos, manchas de tinta velha ou pedaços faltando. O objetivo é completar a pintura ou entender o que está nela, mesmo com esses defeitos.

Até hoje, os "pintores" mais modernos (chamados de Modelos de Espaço de Estado ou Mamba) eram incríveis e rápidos, mas tinham um defeito grave: se você colocasse um pedaço de papel branco (um dado inválido) na tela, eles tentavam pintar sobre o papel branco como se fosse parte da imagem real. Isso estragava tudo, porque o papel branco não tinha informação real, apenas "ruído".

Aqui está a explicação simples do que os autores fizeram:

1. O Problema: O Pintor que ignora os buracos

Os modelos antigos de IA (como CNNs) já tinham uma solução: eles usavam uma "máscara" para dizer: "Ei, pinte apenas onde há tinta real, ignore o papel branco". Eles faziam uma média apenas dos pixels válidos.

Mas os novos modelos, os Mambas, funcionam de forma diferente. Eles leem a imagem como uma sequência de palavras em um livro. Se uma palavra for "vazia" (um buraco na imagem), o modelo tenta ler essa palavra vazia junto com as outras. Como ele não sabe que é vazia, essa "palavra vazia" contamina toda a frase seguinte, estragando o entendimento de toda a imagem. É como se alguém lesse um livro onde algumas páginas foram rasgadas, mas o leitor insistisse em inventar palavras para preencher os rasgos, e essas invenções erradas fizessem o resto da história ficar sem sentido.

2. A Solução: O "PVM" (Mamba Parcial)

Os autores criaram uma nova peça de Lego chamada PVM (Partial Vision Mamba). Pense no PVM como um filtro inteligente ou um guarda-costas que fica na entrada do modelo.

• O Filtro de Entrada (Patch Embedding): Quando a imagem chega, o PVM olha para cada pedaço (patch). Se um pedaço tem 50% de buraco, em vez de jogar fora ou inventar dados, ele usa uma técnica inteligente (uma "média ponderada") para dizer: "Ok, este pedaço tem informação válida, vamos processar apenas o que é real".
• O Token de Máscara: Para os buracos que não têm dados, o PVM não deixa o modelo tentar adivinhar. Ele coloca um "adesivo especial" (um token aprendido) que diz ao modelo: "Aqui não tem nada, ignore este espaço, não deixe isso sujar o resto da história".

3. As Regras do Jogo (O Framework)

Os autores não apenas criaram o PVM, mas escreveram um "manual de instruções" para garantir que ele funcione bem em qualquer lugar. É como se eles dissessem:

• Se você somar duas imagens, só conte o resultado se ambas tiverem dados válidos naquele ponto.
• Se você misturar características (como juntar cores e formas), a "máscara de validade" também deve ser misturada.
• Se um buraco aparecer em qualquer lugar, ele não deve contaminar o vizinho.

4. Onde isso foi testado? (Os 3 Desafios)

Para provar que o PVM funciona, eles o colocaram em três situações diferentes:

• Completar a Profundidade (Depth Completion): Imagine um mapa de um terreno feito por um laser (LiDAR), mas o laser falhou em vários pontos, deixando o mapa cheio de buracos. O PVM conseguiu preencher esses buracos com muito mais precisão do que os modelos antigos, entendendo que os buracos não eram "terreno plano", mas sim "ausência de dado".
  • Resultado: Melhorou a precisão em mais de 23%.
• Restaurar Imagens (Inpainting): Pegar uma foto de uma pessoa com um adesivo grande no rosto e tentar reconstruir o rosto. O PVM conseguiu reconstruir detalhes como o nariz e o cabelo de forma mais realista, sem criar linhas estranhas ou borrões, porque ele sabia exatamente onde estava o adesivo e onde estava a pele real.
• Reconhecer Objetos (Classificação): Tentar identificar um objeto (como um carrinho de compras) mesmo que metade da foto esteja coberta por uma mancha preta. O modelo comum falhava miseravelmente, mas o PVM conseguiu "ver" através da mancha e acertar o nome do objeto com muito mais frequência.

Resumo da Ópera

Pense no PVM como um tradutor inteligente que, ao ler um livro com páginas rasgadas, não tenta inventar palavras para preencher os rasgos. Em vez disso, ele marca os rasgos como "inexistentes" e continua a leitura focando apenas nas palavras reais, garantindo que a história (a imagem) faça sentido do início ao fim.

Isso permite que a tecnologia mais moderna e eficiente (Mamba) seja usada em situações do mundo real, onde as imagens raramente são perfeitas e sempre têm defeitos, buracos ou dados faltando.

Resumo técnico

Título: Inferência Consciente de Máscara com Modelos de Espaço de Estado (PVM)

1. O Problema

Muitas tarefas de visão computacional no mundo real lidam com dados de entrada incompletos ou inválidos devido a erros de sensores, restrições de aplicação ou problemas de privacidade (ex: LiDAR esparsos, regiões anonimizadas, buracos em imagens).

• Limitação das Arquiteturas Atuais: Modelos padrão, como CNNs e os recentes Modelos de Espaço de Estado (SSMs) baseados em Mamba (ex: Vision Mamba, VMamba), são projetados para entradas totalmente válidas. Eles tratam valores inválidos (geralmente preenchidos com zeros ou placeholders) como dados reais. Isso corrompe a extração de características e altera os estados ocultos, levando a falhas na inferência.
• Gap na Literatura: Embora as Convoluções Parciais (PConvs) tenham resolvido esse problema para CNNs ao renormalizar as saídas baseando-se apenas em pixels válidos, não existia um mecanismo equivalente e inerente para arquiteturas sequenciais baseadas em patches e SSMs (como o Mamba). Métodos de Masked Image Modeling (MIM) são estratégias de pré-treinamento e não garantem que o modelo final seja "consciente da máscara" durante a inferência.

2. Metodologia

Os autores propõem uma nova abordagem composta por um componente arquitetural e um framework de design.

A. Framework de Processamento Consciente de Máscara
O trabalho define um framework onde a entrada é um par $(x, m)$, sendo $x$ o tensor de dados e $m$ uma máscara booleana de validade. O princípio central é que todas as operações devem ser condicionadas apenas aos dados válidos. O framework estabelece regras de atualização dinâmica da máscara:

• Operações Elementares: A validade é a interseção lógica (AND) das máscaras de entrada.
• Concatenação e Redimensionamento: A máscara é atualizada conforme a operação.
• Operações com Campo Receptivo (Convoluções, FC, Pooling, SSM):
  • Em uma operação padrão, se qualquer entrada no campo receptivo for inválida, a saída é inválida.
  • Em uma operação parcial (mask-aware), a saída é considerada válida se pelo menos uma entrada no campo receptivo for válida.

B. Partial Vision Mamba (PVM)
O PVM é o componente arquitetural inovador projetado para substituir as camadas padrão do Mamba. Ele resolve dois tipos de invalidade:

1. Invalidade intra-patch: Quando um patch contém dados válidos e inválidos.
   • Solução: Substituição da camada de embedding linear por uma Partial Linear Layer. Esta camada aplica um preenchimento (padding) de média nas posições inválidas antes da projeção linear, garantindo que os tokens resultantes de patches parcialmente válidos sejam tratados como válidos.
2. Invalidade inter-patch: Quando tokens inválidos contaminam a sequência inteira durante o processamento do SSM.
   • Solução: Substituição explícita de tokens inválidos por um Token de Máscara Aprendido (learned masked token), inspirado no BERT. Isso permite que o SSM aprenda a identificar e ignorar esses tokens, evitando a propagação de corrupção para toda a imagem.

C. Regras de Integração
Para que o PVM funcione corretamente, a arquitetura deve seguir regras estritas:

• Camadas anteriores ao PVM devem ser substituídas por suas equivalentes parciais (ex: PConv).
• Em conexões residuais, apenas as posições válidas da entrada devem ser usadas após a fusão.
• O bloco PVM pode atuar como extrator de características parcialmente válidas ou como módulo de completamento de características válidas, dependendo da sua posição na rede.

3. Contribuições Principais

1. PVM (Partial Vision Mamba): O primeiro componente arquitetural baseado em Mamba capaz de processar entradas com dados inválidos de forma arbitrária, garantindo que apenas dados válidos influenciem a saída.
2. Framework Formalizado: Um conjunto de propriedades e regras de design que permitem adaptar qualquer arquitetura baseada em SSM para lidar com dados incompletos.
3. Validação em Três Tarefas Distintas: Demonstração da generalização do método em:
   • Completamento de Profundidade: Tarefa generativa com dados esparsos de LiDAR.
   • Inpainting de Imagem: Preenchimento de buracos arbitrários.
   • Classificação de Imagens com Dados Inválidos: Robustez a oclusões durante a inferência.

4. Resultados Experimentais

Os autores validaram o PVM comparando-o com suas contrapartes "cegas à máscara" (baseadas em VM padrão) e outros baselines.

• Completamento de Profundidade (KITTI-3D):
  • O modelo PVM-DC superou o baseline VM-DC (não consciente de máscara) com uma melhoria relativa de 23% no RMSE (1.38m vs 1.80m).
  • O modelo operou totalmente sem guia de RGB (unguided), dependendo apenas da geometria esparsa, provando a eficácia do mecanismo de máscara.
• Inpainting (FFHQ):
  • Modelos baseados em Mamba (PVM-UNet) superaram significativamente modelos baseados em PConvs em métricas perceptuais (FID e LPIPS), demonstrando a superioridade do contexto global do Mamba.
  • A versão mais complexa, PVM-UNet-N (com conexões de salto conscientes de máscara), obteve os melhores resultados, gerando detalhes mais realistas e menos artefatos visuais.
• Classificação com Dados Inválidos (ImageNet-1k):
  • O modelo PVM-Cls alcançou uma melhoria relativa de 36% na acurácia Top-5 (34.93% vs 25.60%) em comparação ao PlainMamba padrão quando submetido a máscaras irregulares.
• Estudos de Ablação:
  • A estratégia de padding com token aprendido mostrou-se superior ao zero-padding ou mean-padding, embora o SSM tenha capacidade de identificar padrões inválidos mesmo com estratégias subótimas.
  • O modelo demonstrou robustez em diferentes regimes de máscaras (fáceis, difíceis e extremas), não limitando-se apenas ao tipo de máscara usado no treinamento.

5. Significado e Conclusão

Este trabalho preenche uma lacuna crítica na aplicação de Modelos de Espaço de Estado (SSMs) modernos em cenários do mundo real, onde dados incompletos são a norma, não a exceção.

• Inovação: Demonstra que a simples aplicação de SSMs em dados esparsos é insuficiente e que uma adaptação arquitetural explícita (PVM) é necessária para evitar a contaminação de características.
• Generalização: Prova que os princípios de operações parciais, antes restritos a CNNs, podem ser portados com sucesso para arquiteturas de atenção global e SSMs, mantendo a complexidade linear e a eficiência.
• Impacto Futuro: Abre caminho para o uso de modelos Mamba em tarefas críticas como robótica (sensores LiDAR falhos), medicina (imagens com artefatos) e segurança, onde a integridade dos dados de entrada não pode ser garantida.

Em resumo, o PVM transforma o Mamba em uma arquitetura robusta para inferência em dados incompletos, oferecendo ganhos significativos de desempenho sem sacrificar a eficiência computacional.