Fitting Large Nonlinear Mixed Effects Models Using Variational Expectation Maximization

1. Planteamiento del Problema

Los modelos de Efectos Mixtos No Lineales (NLME) son el estándar para analizar datos jerárquicos y longitudinales en farmacometría, Farmacología de Sistemas Cuantitativa (QSP) y ensayos clínicos. Estos modelos estiman los parámetros poblacionales ($\theta$) maximizando la verosimilitud marginal, lo cual requiere integrar los efectos aleatorios específicos del sujeto ($\eta$).

El Desafío Central:

• Intratabilidad Computacional: A medida que los modelos se vuelven más complejos (por ejemplo, incorporando redes neuronales, EDOs de alta dimensión o miles de parámetros), la integración necesaria para calcular la verosimilitud marginal se vuelve computacionalmente prohibitiva.
• Limitaciones de los Métodos Existentes: Los algoritmos tradicionales como la Estimación Condicional de Primer Orden (FOCE) y la Aproximación Estocástica EM (SAEM) dependen de la integración numérica o del muestreo de Cadenas de Markov Monte Carlo (MCMC). Estos métodos escalan mal con el número de parámetros y efectos aleatorios, a menudo fallando en converger o requiriendo tiempos excesivos para modelos con miles de parámetros (por ejemplo, modelos DeepNLME).

2. Metodología: Maximización de Expectativa Variacional (VEM)

Los autores proponen utilizar la Maximización de Expectativa Variacional (VEM) como una alternativa escalable. En lugar de muestrear desde la distribución posterior de los efectos aleatorios (como en MCEM/SAEM), VEM aproxima la posterior utilizando una distribución paramétrica (familia variacional).

Componentes Técnicos Clave:

• Marco de Inferencia Variacional (VI):
  • El objetivo es maximizar la Cota Inferior de la Evidencia (ELBO) en lugar de la verosimilitud marginal directamente.
  • La ELBO se define como: $\text{ELBO} = \mathbb{E}_{q}[\log p(y, \eta|\theta)] - \text{KL}(q(\eta) || p(\eta|y, \theta))$.
  • Maximizar la ELBO es equivalente a minimizar la divergencia de Kullback-Leibler (KL) entre la distribución variacional $q(\eta)$ y la posterior verdadera.
• Truco de Reparametrización:
  • Para permitir la optimización basada en gradientes, los autores utilizan el truco de reparametrización. Los efectos aleatorios se generan como una transformación determinista de una distribución base fija (por ejemplo, $\eta = T(\xi; \kappa)$ donde $\xi \sim \mathcal{N}(0, I)$).
  • Esto permite que los gradientes fluyan a través del proceso de muestreo, habilitando el uso de optimizadores eficientes.
• Optimización Conjunta:
  • A diferencia del EM tradicional que alterna estrictamente entre pasos E y pasos M, la implementación en Pumas realiza una optimización conjunta tanto de los parámetros poblacionales ($\theta$) como de los parámetros variacionales ($\kappa$) simultáneamente utilizando métodos basados en gradientes.
• Diferenciación Automática (AD):
  • Se utiliza AD en modo inverso para modelos con un gran número de parámetros en relación con las salidas (común en DeepNLME). Esto es crucial para calcular eficientemente los gradientes de la ELBO con respecto a miles de parámetros.
  • La implementación soporta tanto el modo directo (para modelos más pequeños) como el modo inverso de AD.
• Estrategias de Optimización:
  • VEM Determinista: Utiliza variables base pre-muestreadas y el optimizador L-BFGS (cuasi-Newton) para una convergencia estable y rápida.
  • VEM Estocástico: Utiliza remuestreo por iteración y el optimizador Adam, adecuado para el procesamiento por mini-lotes de poblaciones grandes.
• Manejo de Restricciones: Para efectos aleatorios restringidos (por ejemplo, log-normal), el método aplica transformaciones de dominio (por ejemplo, transformación logarítmica) antes de aproximar con una familia variacional gaussiana.

