Resumo Técnico: Recompensas de Distância de Conjunto para Geração de Laudos Radiológicos (SDR)

Declaração do Problema

A geração de laudos de radiografia de tórax utilizando modelos de visão-linguagem (VLMs) apresenta um desafio único para o pós-treinamento por Aprendizado por Reforço (RL), especificamente no que diz respeito ao design de funções de recompensa. Abordagens padrão de RL para tarefas de raciocínio frequentemente dependem de modelos de recompensa de processo que verificam etapas intermediárias em uma cadeia causal. No entanto, os laudos radiológicos não seguem uma estrutura de raciocínio sequencial ou causal; em vez disso, consistem em achados desordenados e ortogonais (por exemplo, "o tamanho do coração pode estar aumentado" e "o pulmão direito está limpo" são observações independentes).

Consequentemente, recompensas discretas padrão baseadas em precisão de correspondência exata são incompatíveis com este domínio. Elas induzem ruído ao penalizar laudos que contêm achados corretos, mas em uma ordem diferente ou com pequenas variações de fraseologia, falhando em capturar a realidade clínica de que o conjunto de achados importa mais do que sua sequência. Além disso, as anotações de nível de etapa necessárias para recompensas de processo são frequentemente indisponíveis, e verificar cada etapa via fontes externas é computacionalmente caro.

Metodologia

Os autores propõem as Recompensas de Distância de Conjunto (SDR), um framework que trata os laudos de radiografia de tórax como conjuntos desordenados de embeddings de sentenças, em vez de texto sequencial. A metodologia opera em duas fases: pós-treinamento (GRPO) e escalonamento em tempo de inferência.

1. Representação ao Nível de Sentença

Um laudo de radiografia de tórax $y$ é dividido em duas seções: Achados ($y_F$) e Impressão ($y_I$). Cada seção é segmentada em sentenças individuais. Um transformer de sentença pré-treinado e congelado (especificamente o all-mpnet-base-v2) transforma cada sentença $s$ em um vetor de dimensão fixa $e \in \mathbb{R}^d$. O laudo é, portanto, representado como dois conjuntos desordenados de embeddings:
$$E_F(y) = \{E_\phi(s^F_i) \mid 1 \le i \le n_F\}, \quad E_I(y) = \{E_\phi(s^I_j) \mid 1 \le j \le n_I\}$$
Esta representação é invariante à permutação, alinhando-se com a intuição clínica de que a ordem dos achados não carrecia significado diagnóstico.

2. Recompensas de Distância de Conjunto-para-Conjunto (Pós-Treinamento)

Durante a Otimização de Política Relativa de Grupo (GRPO), o modelo gera um laudo $\hat{y}$. A recompensa semântica é calculada computando a distância entre o conjunto de embeddings gerados e o conjunto de embeddings de referência (verdade fundamental/ground truth).

• Métricas de Distância: Os autores utilizam distâncias simétricas de conjunto-para-conjunto normalizadas para $[0, 1]$, especificamente a distância de Chamfer e a distância de Hausdorff, construídas sobre a distância de cosseno entre embeddings de norma unitária.
  • Chamfer: Médias os custos de vizinho mais próximo, atuando como uma recompensa de "cobertura suave".
  • Hausdorff: Toma o custo máximo de vizinho mais próximo, atuando como uma penalidade de "pior caso" para achados críticos ausentes.
• Formulação da Recompensa: A recompensa total $R(\hat{y}, y)$ é uma soma ponderada de uma recompensa de formato binária (garantindo que a saída siga um template estruturado) e a recompensa semântica contínua derivada da distância de conjunto:
  $$R(\hat{y}, y) = \lambda_{fmt} R_{fmt}(\hat{y}) + \lambda_{sem} (1 - D(E_S(\hat{y}), E_S(y)))$$
  onde $D$ é a distância de conjunto. Isso fornece um sinal contínuo e diferenciável que recompensa a correção parcial e lida com achados desordenados.

3. Escalonamento em Tempo de Inferência e Poda

As mesmas métricas de distância de conjunto são aplicadas no tempo de inferência sem atualizações de gradiente:

• Seleção Best-of-N: Para uma determinada imagem, o modelo gera $K$ laudos candidatos. O conjunto de embeddings de cada candidato é comparado contra os conjuntos de embeddings pré-computados do corpus de treinamento. O candidato com a menor distância para a distribuição de treinamento é selecionado. Isso serve como um prior livre de gradiente para saídas clinicamente plausíveis.
• Poda de Fluxo (Streaming Pruning): Para reduzir o custo computacional, o sistema computa a distância de conjunto contínua de gerações parciais durante a decodificação. Candidatos cuja distância parcial excede um limite são podados precocemente. Este loop de "podar-então-decodificar" reduz a geração de tokens em mais de 50%, mantendo a qualidade comparável à seleção completa de Best-of-N.

