Graph-Constrained Policy Learning for Extreme Clinical Code Prediction

Resumo Técnico: Aprendizado de Política com Restrição de Grafo para Predição de Códigos Clínicos Extremos

Definição do Problema
A predição de códigos clínicos envolve o mapeamento de sumários de alta hospitalar não estruturados para códigos de diagnóstico padronizados ICD-10-CM. Esta tarefa apresenta um desafio duplo: compreender documentos médicos longos e complexos e realizar uma classificação multi-rótulo extrema sobre um espaço de rótulos que é grande (até 15.761 códigos folha), esparso e profundamente hierárquico. A maioria dos sistemas automatizados existentes trata isso como um problema de classificação multi-rótulo plano, pontuando cada código folha de forma independente. Essa abordagem tem dificuldade com rótulos raros devido aos sinais de treinamento limitados e falha em aproveitar as restrições estruturais da hierarquia ICD, levando frequentemente a predições estruturalmente inválidas ou ao desempenho insatisfatório na "cauda longa" de códigos raros.

Metodologia
Os autores propõem reformular a predição de ICD como um problema de travessia com restrição de grafo, modelado como um Processo de Decisão de Markov (MDP) de horizonte finito sobre uma hierarquia ICD-10-CM podada.

• Formulação do MDP:

  • Estado: Compreende a nota de alta, a fronteira atual de nós pendentes de expansão e o conjunto de códigos folha faturáveis registrados até o momento.
  • Ação: A política seleciona um subconjunto de nós filhos válidos da fronteira atual que são suportados pela nota.
  • Restrição: Um ambiente leve impõe restrições de validade rígidas. O modelo só pode selecionar filhos do conjunto exposto pelo nó atual, garantindo que códigos estruturalmente inválidos recebam probabilidade zero por construção.
  • Terminação: O processo continua nível a nível até que os códigos folha faturáveis sejam alcançados e a fronteia esteja vazia.

• Implementação:

  • Arquitetura: Um único modelo de linguagem decodificador apenas (Qwen3-4B-Instruct, quantizado em 4-bit com QLoRA) atua como a política de travessia. Ele opera dentro de uma janela de contexto contínua, condicionando cada decisão ao histórico acumulado completo da travessia.
  • Ambiente: O modelo interage com um ambiente baseado em ferramentas usando expand_level() para revelar filhos e select() para escolher nós válidos. Isso converte o problema de classificação multi-rótulo extrema em uma sequência curta de decisões de subconjuntos esparsos e restritos por grafo.
  • Treinamento: A abordagem principal é o Ajuste Fino Supervisionado (SFT) em "trajetórias de ouro" — sequências determinísticas geradas através do replay dos caminhos ancestrais corretos dos códigos de ouro através do ambiente. O modelo é treinado para imitar essas demonstrações de oráculo.
  • Comparações: O estudo avalia esta política compartilhada contra:
    1. Baselines Planos: CAML, LAAT e PLM-ICD (retreinados nos mesmos splits).
    2. Baseline de Cascata (SFT-3): Três modelos especialistas independentes (capítulo, subseção, código) operando em uma única etapa por nível.
    3. Aprendizado por Reforço (GRPO): Otimização de Política Relativa de Grupo (Group Relative Policy Optimization) aplicada à política compartilhada.

Contribuições Principais e Desenho Experimental
O artigo conduz um estudo fatorial controlado no conjunto de dados MIMIC-IV (122.197 notas validadas) para isolar os efeitos de arquitetura (política compartilhada vs. cascata), algoritmo de treinamento (SFT vs. GRPO) e orçamento de dados (2.000 vs. 5.000 trajetórias de ouro).

• Espaços de Rótulos: Os experimentos são realizados em duas escalas: um subconjunto "Frequente" (50 códigos mais comuns) e a configuração "Todos" (espaço total de 15.761 códigos).
• Métricas: F1 média micro e média macro são utilizadas, juntamente com Path-F1 (crédito parcial hierárquico) e análise de recall por profundidade.

Resultados

• Desempenho vs. Baselines Planos: A política de travessia com restrição de grafo supera significativamente os baselines planos, particularmente na configuração "Todos".
  • No espaço total de 15.761 códigos, o melhor modelo (SFT-1+) atinge 0,527 micro-F1 e 0,256 macro-F1, superando o baseline plano mais forte (LAAT) em +0,044 micro-F1 e +0,157 macro-F1.
  • Baselines planos (CAML, PLM-ICD) colapsam no espaço total (macro-F1 ~0,002), indicando uma incapacidade de discriminar entre 15.761 candidatos sem orientação estrutural.
• Escalabilidade de Dados: Aumentar os dados de treinamento supervisionado de 2.000 para 5.000 notas é a única intervenção que melhora o desempenho de forma consistente e significativa. Os ganhos aproximadamente dobram ao passar da configuração de 50 códigos para a de espaço total.
• Arquitetura: Uma única política de travessia compartilhada iguala a precisão da cascata de três especialistas, mas evita uma falha estrutural crítica: a cascata sofre de estouro de janela de contexto (truncamento) em 28–32% das notas de teste na configuração de espaço total devido às grandes listas de candidatos no nível folha. A política compartilhada possui 0% de truncamento.
• Aprendizado por Reforço: O GRPO não fornece benefício sobre a continuação supervisionada quando treinado com a mesma quantidade de dados. O modelo de RL empata com seu warm-start supervisionado e fica significativamente abaixo do modelo treinado com trajetórias supervisionadas adicionais.

Significância e Alegações
O artigo argumenta que, para a predição de códigos clínicos extremos, a simplicidade e a qualidade dos dados superam a complexidade arquitetônica.

1. Decomposição com Restrição de Grafo: Reformular o problema como uma travessia restrita mitiga efetivamente o "gargalo de códigos raros" ao reduzir o espaço de busca em cada etapa e impor validade estrutural, ofereando ganhos substanciais sobre métodos de pontuação plana.
2. Supervisão sobre RL: Contrariando a tendência de usar aprendizado por reforço para tarefas complexas, o estudo descobre que simplesmente adicionar mais trajetórias supervisionadas de alta qualidade é mais eficaz do que otimizar recompensas via GRPO.
3. Política Unificada vs. Cascata: Uma única política de travessia compartilhada é preferível em relação aos modelos especialistas em cascata. Ela atinge precisão comparável enquanto elimina as limitações de escala (estouro de janela de contexto) inerentes à abordagem de cascata.
4. Recomendação Prática: Profissionais devem implantar uma política de travessia compartilhada única e investir recursos na aquisição de mais dados de anotação supervisionada, em vez de buscar arquiteturas cascata complexas ou estratégias de aprendizado por reforço.

Os autores observam limitações, incluindo o uso de um único modelo de 4B, um design de recompensa específico para RL e um orçamento de anotação relativamente pequeno (5.000 notas) em relação aos dados disponíveis, sugerindo que estudos de escalonamento adicionais são necessários para determinar pontos de saturação.