Reinforcement Learning Improves LLM Accuracy and Reasoning in Disease Classification from Radiology Reports

Resumo Técnico: Aprendizado por Reforço para Melhoria de LLMs Leves em Classificação de Doenças em Laudos Radiológicos

1. Problema e Contexto

A classificação precisa de doenças a partir de laudos radiológicos é fundamental para aplicações clínicas. Embora o Ajuste Fino Supervisionado (SFT) de Grandes Modelos de Linguagem (LLMs) leves tenha demonstrado melhorar a precisão, ele frequentemente degrada a capacidade do modelo de gerar raciocínios explícitos e coerentes. Isso ocorre porque o SFT tradicional foca apenas nos rótulos de saída (doenças), omitindo os passos intermediários de raciocínio.

No domínio médico, a transparência e a explicabilidade são cruciais para a tomada de decisão clínica e para a confiança nas recomendações algorítmicas. Além disso, a dependência de APIs externas para LLMs grandes levanta preocupações de privacidade de dados, enquanto hospedar modelos grandes localmente exige recursos computacionais massivos. LLMs leves (com poucos parâmetros) oferecem um compromisso viável para implantação local, mas enfrentam o desafio de manter o raciocínio após o ajuste fino.

2. Metodologia Proposta

Os autores propõem uma abordagem de duas etapas combinando SFT e Aprendizado por Reforço (RL) para otimizar tanto a precisão da classificação quanto a qualidade do raciocínio, sem a necessidade de supervisão explícita de raciocínio na segunda etapa.

• Etapa 1: Ajuste Fino Supervisionado (SFT)

  • O modelo é ajustado usando apenas os rótulos de doenças (ground-truth) a partir de laudos radiológicos.
  • O objetivo é adaptar o modelo ao domínio da radiologia e melhorar a precisão da classificação.
  • Nesta fase, a seção de raciocínio na saída do modelo é deixada vazia ou não supervisionada, o que pode levar à perda de capacidade de raciocínio (esquecimento catastrófico).

• Etapa 2: Otimização de Política Relativa em Grupo (GRPO)

  • Aplica-se um algoritmo de Aprendizado por Reforço (GRPO) no mesmo conjunto de dados da Etapa 1.
  • Função de Recompensa: Não requer rótulos de raciocínio. A recompensa é baseada em:
    1. Precisão de Classificação: Média de precisão e recall das doenças identificadas.
    2. Correção de Formato: Um incentivo para que o modelo siga o formato <reasoning>...</reasoning><answer>...</answer>. Se a seção de raciocínio estiver vazia ou apenas repetir a lista de doenças, a recompensa de formato é zero.
  • Isso força o modelo a gerar raciocínios explícitos baseados em evidências do texto para maximizar a recompensa total.

• Inferência com Votação Majoritária e Resumo

  • Para cada laudo, o modelo gera cinco inferências independentes.
  • A classificação final da doença é determinada por votação majoritária entre as cinco previsões, aumentando a robustez.
  • Um LLM leve original (não ajustado) é usado para sintetizar as cinco explicações de raciocínio em um único relatório coerente e justificativo.

3. Conjuntos de Dados e Modelos

• Dados: Utilizaram três conjuntos de dados de laudos radiológicos estruturados e anonimizados:
  • MIMIC-CXR: 2.000 relatórios para treino (rótulos gerados por GPT-4o) e subconjuntos de avaliação (100 relatórios anotados por radiologistas e 1.000 "prateados" por GPT-4o).
  • NIH-CXR: 50 relatórios anotados por radiologistas.
  • MIDRC: 50 relatórios anotados por radiologistas.
• Modelos Testados:
  • LLMs Leves: LLaMA 3.1-8B, Qwen 2.5-3B, Phi-3 Mini-128K.
  • Comparativos: MedAlpaca-7B (focado em medicina) e BioClinicalBERT (classificador tradicional).
  • Referência: GPT-4o (usado como modelo "professor" para gerar dados de treino e como avaliador).

4. Métricas de Avaliação

