Resumo Técnico: Pré-treinamento de Redes de Aprendizado de Múltiplas Instâncias com Destilação de Múltiplos Professores a partir de Modelos de Fundação de Lâminas de Patologia

Declaração do Problema

O Aprendizado de Múltiplas Instâncias (MIL - Multiple Instance Learning) estabeleceu-se como o paradigma dominante para a análise de Imagens de Lâmina Inteira (WSIs) na patologia computacional. Nesse framework, uma lâmina é tratada como uma "sacola" (bag) de instâncias de nível de patch, onde um codificador de tiles congelado extrai características e um agregador MIL aprende a combinar essas instâncias em uma predição de nível de lâmina. Apesar da diversidade de arquiteturas de agregadores (ex: baseadas em atenção, Transformers, grafos, modelos de espaço de estados), uma limitação fundamental persiste: os agregadores MIL existentes são tipicamente treinados do zero para cada tarefa downstream.

Essa abordagem "do zero" depende de rótulos limitados de nível de lâmina para aprender simultaneamente representações de instâncias discriminativas e estratégias de agregação eficazes. Consequentemente, esses modelos frequentemente sofrem de otimização instável, sobreajuste (overfitting) a correlações espúrias e baixa transferibilidade entre tarefas. Embora a análise de imagens naturais tenha adotado com sucesso o paradigma de pré-treinar-e-ajustar (pretrain-then-fine-tune) usando grandes conjuntos de dados (ex: ResNet, ViT), aplicar isso ao MIL na patologia enfrenta duas barreiras críticas:

1. Escassez de Dados: Diferente de imagens de nível de tile que podem ser recortadas aos milhões, conjuntos de dados de nível de lâmina de alta qualidade e disponíveis publicamente são limitados a milhares ou dezenas de milhares de amostras, insuficientes para treinar grandes modelos de fundação do zero.
2. Capacidade do Modelo: Os agregadores MIL convencionais são módulos leves e específicos para tarefas, com poucos parâmetros. Eles carecem de capacidade para se beneficiar do pré-treinamento supervisionado convencional da mesma forma que os grandes codificadores de escala, e o ajuste fino direto de grandes codificadores de lâmina para diversas tarefas clínicas pode ser computacionalmente caro e instável na otimização.

Metodologia

Os autores propõem um framework de destilação de múltiplos professores para pré-treinar agregadores MIL leves, transferindo conhecimento de modelos de fundação de lâminas de patologia pré-treinados e de grande escala. O núcleo do método envolve o uso de dois codificadores de lâmina de última geração, TITAN e CARE, como modelos "professores" complementares para guiar o treinamento de várias arquiteturas de estudantes de MIL.

1. Configuração de Destilação de Múltiplos Professores

• Professores:
  • TITAN: Fornece semântica de lâmina inteira generalista e transferível, aprendida através de pré-treinamento multimodal de larga escala.
  • CARE: Oferece representações regionais morfologicamente conscientes e molecularmente informadas.
  • A combinação desses professores fornece sinais supervisores diversos, reduzindo o viés inerente ao depender de uma única representação.
• Estudantes: Nove diversas arquiteturas de MIL (incluindo ABMIL, CLAM, TransMIL, WiKG, DAGMIL, GDFMIL, 2DMamba, AMDMIL e AEMIL) servem como estudantes. Os modelos estudantes são inicializados com pesos aleatórios e treinados para imitar os embeddings de nível de lâmina produzidos pelos professores congelados.
• Construção de Dados: Para evitar a irregularidade computacional de WSIs completas, o pré-treinamento é conduzido em Regiões de Interesse (ROIs) de alta resolução do dataset TCGA-UT-8K. Cada ROI é dividida em uma grade regular de $16 \times 16$ patches de 256 patches, garantindo tamanhos de lote consistentes e consumo de memória estável.

2. Perda de Destilação Normalizada por Dispersão Angular (Perda ADN)

Um desafio central na destilação de múltiplos professores é que diferentes professores podem exibir níveis variados de concentração em suas distribuições de representação na hiperesfera. Se um professor produz embeddings fortemente agrupados enquanto outro é mais disperso, seus sinais de supervisão podem ter escalas diferentes, levando a uma otimização desequilibrada.

Para abordar isso, os autores introduzem a Perda de Destilação Normalizada por Dispersão Angular (Perda ADN).

• Mecanismo: A perda mede a distância angular entre a predição do estudante e o alvo do professor. Crucialmente, este erro é normalizado pela dispersão angular corrente ($\sigma_k$) da distribuição de representação do professor específico.
• Fórmula: Para um professor $k$, a perda é computada como:
  $$L^{(k)}_{distill} = \frac{1}{B} \sum_{i=1}^{B} \frac{\Theta^{(k)}_i}{\sigma_k + \epsilon}$$
  onde $\Theta^{(k)}_i$ é a distância angular, e $\sigma_k$ é o desvio padrão estimado das embeddings do professor em torno de seu centro de atualização.
