Fitting Large Nonlinear Mixed Effects Models Using Variational Expectation Maximization

Mohamed Tarek 와 Pedro Afonso 가 작성한 논문 "Variational Expectation Maximization 을 이용한 대규모 비선형 혼합 효과 모델 피팅"에 대한 상세한 기술적 요약입니다.

1. 문제 제기

비선형 혼합 효과 (NLME) 모델은 약동약력학, 정량적 시스템 약리학 (QSP), 그리고 임상 시험에서 계층적 및 종단 데이터를 분석하는 표준 방법입니다. 이러한 모델은 개체별 무작위 효과 ($\eta$) 를 적분하여 **주변 우도 (marginal likelihood)**를 최대화함으로써 모수 ($\theta$) 를 추정합니다.

핵심 과제:

• 계산적 비실용성: 모델이 더 복잡해짐에 따라 (예: 신경망 통합, 고차원 ODE 포함, 또는 수천 개의 모수), 주변 우도를 계산하는 데 필요한 적분은 계산적으로 불가능해집니다.
• 기존 방법의 한계: 1 차 조건부 추정 (FOCE) 및 확률적 근사 EM(SAEM) 과 같은 전통적인 알고리즘은 수치 적분 또는 마코프 연쇄 몬테 카를로 (MCMC) 샘플링에 의존합니다. 이러한 방법들은 모수와 무작위 효과의 수에 따라 확장성이 낮아, 수천 개의 모수를 가진 모델 (예: DeepNLME 모델) 의 경우 수렴에 실패하거나 과도한 시간이 소요되는 경우가 많습니다.

2. 방법론: 변분 기대 최대화 (VEM)

저자들은 확장 가능한 대안으로 **변분 기대 최대화 (VEM)**를 제안합니다. MCEM/SAEM 과 같이 무작위 효과의 사후 분포에서 샘플링하는 대신, VEM 은 매개변수 분포 (변분족) 를 사용하여 사후 분포를 근사합니다.

주요 기술 구성 요소:

• 변분 추론 (VI) 프레임워크:
  • 목표는 주변 우도를 직접 최대화하는 것이 아니라 **증거 하한 (ELBO)**을 최대화하는 것입니다.
  • ELBO 는 다음과 같이 정의됩니다: $\text{ELBO} = \mathbb{E}_{q}[\log p(y, \eta|\theta)] - \text{KL}(q(\eta) || p(\eta|y, \theta))$.
  • ELBO 를 최대화하는 것은 변분 분포 $q(\eta)$와 실제 사후 분포 간의 Kullback-Leibler (KL) 발산을 최소화하는 것과 동일합니다.
• 재매개변수화 기법 (Reparameterization Trick):
  • 경사 기반 최적화를 가능하게 하기 위해 저자들은 재매개변수화 기법을 사용합니다. 무작위 효과는 고정된 기본 분포의 결정론적 변환으로 생성됩니다 (예: $\eta = T(\xi; \kappa)$, 여기서 $\xi \sim \mathcal{N}(0, I)$).
  • 이를 통해 샘플링 과정을 통해 기울기가 흐를 수 있게 되어 효율적인 최적화 도구를 사용할 수 있습니다.
• 공동 최적화 (Joint Optimization):
  • 전통적인 EM 이 E-단계와 M-단계를 엄격하게 번갈아 수행하는 것과 달리, Pumas 구현체는 경사 기반 방법을 사용하여 모수 ($\theta$) 와 변분 모수 ($\kappa$) 를 동시에 공동 최적화합니다.
• 자동 미분 (AD):
  • 역방향 모드 AD는 출력에 비해 모수가 많은 모델 (DeepNLME 에서 일반적) 에 활용됩니다. 이는 수천 개의 모수에 대한 ELBO 의 기울기를 효율적으로 계산하는 데 필수적입니다.
  • 구현체는 작은 모델용 순방향 모드와 역방향 모드 AD 를 모두 지원합니다.
• 최적화 전략:
  • 결정론적 VEM: 사전 샘플링된 기본 변수와 L-BFGS 최적화 도구 (준뉴턴) 를 사용하여 안정적이고 빠른 수렴을 달성합니다.
  • 확률적 VEM: 반복마다 리샘플링을 사용하고 Adam 최적화 도구를 사용하며, 대규모 인구의 미니배치 처리에 적합합니다.
