Resumo Técnico: RD-ViT: Vision Transformer Recorrente-Profundo para Segmentação Semântica com Dependência Reduzida de Dados

Declaração do Problema

Os Vision Transformers (ViTs) alcançaram desempenho de última geração em tarefas de previsão densa, como segmentação semântica. No entanto, os ViTs padrão sofrem de alta dependência de dados devido ao seu design arquitetônico: cada camada em uma pilha profunda possui parâmetros únicos que devem ser aprendidos independentemente. Isso resulta em uma contagem de parâmetros que escala linearmente com a profundidade ($O(L \cdot d^2)$), tornando os ViTs propensos a sobreajuste quando os dados de treinamento são escassos. Essa limitação é particularmente crítica na imagem médica, onde conjuntos de dados rotulados são frequentemente pequenos (por exemplo, o benchmark de ressonância magnética cardíaca ACDC contém apenas 100 volumes de pacientes) e a rotulagem é cara e demorada. O artigo aborda a necessidade de arquiteturas que possam alcançar alta precisão de segmentação com dados de treinamento limitados, mantendo a eficiência computacional.

Metodologia

Os autores propõem o RD-ViT (Vision Transformer Recorrente-Profundo), uma arquitetura que adapta o conceito de Transformer Recorrente-Profundo (RDT) do projeto OpenMythos para tarefas de previsão densa. Em vez de uma pilha profunda de blocos transformadores únicos, o RD-ViT emprega um único bloco transformador compartilhado, repetido $T$ vezes. A arquitetura consiste em quatro componentes principais:

1. Prelúdio e Coda: Blocos ViT padrão para extração inicial de recursos e refinamento final, respectivamente.
2. Laço Recorrente: A inovação central onde um único bloco ViT é aplicado iterativamente. Este laço incorpora vários mecanismos especializados:
   • Injeção Estável LTI: Para garantir convergência durante recorrências de longo horizonte, a matriz de transição de estado $A$ é parametrizada como $A = \exp(-\exp(\log d_t + \log A))$. Esta construção garante que todos os autovalores estejam estritamente dentro de $(0, 1)$, assegurando que a recorrência seja uma contração. Isso permite extrapolação de profundidade estável (inferência com mais laços do que no treinamento).
   • Tempo de Computação Adaptativo (ACT): Um mecanismo de parada aprendido atribui uma probabilidade de parada a cada token de patch em cada iteração. Tokens em regiões "fáceis" (por exemplo, fundos homogêneos) param cedo, enquanto tokens em regiões "difíceis" (por exemplo, limites) passam por mais iterações. Isso cria um mapa de computação espacial, alocando recursos dinamicamente.
   • LoRA Específico por Profundidade: Para permitir que o bloco compartilhado se comporte de maneira diferente em várias profundidades sem duplicar pesos de posto completo, adaptadores de Adaptação de Baixo Posto (LoRA) com escalonamento indexado por laço são aplicados.
   • Mistura de Especialistas (MoE): Camadas MoE esparsas opcionais substituem a Rede de Alimentação Direta (FFN) padrão. Um roteador direciona tokens para sub-redes de especialistas específicas, permitindo computação condicional e especialização específica de classe dentro do bloco compartilhado.
3. Cabeça de Segmentação: Faz o upsampling de recursos em nível de patch para previsões em nível de pixel/vóxel usando convoluções transpostas.

A arquitetura é agnóstica à dimensão, suportando entradas 2D (via incorporação de patch Conv2d) e 3D (via Conv3d).

Contribuições Principais

1. Adaptação de Arquitetura: Uma arquitetura completa de segmentação 2D/3D integrando injeção estável LTI, ACT, LoRA específico por profundidade e MoE para previsão densa.
2. Validação de Eficiência de Dados: Validação empírica no conjunto de dados de ressonância magnética cardíaca ACDC mostrando que o RD-ViT pode superar ou igualar ViTs padrão com significativamente menos parâmetros e menos dados, particularmente em configurações 2D.
3. Especialização Emergente: Evidência de que os especialistas MoE se especializam espontaneamente para diferentes estruturas anatômicas (por exemplo, Ventrículo Direito, Miocárdio, Ventrículo Esquerdo) sem supervisão explícita de roteamento, fornecendo um viés indutivo estrutural para segmentação multi-classe.
4. Extrapolação de Profundidade: Demonstração de que o modelo pode inferir com segurança com mais iterações de laço ($T$) do que as usadas durante o treinamento sem degradação de desempenho, graças à garantia de estabilidade LTI.
5. Generalização: Aplicação bem-sucedida da mesma arquitetura a uma tarefa fundamentalmente diferente: segmentação de instância de dentes 3D e numeração a partir de volumes de TCFC, aproveitando a atenção global para raciocínio espacial de longo alcance.

