Fitting Large Nonlinear Mixed Effects Models Using Variational Expectation Maximization

1. Formulação do Problema

Modelos de Efeitos Mistos Não Lineares (NLME) são o padrão para análise de dados hierárquicos e longitudinais em farmacometria, Farmacologia de Sistemas Quantitativa (QSP) e ensaios clínicos. Esses modelos estimam parâmetros populacionais ($\theta$) maximizando a verossimilhança marginal, o que requer a integração dos efeitos aleatórios específicos do sujeito ($\eta$).

O Desafio Central:

• Intratabilidade Computacional: À medida que os modelos se tornam mais complexos (por exemplo, incorporando redes neurais, EDOs de alta dimensão ou milhares de parâmetros), a integração necessária para calcular a verossimilhança marginal torna-se computacionalmente proibitiva.
• Limitações dos Métodos Existentes: Algoritmos tradicionais como Estimação Condicional de Primeira Ordem (FOCE) e Aproximação Estocástica EM (SAEM) dependem de integração numérica ou amostragem de Cadeias de Markov Monte Carlo (MCMC). Esses métodos escalam mal com o número de parâmetros e efeitos aleatórios, frequentemente falhando em convergir ou levando tempo excessivo para modelos com milhares de parâmetros (por exemplo, modelos DeepNLME).

2. Metodologia: Maximização Expectativa Variacional (VEM)

Os autores propõem o uso da Maximização Expectativa Variacional (VEM) como uma alternativa escalável. Em vez de amostrar a partir da distribuição posterior dos efeitos aleatórios (como em MCEM/SAEM), o VEM aproxima a posterior usando uma distribuição paramétrica (família variacional).

Componentes Técnicos Chave:

• Estrutura de Inferência Variacional (VI):
  • O objetivo é maximizar o Limite Inferior da Evidência (ELBO) em vez de maximizar diretamente a verossimilhança marginal.
  • O ELBO é definido como: $\text{ELBO} = \mathbb{E}_{q}[\log p(y, \eta|\theta)] - \text{KL}(q(\eta) || p(\eta|y, \theta))$.
  • Maximizar o ELBO é equivalente a minimizar a divergência de Kullback-Leibler (KL) entre a distribuição variacional $q(\eta)$ e a posterior verdadeira.
• Truque de Reparametrização:
  • Para permitir otimização baseada em gradiente, os autores utilizam o truque de reparametrização. Os efeitos aleatórios são gerados como uma transformação determinística de uma distribuição base fixa (por exemplo, $\eta = T(\xi; \kappa)$ onde $\xi \sim \mathcal{N}(0, I)$).
  • Isso permite que os gradientes fluam através do processo de amostragem, habilitando o uso de otimizadores eficientes.
• Otimização Conjunta:
  • Diferentemente do EM tradicional, que alterna estritamente entre passos E e passos M, a implementação no Pumas realiza otimização conjunta tanto dos parâmetros populacionais ($\theta$) quanto dos parâmetros variacionais ($\kappa$) simultaneamente, usando métodos baseados em gradiente.
• Diferenciação Automática (AD):
  • A AD em modo reverso é utilizada para modelos com um grande número de parâmetros em relação às saídas (comum em DeepNLME). Isso é crucial para calcular eficientemente os gradientes do ELBO em relação a milhares de parâmetros.
  • A implementação suporta tanto o modo direto (para modelos menores) quanto o modo reverso de AD.
• Estratégias de Otimização:
  • VEM Determinístico: Usa variáveis base pré-amostradas e o otimizador L-BFGS (quase-Newton) para convergência estável e rápida.
  • VEM Estocástico: Usa reamostragem por iteração e o otimizador Adam, adequado para processamento em mini-lotes de grandes populações.
• Tratamento de Restrições: Para efeitos aleatórios restritos (por exemplo, log-normal), o método aplica transformações de domínio (por exemplo, transformação logarítmica) antes de aproximar com uma família variacional Gaussiana.

