Resumo Técnico: CLIN-JEPA

Declaração do Problema

A modelagem de Registros Eletrônicos de Saúde (RES) na Unidade de Terapia Intensiva (UTI) enfrenta uma limitação fundamental: as abordagens existentes tratam os dados do paciente ou como sequências de texto estáticas (modelos autorregressivos de tokens) que falham em modelar estados fisiológicos contínuos, ou como modelos fundacionais por tarefa que exigem engenharia de características específica e ajuste fino para cada tarefa downstream. Nenhuma das abordagens fornece uma representação latente unificada capaz de tanto prever trajetórias de pacientes quanto apoiar diversas tarefas de previsão de risco sem adaptação por tarefa.

Embora as Arquiteturas Preditivas de Incorporação Conjunta (JEPA) tenham permitido com sucesso o planejamento no espaço latente em robótica e visão, adaptar esse paradigma para RES clínicas permanece um desafio aberto. Projetos anteriores de JEPA sofrem de lacunas arquitetônicas específicas quando aplicados à simulação clínica:

• I-JEPA e V-JEPA: Descartam o preditor após o pré-treinamento, tornando impossível a simulação de trajetórias a jusante.
• V-JEPA 2-AC: Mantém o preditor, mas o treina sobre um codificador congelado. Consequentemente, o codificador não está ciente do "sinal de simulação" (rollout signal) necessário para a inferência, levando a uma desconexão entre o espaço de representação e a dinâmica de previsão.

O co-treinamento ingênuo de um codificador e um preditor sob um objetivo JEPA compartilhado é instável, propenso ao colapso de representação (onde o codificador degenera para saídas constantes) e à deriva online/alvo (onde o preditor diverge à medida que acumula erros em um espaço latente em movimento).

Metodologia: Framework CLIN-JEPA

Os autores propõem o CLIN-JEPA, um framework de co-treinamento multifase projetado para treinar estávelmente uma única base que simultaneamente prevê trajetórias de pacientes e serve a tarefas de previsão de risco downstream.

1. Representação de Dados e Codificador

• Formato de Entrada: As internações de pacientes são modeladas como uma sequência de pares estado-ação em intervalos de 1 hora (até 72 horas), precedidos por dados demográficos estáticos. Diferentemente da vetorização numérica, o modelo utiliza texto estruturado em linguagem natural para representar estados (sinais vitais, exames laboratoriais, escores de gravidade) e ações (intervenções, doses). Isso preserva a legibilidade humana, lida naturalmente com valores ausentes (por omissão) e evita engenharia de características específica de modalidade.
• Arquitetura do Codificador: O sistema utiliza um modelo de linguagem Qwen3-8B como codificador base. Os pesos base são congelados, e adaptadores LoRA leves são treinados para mapear entradas de texto (demografia $d$, estado $s_t$, ação $a_t$) em um espaço latente de 4096 dimensões.

2. Preditor de Trajetória Latente

• Arquitetura: Um codificador Transformer de 92 milhões de parâmetros opera inteiramente dentro do espaço latente de 4096 dimensões.
• Mecanismo: Ele consome uma sequência intercalada de dados demográficos, estados e ações. Usando atenção causal em blocos (bidirecional dentro de um passo de tempo, estritamente causal entre passos de tempo), ele prevê o próximo estado latente $\hat{z}_{t+1}$.
• Inferência: O preditor é mantido no momento da inferência. Ele realiza simulação autorregressiva (autoregressive rollout), alimentando suas próprias previsões contínuas de incorporação de volta como entradas para etapas subsequentes para simular trajetórias de pacientes ao longo de um horizonte de $H$ horas.

3. Currículo de Co-treinamento em Cinco Fases

Para abordar a instabilidade do co-treinamento ingênuo, os autores introduzem um cronograma de pré-treinamento em cinco fases:

