Entropy Minimization without Model Collapse: Mitigating Prediction Bias in Medical Imaging

Resumo Técnico: Minimização de Entropia sem Colapso de Modelo

Problema
A Adaptação em Tempo de Teste (TTA) é essencial para implantar modelos de aprendizado de máquina em cenários do mundo real onde as distribuições de dados sofrem mudanças devido a alterações em sensores, ambientes ou protocolos de aquisição. Na imagem médica, essas mudanças são frequentemente sutis e heterogêneas (por exemplo, diferentes scanners ou protocolos de coloração), mas podem causar uma degradação significativa no desempenho. A Minimização de Entropia (EM) tornou-se o objetivo dominante para a TTA, visando reduzir a incerteza da predição em dados de teste não rotulados. No entanto, a EM sofre de um modo de falha crítico conhecido como colapso de modelo, onde o modelo converge para a predição de uma única classe, independentemente da entrada.

Os autores identificam que a causa raiz deste colapso não é meramente o processo de adaptação em si, mas um fenômeno que eles denominam viés de predição. Sob mudanças de distribuição, os agrupamentos de características (clusters) correspondentes a classes distintas no espaço de representação podem derivar e se fundir parcialmente, enquanto a fronteira de decisão do classificador permanece fixa. Isso induz um desvio na distribuição de classes preditas, onde algumas classes tornam-se super-representadas e outras são suprimidas. A EM padrão amplifica esse viés existente ao estreitar os clusters fundidos, criando um ciclo de feedback positivo que conduz o modelo a uma solução trivial de classe única.

Metodologia: Redução de Viés de Mudança de Distribuição (DSBR)
Para mitigar este modo de falha, os autores propõem a Redução de Viés de Mudança de Distribuição (DSBR), um objetivo de correção de viés projetado para operar exclusivamente em tempo de teste, sem acesso aos dados de origem ou rótulos.

1. Mecanismo Central: O DSBR aborda o viés de predição equalizando a contribuição de cada classe predita para a perda de entropia não supervisionada. Em vez de permitir que clusters dominantes ditem as atualizações de gradiente, o DSBR repondera a perda para evitar que qualquer classe única sobrecarregue o sinal de otimização.
2. Formulação Matemática: O método define uma distribuição alvo equilibrada $p^*_X$ que atribui massa de probabilidade igual ($1/K$) a cada cluster predito. Usando reponderação por importância, a perda para uma amostra $x$ é escalonada por $1 / (K \cdot \hat{p}_{\hat{Y}}(\hat{y}(x)))$, onde $\hat{p}_{\hat{Y}}$ é a frequência estimada da classe predita. Isso garante que amostras em clusters sub-representados contribuam mais para o gradiente, enquanto amostras em clusters sobre-representados contribuam menos, incentivando ativamente a separação dos clusters fundidos.
3. Implementação: Como a verdadeira distribuição de classes preditas é desconhecida em tempo de teste, o DSBR a aproxima usando uma Média Móvel Exponencial (EMA) das frequências de cluster entre os lotes (batches). Isso é particularmente crucial para conjuntos de dados com alto número de classes (ex: ImageNet com $K=1000$), onde tamanhos de lote pequenos levariam a estimativas degeneradas por lote.

Principais Contribuições

• Identificação do Viés de Predição: O artigo estabelece que o viés de predição — um desvio na distribuição de classes preditas presente antes da adaptação começar devido à fusão de clusters de características — é a causa raiz do colapso de modelo na EM. Demonstra que a EM não cria este viés, mas o amplifica até que o colapso ocorra.
• Algoritmo DSBR: Os autores introduzem o DSBR, um método que visa especificamente este modo de falha ao reponderar a perda de entropia para equilibrar as contribuições das classes. Diferente de outras abordagens, ele não depende de filtragem de amostras, aumentos específicos de domínio ou regularizadores complexos.
• Avaliação Abrangente: O estudo projeta cenários de adaptação desafiadores em quatro conjuntos de dados de imagem médica (WILDSCamelyon, Histopantum, MammoBench, GliomaMRI) e no benchmark de visão natural ImageNet-C.

Resultos

• Prevenção de Colapso: Nos benchmarks de imagem médica, métodos EM padrão (como o Tent) frequentemente colapsam para uma acurácia de nível de acaso em pares específicos de domínio origem-alvo. O DSBR previne consistentemente este colapso, alcançando acurácia próxima da ótima em todas as configurações.
• Ganhos de Desempenho: O DSBR supera os estados da arte de referência (incluindo Tent, SAR, ROID e DeYO). Em conjuntos de dados médicos, alcançou uma melhoria média de +4,2% (ResNet) e +2,6% (ViT) sobre as bases não adaptadas, superando significativamente os próximos melhores métodos.
• Robustez: No ImageNet-C, o DSBR demonstrou estabilidade superior sob corrupções severas, corrupções mistas e mudanças extremas de rótulos. Ele lidou eficazmente com tamanhos de lote tão pequenos quanto 1, desde que o decaimento da EMA fosse ajustado apropriadamente, enquanto outros métodos frequentemente falhavam ou colapsavam.
• Validação do Mecanismo: Visualizações de embeddings de espaço de características confirmam que, enquanto a EM causa a fusão e o colapso dos clusters, o DSBR mantém os clusters bem separados durante todo o processo de adaptação.

Significância
O artigo argumenta que estabilizar a minimização de entropia requer uma compreensão mecânica e fundamentada de por que ela falha. Ao isolar o viés de predição como o principal motor do colapso, os autores mostram que corrigir esta única quantidade é suficiente para estabilizar a TTA. A significância deste trabalho reside na demonstração de que a adaptação em tempo de teste robusta em domínios críticos como a imagem médica não requer heurísticas complexas ou filtragem de dados, mas sim uma correção direta da dinâmica de otimização que leva ao colapso do modelo. O DSBR oferece uma solução simples, com poucos hiperparâmetros, que opera inteiramente em tempo de teste, tornando-o altamente aplicável para implantação clínica onde a confiabilidade é primordial.