Resultados Experimentais

Segmentação Cardíaca (Conjunto de Dados ACDC)

• Volumétrico 3D:
  • RD-ViT Tiny + MoE (3,0 milhões de parâmetros) alcançou uma pontuação Dice de 0,812, atingindo 99,4% da linha de base ViT Padrão (0,817, 5,7 milhões de parâmetros) enquanto usava apenas 53% dos parâmetros.
  • Ablação: O MoE forneceu o maior ganho (+0,8 pontos percentuais sobre o RD-ViT Tiny). ACT e LoRA forneceram melhorias menores, mas consistentes (+0,3 e +0,2 pontos percentuais, respectivamente).
  • Extrapolação de Profundidade: O desempenho permaneceu estável ao aumentar os laços de inferência de $T=8$ (treinamento) para $T=16$, confirmando a estabilidade da injeção LTI.
• Nível de Fatia 2D:
  • O RD-ViT superou o ViT Padrão em todas as frações de dados. Com 10% dos dados de treinamento, o RD-ViT alcançou 0,774 de Dice versus 0,762 para o ViT Padrão. Com 100% dos dados, alcançou 0,882 versus 0,872.
  • Os autores atribuem isso ao contexto espacial mais rico em fatias 2D, permitindo que o refinamento iterativo seja mais eficaz do que no ambiente 3D de menor resolução.

Segmentação de Dentes (Conjunto de Dados ToothFairy2)

• Aplicado à segmentação de instância de dentes 3D e numeração FDI (33 classes).
• Um modelo RD-ViT Tiny + MoE de fase única (5,4 milhões de parâmetros) alcançou:
  • 0,797 de Dice de Semente (dente versus fundo).
  • 89,1% de precisão de classificação FDI.
• Os resultados sugerem que a autoatenção global do transformador captura efetivamente o contexto de quadrante da mandíbula de longo alcance necessário para a numeração de dentes, uma tarefa onde os campos receptivos locais de CNNs frequentemente lutam.

Eficiência Computacional

• Dinâmica ACT: O tempo de ponderação médio (iterações efetivas) diminuiu de ~2,6 para ~1,4 durante o treinamento, indicando que o modelo aprendeu a processar regiões fáceis mais rapidamente.
• Alocação de Computação: Mapas de parada ACT mostraram maior alocação de computação em limites cardíacos e menor alocação em regiões homogêneas.

Significado e Alegações

O artigo afirma que o RD-ViT oferece uma estratégia arquitetônica viável para segmentação eficiente em parâmetros em domínios médicos limitados por dados. Ao combinar compartilhamento de pesos (recorrência) com computação condicional (MoE), o modelo alcança desempenho próximo ao de última geração com aproximadamente metade dos parâmetros de um ViT padrão.

Os autores enfatizam que a injeção estável LTI é crítica para permitir a extrapolação de profundidade, fornecendo um "controle livre no tempo de inferência" para trocar velocidade e precisão sem retreinamento. Além disso, a especialização espontânea dos especialistas MoE sugere que arquiteturas recorrentes-profundas podem decompor naturalmente problemas complexos de segmentação em subtarefas específicas de classe.

O trabalho destaca uma dicotomia entre desempenho 2D e 3D: enquanto o RD-ViT se destaca em 2D devido ao rico contexto espacial, o desempenho 3D é atualmente mais limitado pela capacidade do modelo do que pela própria arquitetura recorrente. No entanto, a transferência bem-sucedida para imagem dentária demonstra o potencial agnóstico de domínio da arquitetura.

O artigo conclui que, embora o RD-ViT ainda não supere pipelines altamente otimizados e pesados em dados em todos os cenários 3D, ele reduz significativamente a dependência de dados e oferece um novo paradigma para segmentação de imagens médicas eficiente, interpretável e adaptável. Todo o código, notebooks e resultados foram lançados publicamente.