3. Contribuições Principais

1. Algoritmo Escalável para NLME: O artigo introduz uma variante do VEM especificamente adaptada para modelos NLME, capaz de ajustar modelos com >15.000 parâmetros populacionais e 16+ efeitos aleatórios, um regime onde os métodos tradicionais falham.
2. Implementação no Pumas: O algoritmo foi implementado no software estatístico Pumas, aproveitando as capacidades de computação de alto desempenho e diferenciação automática da linguagem Julia.
3. Conexões Teóricas: Os autores esclarecem a relação entre VEM, Autoencoders Variacionais (VAEs), PCA Probabilístico (PPCA) e a aproximação de Laplace. Eles demonstram que o método de Laplace é essencialmente uma aproximação de Taylor de segunda ordem do objetivo do VEM.
4. Validação Empírica: O artigo fornece benchmarks rigorosos comparando o VEM com o FOCE em modelos de pequena escala (Warfarina) e de grande escala (DeepNLME).

4. Resultados Experimentais

Os autores conduziram experimentos em dois modelos usando o software Pumas:

A. Modelo de Warfarina (Validação em Pequena Escala)

• Configuração: Modelo PK-PD padrão com 31 sujeitos e 464 observações.
• Comparação: VEM vs. FOCE.
• Resultados:
  • Precisão: O VEM produziu estimativas de parâmetros e log-verossimilhanças quase idênticas ao FOCE (Log Verossimilhança: -994,5 vs -995,3).
  • Velocidade: O VEM foi significativamente mais rápido (16,6 segundos vs. 44,4 segundos).
  • Conclusão: O VEM é uma alternativa correta e eficiente para modelos farmacométricos padrão.

B. Modelo Friberg DeepNLME (Escalabilidade em Grande Escala)

• Configuração: Um modelo artificialmente grande incorporando redes neurais (MLPs) com 15.410 parâmetros populacionais e 16 efeitos aleatórios.
• Comparação: VEM (5 iterações) vs. FOCE (tentativa de 1 iteração).
• Resultados:
  • VEM: Completou 5 iterações em 3092 segundos (618 segundos/iteração) usando AD em modo reverso.
  • FOCE: Falhou em completar uma única iteração dentro de 24 horas.
  • Conclusão: O VEM escala com sucesso para modelos complexos e de alta dimensão onde o FOCE se torna intratável computacionalmente.

5. Significado e Implicações

• Habilitando DeepNLME e QSP: A capacidade de ajustar modelos com milhares de parâmetros abre portas para DeepNLME (integração de redes neurais em modelos mecanísticos) e modelos complexos de Farmacologia de Sistemas Quantitativa (QSP), que anteriormente eram computacionalmente caros demais para serem ajustados como NLMEs.
• Mudança no Paradigma de Otimização: O artigo defende uma mudança da amostragem baseada em MCMC (SAEM/MCEM) para inferência variacional baseada em gradiente para problemas de grande escala, aproveitando infraestrutura moderna de aprendizado profundo (AD em modo reverso, compatibilidade com GPU/TPU).
• Limitações Práticas: Os autores reconhecem que o VEM é sensível a hiperparâmetros (número de amostras, escolha da família variacional) e requer etapas separadas para calcular Estimativas de Bayes Empírico (EBEs) e aproximações de log-verossimilhança pós-otimização, ao contrário do FOCE, que faz isso durante a estimação.
• Direções Futuras: O método está posicionado para se tornar uma ferramenta padrão para ajustar grandes modelos de biomarcadores conjuntos e sistemas biológicos complexos onde os métodos tradicionais atingem um "muro" de custo computacional.

Em resumo, este artigo demonstra que a Maximização Expectativa Variacional, quando combinada com diferenciação automática em modo reverso, fornece um framework robusto, escalável e preciso para ajustar a próxima geração de modelos de efeitos mistos não lineares complexos e de alta dimensão em farmacometria.