Principais Contribuições

1. Funções de Recompensa de Distância de Conjunto para GRPO: O artigo aborda a inviabilidade da modelagem de recompensa de processo para radiologia ao tratar os laudos como conjuntos desordenados. Usando distâncias de Chamfer e Hausdorff como recompensas contínuas, o método supera consistentemente o Ajuste Fino Supervisionado (SFT) e o GRPO de correspondência exata discreta em três backbones de VLM (Qwen3-VL-2B/4B, Gemma3-4B) e dois datasets (MIMIC-CXR, ReXGradient).
2. Seleção Guiada por Distância de Conjunto no Tempo de Teste: Os autores demonstram que as mesmas métricas de distância podem servir como critério de seleção tanto para modelos ajustados quanto para LLMs generalistas de código fechado (Mistral-Small, Gemini-2.5 Flash-Lite, GPT-4o-mini, GPT-5 mini). Esta abordagem melhora o desempenho em relação à seleção aleatória sem exigir acesso aos gradientes do modelo ou ajuste fino específico do domínio.
3. Poda Baseada em Distância On-the-Fly: O artigo introduz um método para podar candidatos de baixa pontuação durante a geração. Isso reduz o total de tokens gerados em >50% enquanto preserva a qualidade da seleção final, ofereando uma forma mais eficiente de escalonamento em tempo de teste.

Resultados Experimentais

A avaliação abrange métricas principais incluindo BERTScore, RadGraph F1 e CheXbert F1.

• Desempenho de Pós-Treinamento:
  • No dataset ReXGradient, a recompensa baseada em Chamfer alcançou o maior desempenho em todos os modelos e métricas.
  • No MIMIC-CXR, tanto as recompensas de Chamfer quanto as de Hausdorff atingiram o maior desempenho médio.
  • As melhorias relativas sobre SFT e GRPO de correspondência exata foram significativas: melhorias médias de 6,80% (BERTScore), 7,82% (RadGraph F1) e 4,45% (CheXbert F1).
• Seleção em Tempo de Inferência:
  • A aplicação da seleção por distância de conjunto em LLMs de código fechado rendeu ganhos substanciais sobre a seleção aleatória. Por exemplo, no BERTScore, as melhorias variaram de +11,2% a +18,7%, dependendo do modelo e do template de prompt.
  • O método melhorou com sucesso o desempenho de modelos generalistas (ex: GPT-4o mini, Gemini) que não foram ajustados finamente para o domínio médico específico.
• Eficiência de Poda:
  • A estratégia de poda economizou de 42,1% a 60,1% dos tokens de geração.
  • As gerações podadas mantiveram qualidade comparável à seleção completa de Best-of-N, com melhorias relativas médias sobre a seleção aleatória de +12,7% (BS-F1), +17,1% (RG-F1) e +6,2% (CXB-F1).

Significância e Alegações

O artigo afirma estabelecer as recompensas de distância de conjunto como um sinal unificado tanto para o pós-treinamento quanto para o escalonamento em tempo de teste em laudos de radiografia de tórax.

• Sinal Unificado: A mesma formulação matemática (distância de conjunto-para-conjunto) efetivamente impulsiona o processo de pós-treinamento por RL e serve como um critério de seleção robusto e livre de gradiente no tempo de inferência.
• Adaptação de Domínio: Ao reconhecer a natureza desordenada dos achados radiológicos, o método evita o ruído introduzido por restrições sequenciais, tornando-se especificamente adequado para a geração de laudos médicos, onde os achados são frequentemente independentes.
• Eficiência: O trabalho demonstra que a geração de laudos de alta qualidade não requer necessariamente a geração de todos os tokens para todos os candidatos; a poda precoce baseada na distância de conjunto oferece um caminho viável para reduzir os custos de inferência.
• Generalização: A capacidade de aplicar essas regras de seleção a LLMs de código fechado sugere um caminho para melhorar os outputs clínicos de IA sem a necessidade de ajuste fino específico do domínio ou acesso aos pesos do modelo.

Os autores concluem que sua abordagem fornece um sinal de recompensa contínuo e prático que se alinha com a estrutura clínica dos laudos radiológicos, superando mecanismos tradicionais de recompensa discretos ou sequenciais.