3. Contribuciones Clave

1. Algoritmo Escalable para NLME: El artículo introduce una variante de VEM específicamente adaptada para modelos NLME, capaz de ajustar modelos con >15,000 parámetros poblacionales y 16+ efectos aleatorios, un régimen donde los métodos tradicionales fallan.
2. Implementación en Pumas: El algoritmo está implementado en el software estadístico Pumas, aprovechando las capacidades de computación de alto rendimiento y la diferenciación automática de Julia.
3. Conexiones Teóricas: Los autores aclaran la relación entre VEM, Autoencoders Variacionales (VAE), PCA Probabilístico (PPCA) y la aproximación de Laplace. Demuestran que el método de Laplace es esencialmente una aproximación de Taylor de segundo orden del objetivo de VEM.
4. Validación Empírica: El artículo proporciona benchmarks rigurosos comparando VEM contra FOCE en modelos tanto a pequeña escala (Warfarina) como a gran escala (DeepNLME).

4. Resultados Experimentales

Los autores realizaron experimentos en dos modelos utilizando el software Pumas:

A. Modelo de Warfarina (Validación a Escala Pequeña)

• Configuración: Modelo PK-PD estándar con 31 sujetos y 464 observaciones.
• Comparación: VEM vs. FOCE.
• Resultados:
  • Precisión: VEM produjo estimaciones de parámetros y log-verosimilitudes casi idénticas a FOCE (Log Verosimilitud: -994.5 vs -995.3).
  • Velocidad: VEM fue significativamente más rápido (16.6 segundos vs. 44.4 segundos).
  • Conclusión: VEM es una alternativa correcta y eficiente para modelos farmacométricos estándar.

B. Modelo DeepNLME de Friberg (Escalabilidad a Gran Escala)

• Configuración: Un modelo artificialmente grande que incorpora redes neuronales (MLP) con 15,410 parámetros poblacionales y 16 efectos aleatorios.
• Comparación: VEM (5 iteraciones) vs. FOCE (intento de 1 iteración).
• Resultados:
  • VEM: Completó 5 iteraciones en 3092 segundos (618 segundos/iteración) utilizando AD en modo inverso.
  • FOCE: No logró completar una sola iteración dentro de las 24 horas.
  • Conclusión: VEM escala con éxito a modelos complejos y de alta dimensión donde FOCE se vuelve computacionalmente intratable.

5. Significado e Implicaciones

• Habilitación de DeepNLME y QSP: La capacidad de ajustar modelos con miles de parámetros abre la puerta a DeepNLME (integración de redes neuronales en modelos mecanísticos) y modelos complejos de Farmacología de Sistemas Cuantitativa (QSP), que anteriormente eran demasiado costosos computacionalmente para ajustarse como NLME.
• Cambio en el Paradigma de Optimización: El artículo aboga por un cambio desde el muestreo basado en MCMC (SAEM/MCEM) hacia la inferencia variacional basada en gradientes para problemas a gran escala, aprovechando la infraestructura moderna de aprendizaje profundo (AD en modo inverso, compatibilidad con GPU/TPU).
• Limitaciones Prácticas: Los autores reconocen que VEM es sensible a los hiperparámetros (número de muestras, elección de la familia variacional) y requiere pasos separados para calcular las Estimaciones Bayesianas Empíricas (EBE) y aproximaciones de log-verosimilitud post-optimización, a diferencia de FOCE que realiza esto durante la estimación.
• Direcciones Futuras: El método está posicionado para convertirse en una herramienta estándar para ajustar modelos conjuntos de biomarcadores grandes y sistemas biológicos complejos donde los métodos tradicionales chocan con un "muro" de costo computacional.

En resumen, este artículo demuestra que la Maximización de Expectativa Variacional, cuando se combina con diferenciación automática en modo inverso, proporciona un marco robusto, escalable y preciso para ajustar la próxima generación de modelos de efectos mixtos no lineales complejos y de alta dimensión en farmacometría.