• 제약 조건 처리: 제약이 있는 무작위 효과 (예: 로그정규분포) 의 경우, 가우시안 변분족으로 근사하기 전에 도메인 변환 (예: 로그 변환) 을 적용합니다.

3. 주요 기여

1. 확장 가능한 NLME 알고리즘: 이 논문은 15,000 개 이상의 모수와 16 개 이상의 무작위 효과를 가진 모델을 피팅할 수 있는 NLME 모델에 특화된 VEM 변형을 소개합니다. 이는 전통적인 방법이 실패하는 영역입니다.
2. Pumas 구현: 이 알고리즘은 Julia 의 고성능 컴퓨팅 기능과 자동 미분을 활용하여 Pumas 통계 소프트웨어에 구현되었습니다.
3. 이론적 연결: 저자들은 VEM, 변분 오토인코더 (VAE), 확률적 PCA(PPCA), 그리고 라플라스 근사 사이의 관계를 명확히 합니다. 그들은 라플라스 방법이 본질적으로 VEM 목적 함수의 2 차 테일러 근사임을 입증합니다.
4. 실증적 검증: 이 논문은 소규모 (Warfarin) 및 대규모 (DeepNLME) 모델 모두에서 FOCE 와 VEM 을 비교하는 엄격한 벤치마크를 제공합니다.

4. 실험 결과

저자들은 Pumas 소프트웨어를 사용하여 두 가지 모델에 대해 실험을 수행했습니다.

A. Warfarin 모델 (소규모 검증)

• 설정: 31 명의 대상자와 464 개의 관측치를 가진 표준 PK-PD 모델.
• 비교: VEM 대 FOCE.
• 결과:
  • 정확도: VEM 은 FOCE 와 거의 동일한 모수 추정치와 로그 우도를 생성했습니다 (로그 우도: -994.5 대 -995.3).
  • 속도: VEM 은 훨씬 빨랐습니다 (16.6 초 대 44.4 초).
  • 결론: VEM 은 표준 약동약력학 모델에 대한 정확하고 효율적인 대안입니다.

B. DeepNLME Friberg 모델 (대규모 확장성)

• 설정: 15,410 개의 모수와 16 개의 무작위 효과를 포함하는 신경망 (MLP) 을 통합한 인위적으로 큰 모델.
• 비교: VEM(5 회 반복) 대 FOCE(1 회 반복 시도).
• 결과:
  • VEM: 역방향 모드 AD 를 사용하여 약 3092 초 (~618 초/반복) 에 5 회 반복을 완료했습니다.
  • FOCE: 24 시간 이내에 단일 반복도 완료하지 못했습니다.
  • 결론: VEM 은 FOCE 가 계산적으로 비실용적이 되는 복잡하고 고차원적인 모델까지 성공적으로 확장됩니다.

5. 중요성 및 함의

• DeepNLME 및 QSP 가능화: 수천 개의 모수를 가진 모델을 피팅할 수 있는 능력은 DeepNLME(기계적 모델에 신경망 통합) 및 복잡한 **정량적 시스템 약리학 (QSP)**모델의 문을 엽니다. 이는 이전에는 NLME 로 피팅하기에는 계산 비용이 너무 많이 들어 불가능했던 영역입니다.
• 최적화 패러다임의 전환: 이 논문은 대규모 문제에 대해 MCMC 기반 샘플링 (SAEM/MCEM) 에서 현대적인 딥러닝 인프라 (역방향 모드 AD, GPU/TPU 호환성) 를 활용하는 경사 기반 변분 추론으로의 전환을 옹호합니다.
• 실용적 한계: 저자들은 VEM 이 하이퍼파라미터 (샘플 수, 변분족 선택) 에 민감하며, FOCE 는 추정 중에 수행하는 것과 달리 최적화 후 경험적 베이즈 추정치 (EBE) 와 로그 우도 근사를 별도로 계산해야 한다는 점을 인정합니다.
• 향후 방향: 이 방법은 전통적인 방법이 계산 비용의 "벽"에 부딪히는 대규모 결합 바이오마커 모델 및 복잡한 생물학적 시스템을 피팅하는 표준 도구가 될 준비가 되어 있습니다.

요약하자면, 이 논문은 역방향 모드 자동 미분과 결합된 변분 기대 최대화가 약동약력학에서 차세대 복잡하고 고차원적인 비선형 혼합 효과 모델을 피팅하기 위한 강력하고 확장 가능하며 정확한 프레임워크를 제공함을 보여줍니다.