• Efeito: Esta normalização garante que um professor com embeddings fortemente agrupados (pequeno $\sigma_k$) não domine as atualizações de gradiente simplesmente porque sua magnitude de erro bruta é pequena. Isso equilibra a influência de professores com diferentes características de representação.

3. Adaptação Downstream

Após o pré-treinamento, as cabeças de projeção específicas do professor são descartadas. Os pesos do agregador MIL pré-treinado ($\theta^*$) servem como uma inicialização reutilizável para tarefas downstream. Os autores avaliam dois protocolos de adaptação:

• Ajuste Fino de Parâmetros Completos (Full-Parameter Fine-tuning): Tanto o agregador MIL quanto a cabeça de classificação específica da tarefa são atualizados usando rótulos de nível de lâmina da tarefa downstream.
• Probing Linear (Linear Probing): O agregador MIL pré-treinado é congelado, e apenas a cabeça de classificação linear é treinada, avaliando diretamente a qualidade das representações de nível de lâmina aprendidas.

Principais Contribuições

1. Pré-treinamento de MIL Viável: O artigo demonstra um caminho prático para pré-treinar agregadores MIL leves, apesar da escassez de dados rotulados de nível de lâmina em larga escala, ao alavancar a destilação de conhecimento de modelos de fundação existentes.
2. Framework de Múltiplos Professores: Introduz um framework que utiliza conhecimento complementar de TITAN e CARE, mostrando que combinar professores heterogêneos pode gerar melhor generalização do que a destilação de um único professor.
3. Perda ADN: A proposta da Perda de Dispersão Angular Normalizada aborda o desafio específico de equilibrar sinais de supervisão de professores com diferentes distribuições de representação.
4. Avaliação Abrangente: O método é sistematicamente avaliado em 15 tarefas downstream em 6 conjuntos de dados de referência, cobrindo subtipagem de câncer, predição de biomarcadores e detecção de mutações.

Resultados Experimentais

Os autores conduziram experimentos extensos comparando seus modelos de MIL pré-treinados por destilação contra:

• Modelos treinados do zero.
• Ajuste fino direto dos modelos de fundação (TITAN e CARE).

Principais Descobertas:

• Ganhos de Desempenho: A inicialização pré-treinada melhora consistentemente os agregadores MIL em relação ao treinamento do zero. Isso é particularmente evidente nos cenários de probing linear e few-shot.
• Superando os Professores: Em muitos casos, os agregadores MIL leves pré-treinados superam os pesados codificadores de lâmina (TITAN e CARE) quando adaptados para tarefas downstream, sugerindo que o conhecimento destilado é efetivamente comprimido e especializado para a tarefa de agregação.
• Eficiência Few-Shot: Em configurações de poucas amostras (few-shot settings) (1 a 16 amostras por classe), os modelos inicializados por destilação mostram ganhos significativos (ex: até 20,55 pontos percentuais no BRACS-Coarse) em comparação ao treinamento do zero, destacando melhoria na eficiência de dados e estabilidade de otimização.
• Agnosticismo de Arquitetura: Os benefícios foram observados em várias arquiteturas (atenção, Transformer, grafo, espaço de estados), embora algumas arquiteturas (como GDFMIL e 2DMamba) tenham mostrado desempenho mais variável dependendo da tarefa e do protocolo de adaptação.
• Eficiência de Memória: Ao contrário dos grandes modelos de fundação que sofrem explosão de memória conforme o número de patches aumenta, os agregadores MIL leves mantem o uso de memória estável mesmo com um grande número de patches (ex: 30k patches), tornando-os escaláveis para análise de WSIs de gigapixels.

Significância e Alegações

O artigo afirma que este trabalho fornece uma alternativa prática ao treinamento convencional de MIL do zero. Ao transferir priors de nível de lâmina de grandes modelos de fundação para agregadores leves via destilação, a abordagem permite:

• Priors de Nível de Lâmina Transferíveis: Modelos MIL podem adquirir representações estáveis e generalizáveis sem exigir conjuntos de dados massivos rotulados para seu próprio pré-treinamento.
• Eficiência: Os modelos resultantes permanecem computacionalmente eficientes e escaláveis, adequados para ambientes clínicos com recursos computacionais ou orçamentos de anotação limitados.
• Robustez: O método melhora a velocidade de convergência e reduz o sobreajuste, especialmente em regimes de escassez de dados como o aprendizado few-shot.

Os autores posicionam este trabalho como um passo para remodelar a patologia computacional, tornando o paradigma "pré-treinar-e-ajustar" acessível ao MIL, um componente crítico da análise de WSI que anteriormente era limitado ao treinamento do zero. Sugere-se que trabalhos futuros escalem os dados de pré-treinamento e avaliem o framework em uma gama ainda mais ampla de tarefas downstream.