VariViT: A Vision Transformer for Variable Image Sizes

1. Problema e Motivação

Os Vision Transformers (ViTs) tornaram-se o estado da arte em aprendizado de representação de imagens, superando redes convolucionais (CNNs) em muitas tarefas. No entanto, a aplicação de ViTs na análise de imagens médicas enfrenta desafios significativos:

• Restrição de Tamanho Fixo: Os ViTs tradicionais exigem que as imagens sejam divididas em patches de tamanho fixo, necessitando que a imagem de entrada seja redimensionada, preenchida (padding) ou recortada para um tamanho pré-definido.
• Perda de Informação e Artefatos: Imagens médicas, especialmente tumores, possuem formas irregulares e tamanhos variados. O redimensionamento (resizing) pode introduzir artefatos, distorcer características anatômicas críticas e degradar informações essenciais para o diagnóstico.
• Razão Sinal-Ruído Variável: Um recorte fixo (bounding box) em tumores de tamanhos diferentes resulta em proporções variáveis entre o tumor (fundo) e o tecido saudável (fundo). Isso faz com que o modelo desperdice recursos computacionais focando em áreas de fundo irrelevantes, em vez do tumor.
• Custo Computacional: Imagens grandes são computacionalmente caras, enquanto imagens muito pequenas podem levar à perda de detalhes. Existe um trade-off entre precisão e eficiência.

2. Metodologia: VariViT

O authors propõem o VariViT, um modelo ViT aprimorado projetado para lidar com imagens de tamanhos variáveis mantendo um tamanho de patch consistente. A arquitetura introduz duas inovações principais:

A. Estratégia de Redimensionamento de Posição "Center and Select"

Para resolver o problema de embeddings posicionais em tamanhos variáveis sem usar interpolação (que pode causar perda de informação):

• Alinhamento Central: Em tarefas de classificação de tumores, os recortes 3D são alinhados pelo centro de massa do tumor. Isso estabelece um ponto de referência consistente.
• Mecanismo: O modelo inicializa uma grade de embeddings posicionais fixos (usando codificação senoidal) para o maior tamanho de imagem possível.
• Seleção Dinâmica: Para imagens menores, em vez de interpolar a grade inteira, o método seleciona um subconjunto de embeddings centrado na grade original. Isso preserva a informação posicional original com baixo custo computacional, evitando aproximações.

B. Estratégias de Batching (Agrupamento)

Para lidar com a variabilidade de tamanhos dentro de um batch de treinamento, o paper propõe duas abordagens:

1. Amostrador de Batch Personalizado (CBS - Custom Batch Sampler): Agrupa imagens do mesmo tamanho dentro de um único batch. Os tamanhos podem variar entre os batches, mas a consistência interna permite processamento eficiente.
2. Acumulação de Gradientes (GA - Gradient Accumulation): Permite batches com imagens de tamanhos variados. O modelo acumula gradientes de vários mini-batches (com tamanhos diferentes) antes de atualizar os pesos, simulando um batch grande e estável.

3. Contribuições Principais

1. Estratégia de Embedding Posicional Flexível: Uma nova abordagem "Center and Select" adaptada para diferentes tamanhos de imagem, eliminando a necessidade de interpolação de grades 3D complexas.
2. Otimização Computacional: Implementação de estratégias de batching (CBS e GA) que reduzem a complexidade computacional e aceleram o treinamento e a inferência.
3. Validação em Tarefas Médicas Desafiadoras: Demonstração da aplicabilidade do VariViT em duas tarefas críticas com tamanhos de tumor altamente variáveis:
   • Predição do genótipo de glioma (status de mutação IDH).
   • Classificação de tumores cerebrais (primário vs. metastático).

4. Resultados Experimentais

Os experimentos foram realizados em dois conjuntos de dados de MRI cerebral 3D (gliomas e tumores cerebrais), comparando o VariViT com ResNet-18, ViT Vanilla e Pix2Struct.

• Desempenho em Predição de IDH (Glioma):

  • O VariViT-GA superou todos os baselines, alcançando um F1-score de 75,5% e um AUC de 0,942.
  • O VariViT-CBS também superou o ResNet-18 e o ViT Vanilla, com tempos de treinamento significativamente menores.
  • Visualizações t-SNE mostraram uma separação de clusters mais clara no modelo VariViT, indicando aprendizado de características mais discriminativas.

• Desempenho em Classificação de Tipo de Tumor:

  • O VariViT-GA obteve o melhor desempenho com F1-score de 76,3% e MCC de 0,706, superando o Pix2Struct e o ResNet-18.
  • O VariViT-CBS superou o Pix2Struct e teve desempenho similar ao ResNet-18, mas com maior velocidade.

• Eficiência Computacional:

  • A estratégia de batching proposta reduziu o tempo de computação em até 30% em comparação com arquiteturas convencionais.
  • O modelo possui aproximadamente 28 milhões de parâmetros, comparável ao ResNet-18 (33M) e Pix2Struct (34,8M).

• Estudos de Ablação:

  • A estratégia "Center and Select" superou outras abordagens de embeddings posicionais (fixos independentes, interpolados fixos ou interpolados aprendidos) na tarefa de classificação de IDH.

5. Significado e Conclusão

O VariViT demonstra que é possível adaptar a arquitetura Vision Transformer para imagens médicas de tamanhos variáveis sem sacrificar a qualidade da representação ou introduzir artefatos de redimensionamento.

• Foco no ROI: Ao permitir tamanhos de recorte variáveis, o modelo concentra-se nas regiões de interesse (tumores), ignorando o excesso de fundo, o que melhora a eficiência do aprendizado de características.
• Viabilidade Prática: A redução de 30% no tempo de computação e a eliminação de etapas de pré-processamento destrutivas (como redimensionamento agressivo) tornam o modelo altamente viável para aplicações clínicas reais.
• Generalização: Embora testado em MRI cerebral, a metodologia é versátil e pode ser adaptada para outras modalidades de imagem e regiões anatômicas, especialmente em cenários onde a variabilidade de tamanho é alta.

Em suma, o VariViT resolve o dilema clássico entre a flexibilidade de entrada e a eficiência computacional em Transformers, oferecendo um estado da arte superior para a análise de tumores cerebrais.