1. Aquecimento (Warmup): O preditor é treinado sozinho contra um codificador Supervisionado Ajustado (SFT) congelado. Isso impede que o preditor "frio" arraste imediatamente o codificador em direção ao colapso de representação.
2. Co-treinamento: Os adaptadores LoRA do codificador são desbloqueados e treinados conjuntamente com o preditor sob um objetivo JEPA compartilhado. A simulação forçada pelo professor (teacher-forced rollout) é usada para garantir que o codificador aprenda representações dinamicamente fundamentadas para o preditor.
3. Alinhamento: O codificador online é re-congelado, e o codificador alvo com Média Móvel Exponencial (EMA) acompanha suavemente os parâmetros do codificador online.
4. Sincronização Rígida (Hard Sync): Uma cópia instantânea de parâmetros alinha os codificadores alvo e online, eliminando discrepâncias residuais.
5. Finalização: O preditor é treinado no codificador estabilizado usando simulação autorregressiva nativa (mimetizando condições de inferência) para fazer o preditor convergir completamente nas representações refinadas.

O objetivo de treinamento minimiza a distância $\ell_1$ entre a incorporação do próximo estado prevista e a incorporação do próximo estado real pelo codificador alvo, combinada com uma perda de simulação de curto horizonte para encorajar estabilidade multi-etapa.

Contribuições Principais

1. Currículo de Co-treinamento Estável: Um cronograma inovador de cinco fases que co-treina com sucesso um codificador baseado em modelo de linguagem de grande escala e um preditor de trajetória latente, prevenindo colapso de representação e deriva.
2. Representação Unificada Baseada em Texto: Um framework que consome texto clínico bruto sem imputação, normalização ou características criadas manualmente, servindo a todas as tarefas downstream a partir de um único conjunto de incorporações latentes.
3. Primeiro JEPA Clínico com Preditor Retido: Diferentemente de trabalhos anteriores que descartam o preditor, o CLIN-JEPA o mantém para permitir simulação estável e autorregressiva do estado do paciente.

Resultados Experimentais

O framework foi avaliado no conjunto de dados MIMIC-IV ICU (84.497 internações) em três eixos:

1. Estabilidade de Treinamento e Convergência

• Deriva de Simulação: Ao longo de um horizonte autorregressivo de 48 horas, o erro de previsão do CLIN-JEPA convergiu (diminuiu 15,7% em relação à linha de base de 1 hora). Em contraste, linhas de base e ablações (removendo aquecimento ou alinhamento) divergiram catastróficamente, com aumentos de deriva variando de +3% a +4951%.
• Colapso de Representação: Apenas o currículo completo do CLIN-JEPA manteve um desvio padrão saudável nas saídas do codificador ($z_{std}$), evitando o colapso observado em variantes que careciam da fase de aquecimento.

2. Diagnóstico de Geometria Latente

• Discriminação Clínica: O codificador aprendeu uma geometria onde pacientes em deterioração (aumento do escore SOFA $\ge$ 3) deslocaram-se 4,83× mais no espaço latente do que pacientes estáveis. Codificadores de linha de base mostraram taxas de separação significativamente menores (2,62× e 1,03×).
• Dinâmica Temporal: O poder discriminativo do CLIN-JEPA cresceu com o horizonte de observação, atingindo o pico na hora 66, refletindo uma internalização da escala de tempo lenta e integrativa da progressão na UTI.

3. Desempenho Multi-tarefa Downstream

Avaliado em dois benchmarks (ICareFM Early Event Prediction e 8 tarefas de risco binárias) sem ajuste fino por tarefa:

• ICareFM EEP: Alcançou uma AUROC média de 0,851, superando a linha de base mais forte (LightGBM, 0,827) e ablações do codificador.
• Tarefas de Risco Binárias: Alcançou uma AUROC média de 0,883, superando a linha de base LSTM (0,865).
• Valor do Preditor: A inclusão da simulação autorregressiva (previsões futuras) melhorou consistentemente o desempenho em relação a entradas apenas históricas, particularmente para tarefas que exigem a composição de informações de trajetória heterogêneas (por exemplo, lesão renal, sepse, descompensação).

Significância

O artigo afirma que o CLIN-JEPA é o primeiro framework a fechar com sucesso a lacuna entre o pré-treinamento JEPA e a inferência clínica ao reter o preditor. Ele demonstra que uma única base clínica pode entregar simulação autorregressiva estável, geometria latente clinicamente discriminativa e desempenho multi-tarefa competitivo simultaneamente. O trabalho estabelece que o co-treinamento do codificador e do preditor sob um objetivo compartilhado, guiado por um currículo específico, é necessário para fundamentar representações para simulação de longo horizonte em sistemas dinâmicos de alta dimensão como a UTI.