• Classificação: Precisão, Recall e F1-Score micro-médiados.
• Qualidade do Raciocínio:
  • Recall de Raciocínio: Proporção de doenças verdadeiras mencionadas no texto de raciocínio gerado.
  • Abrangência do Raciocínio (Comprehensiveness): Proporção de doenças previstas que são suportadas por evidências explícitas no raciocínio.
• Validação Humana: Dois radiologistas certificados avaliaram a qualidade do raciocínio em 50 casos, preferindo a saída do modelo SFT+GRPO na maioria das vezes.

5. Resultados Principais

• Desempenho de Classificação:
  • O SFT superou consistentemente os modelos de base (baselines).
  • A adição do GRPO melhorou ainda mais o F1 micro em 8 de 9 coortes avaliadas.
  • O modelo LLaMA 8B + SFT + GRPO alcançou uma precisão de classificação quase equivalente à do GPT-4o (modelo professor).
  • Em comparação, o BioClinicalBERT teve desempenho inferior ao LLaMA em benchmarks gerais, embora tenha superado o Phi-3 em alguns casos.
• Recuperação e Melhoria do Raciocínio:
  • O SFT puro causou uma perda quase total da capacidade de raciocínio no modelo LLaMA (o modelo apenas listava doenças sem justificativa).
  • O GRPO restaurou e melhorou significativamente o raciocínio. Por exemplo, no conjunto NIH-CXR, o recall de raciocínio do LLaMA saltou de 0.00 (após SFT) para 0.567 (após GRPO).
  • O Qwen e o Phi-3 também mostraram ganhos substanciais em recall e abrangência de raciocínio com o GRPO.
• Comparação com SFT Adicional:
  • Comparado a um cenário onde se faz SFT adicional com dados de raciocínio (que exige anotação humana cara), o SFT+GRPO obteve desempenho de classificação superior ou comparável, mas com menor custo de dados (não precisou de novos rótulos de raciocínio).
  • O SFT+GRPO gerou raciocínios mais robustos do que apenas SFT em doenças, embora o SFT com dados de raciocínio explícito tenha tido ligeira vantagem em algumas métricas de recall de raciocínio, o GRPO foi mais eficiente em melhorar a precisão geral sem dados adicionais.

6. Contribuições Chave

1. Abordagem Híbrida SFT+GRPO: Demonstrou que é possível recuperar e melhorar o raciocínio de LLMs leves após o ajuste fino supervisionado, utilizando apenas recompensas baseadas em regras (precisão e formato), sem necessidade de anotação de raciocínio.
2. Eficiência de Dados: O método elimina a dependência de grandes conjuntos de dados anotados com raciocínios detalhados, que são caros e difíceis de obter na medicina.
3. Validação em Múltiplos Modelos e Conjuntos: Resultados consistentes em três modelos diferentes (LLaMA, Qwen, Phi) e três conjuntos de dados distintos, validando a generalização da abordagem.
4. Mecanismo de Inferência Robusto: A combinação de votação majoritária e síntese de raciocínio demonstrou aumentar a confiabilidade clínica das previsões.

7. Significado e Limitações

• Significado: Este trabalho oferece um caminho prático para implantar LLMs leves em ambientes hospitalares locais (preservando a privacidade) que não apenas diagnosticam com alta precisão, mas também fornecem justificativas clínicas explicáveis, superando a "caixa preta" comum em modelos de IA médica.
• Limitações:
  • A função de recompensa do GRPO não supervisiona explicitamente atributos como o comprimento do raciocínio ou a alinhamento semântico profundo (apenas formato e acurácia).
  • A avaliação do raciocínio depende de LLMs (GPT-4o/Gemini) como juízes, o que pode introduzir viés, embora tenha sido validado por radiologistas humanos.
  • O método não fornece estimativas de incerteza calibradas para as previsões.

Em resumo, o estudo prova que o Aprendizado por Reforço (GRPO) é uma ferramenta poderosa para alinhar LLMs leves com as necessidades de transparência e precisão da radiologia, permitindo diagnósticos automatizados de alta qualidade sem a necessidade de recursos computacionais massivos ou anotações de raciocínio extensivas.