MedSAE: Dissecting MedCLIP Representations with Sparse Autoencoders

Resumo Técnico: MEDSAE – Dissecando Representações MedCLIP com Autoencoders Esparsos

Declaração do Problema
A integração da Inteligência Artificial (IA) na assistência à saúde, particularmente por meio de Modelos de Linguagem Multimodais de Grande Escala (MLLMs) como o MedCLIP, melhorou as capacidades diagnósticas, mas introduziu desafios significativos quanto à interpretabilidade. Embora esses modelos alcancem alto desempenho em tarefas como análise de radiografias de tórax, seus mecanismos internos permanecem "caixas pretas" opacas. Eles operam em um regime de superposição, onde o número de características aprendidas excede a dimensionalidade do espaço de representação, fazendo com que múltiplas características sejam multiplexadas através de neurônios compartilhados. Isso complica a atribuição de conceitos clínicos específicos a ativações individuais do modelo. Técnicas existentes de interpretabilidade mecânica não foram suficientemente escaladas ou adaptadas para tarefas complexas de visão e linguagem médica, deixando uma lacuna entre a IA médica de alto desempenho e a transparência necessária para a confiabilidade clínica.

Metodologia
Os autores propõem o MedSAE, um framework que aplica Autoencoders Esparsos (SAEs) ao espaço latente do MedCLIP para desentrelaçar características superpostas em conceitos aproximadamente monossêmicos compreensíveis por humanos. O pipeline consiste em três etapas principais:

1. Extração de Embeddings e Treinamento:

   • Os autores utilizam o codificador de visão do MedCLIP (congelado) para gerar embeddings a partir do conjunto de dados CheXpert (radiografias de tórax).
   • Para abordar possíveis desalinhamentos entre modalidades, os embeddings são normalizados centralizando (subtraindo as médias por modalidade) e escalando para garantir normas consistentes.
   • Um SAE baseado em ReLU é treinado nesses embeddings normalizados. O SAE mapeia vetores de entrada densos ($x$) para uma representação latente esparsa ($z$) usando uma arquitetura codificador-decodificador com uma penalidade $L_1$ para impor esparsidade:
     $$L := \|x - \hat{x}\|_2^2 + \lambda\|z\|_1$$
   • O objetivo é aprender um conjunto sobrecompleto de direções latentes que separem as características multiplexadas.

2. Avaliação da Monossignificância:

   • Para validar que os neurônios do SAE correspondem a conceitos únicos, os autores calculam a correlação de Pearson entre a ativação de cada neurônio e rótulos clínicos codificados como multi-hot (por exemplo, Edema, Pneumonia).
   • Eles calculam a entropia da distribuição de correlação para cada neurônio em todos os rótulos. Uma entropia mais baixa indica que um neurônio está fortemente concentrado em um único conceito (alta monossignificância), enquanto uma entropia mais alta sugere ativação distribuída entre múltiplos conceitos.

3. Nomeação Automatizada e Validação:

   • Reconhecendo a dificuldade de construir um vocabulário médico robusto para correspondência baseada em decodificador, os autores empregam o MedGemma, um modelo fundacional de visão e linguagem, para identificação automatizada de conceitos.
   • Para cada neurônio do SAE, as imagens com maior ativação são enviadas ao MedGemma por meio de prompts estruturados para gerar uma explicação de conceito compartilhada (por exemplo, "Edema pulmonar grave").
   • Validação: Para garantir que esses nomes gerados não sejam alucinações, uma tarefa de detecção é realizada. O MedGemma recebe a descrição do conceito gerada e um conjunto balanceado de imagens ativadoras e não ativadoras, sendo então solicitado a classificar quais imagens correspondem ao conceito. Isso fornece uma medida quantitativa de alinhamento semântico.

Principais Contribuições

• Framework de Avaliação Inovador: O artigo introduz um método de avaliação quantitativa que combina métricas baseadas em correlação com nomeação e validação automatizada de conceitos usando o MedGemma. Isso vai além da simples análise de ativação para fornecer rótulos interpretáveis por humanos.
• Validação Empírica da Linearidade: Os autores demonstram que os embeddings do MedCLIP satisfazem aproximadamente a hipótese de superposição linear. Por meio de uma verificação de sanidade envolvendo imagens compostas, eles mostraram que embeddings compostos podem ser linearmente reconstruídos a partir de seus componentes com alta similaridade de cosseno (0,97), justificando a aplicação de ferramentas de interpretabilidade linear como SAEs.
• Desentrelaçamento de Características: O estudo mostra que os SAEs podem extrair com sucesso características clinicamente significativas e interpretáveis das representações internas de um modelo de visão e linguagem médico, superando os embeddings brutos do MedCLIP em termos de seletividade de classe.

Resultados
Experimentos conduzidos no conjunto de dados CheXpert produzem as seguintes descobertas:

• Melhoria na Monossignificância: Os neurônios do SAE exibem entropia mais baixa (2,25) em comparação com os embeddings brutos do MedCLIP (2,38), indicando um alinhamento mais forte com conceitos únicos.
• Maior Precisão de Interpretabilidade: O framework de nomeação automatizada identificou 21 neurônios distintos do SAE com precisão de detecção acima de 70% (por exemplo, "Derrame pleural do lado direito com possíveis loculações", "Cardiomegalia grave com congestão vascular pulmonar significativa"). Em contraste, os embeddings brutos do MedCLIP produziram apenas duas interpretações de alta precisão.
• Robustez das Descobertas: Uma condição de controle onde pares de neurônio-conceito foram embaralhados aleatoriamente resultou em uma queda significativa no desempenho de detecção, confirmando que as interpretações significativas surgem da estrutura do SAE e não de artefatos estatísticos.
• Equilíbrio Esparsidade-Reconstrução: O modelo selecionado alcançou um equilíbrio favorável com uma taxa de ativação $L_0$ de 0,20% e uma fração de variância explicada (FVE) de 0,98, apesar de uma taxa de 30% de neurônios "mortos".

Significado e Alegações
O artigo posiciona o MedSAE como um passo escalável em direção a representações clinicamente confiáveis na IA médica. Ao conectar modelos de alto desempenho com transparência, o trabalho oferece um caminho para reverter mecanismos computacionais em conceitos médicos compreensíveis por humanos. Os autores afirmam que esta é a primeira análise empírica de interpretabilidade mecânica em MLLMs médicos usando SAEs.

No entanto, os autores mantêm um tom modesto em relação às limitações. Eles reconhecem que sua abordagem depende da hipótese de representação linear, que pode não se aplicar universalmente, e que sua análise atual é limitada a uma única camada (MedCLIP-ResNet). Eles observam que unidades inativas podem refletir tensões entre capacidade e esparsidade, e não falhas da abordagem, e que a validação com loop humano por especialistas médicos permanece necessária para avaliar a relevância clínica e a segurança. O trabalho é apresentado como um passo fundamental para motivar pesquisas futuras em sistemas de IA médica transparentes, robustos e confiáveis.