Uncertainty of Treatment Effect

! 개인 연구 미팅 준비하면서 정리한 내용으로, 의식의 흐름 전개가 있을 수 있음 ! 

Potential outcome은 Rubin의 Potential Outcome Framework에서 나온 개념이다. 이 개념은 한 사람에게 가능한 모든 treatment 시나리오에 따른 결과를 말한다. 가령 어떤 치료가 있을 때 Y(1)은 치료를 받았을 때의 potential outcome이고, Y(0)은 치료를 받지 않았을 때의 potential outcome이다.

Observational data에서는 두 potential outcome 중 단 하나만 관측이 가능하다. 이 경우 potential outcome 중 실제로 관측되지 않은 것을 counterfactual outcome이라 한다. 가령 어떤 환자가 실제로 치료를 받았다면, 우리는 Y(1)은 알지만, Y(0)은 관측할 수 없다. 이떄 Y(0)은 counterfactual outcome이 된다.   

전체 인구에 대해 치료의 평균 효과는 다음처럼 정의된다. 

- Average Treatment Effect (ATE): E[Y(1) - Y(0)]
- Conditional Average Treatment Effect (CATE): E[Y(1) - Y(0) | X= x]

산업에서 decision making을 하기 위해 이렇게 averaged causal quantities를 사용하게 되면 오해가 생길 수 있고, 특정 분야 (의료)에서는 위험할 수도 있다. 

ATE나 CATE는 모두 기대값만을 알려준다. 두 환자 모두 CATE가 2라고 하더라도, 실제 분포는 다음처럼 완전히 다를 수 있다.
- A 환자 군: Y(1) ~ N(5, $0.1^2$), Y(0) ~ N(3, $0.1^2$) $\rightarrow$ CATE = 2, 효과 매우 확실
- B 환자 군: Y(1) ~ N(6, $4^2$), Y(0) ~ N(4, $4^2$) $\rightarrow$ CATE = 2, 효과 불확실

이렇게 어떤 치료가 평균적으로는 효과가 있다고 해도, 어떤 사람에게는 전혀 안듣거나, 오히려 해로울 수도 있다. ( 평균이라는 것이 예외, 변동성을 고려할 수 없기 때문에 특정 환자에게는 치료효과가 없어도 평균을 내게 되면 커보일 수 있다. 꼬리분포) 그래서 결과의 분포를 알아야 한다. 이를 통해 어떤 결과가 얼마나 자주 일어날지까지 알 수 있다. 특히 의료에서는 "이 치료가 성공할 확률은 70%" 처럼 성공/실패 확률을 알려줄 수 있어야 한다.

따라서 분포 전체를 추정해야만 확률, 신뢰구간에 기반한 의사결정이 가능하다. 
예시 1 ) 치료를 했을 때 종양의 크기가 $\tau$ 이하가 될 확률: P(Y(1) $<= \tau$ | X = x)을 구해서 임계 확률을 넘는 경우 치료를 권장 
예시 2 ) 치료의 부작용이 너무 클 확률이 높은 경우 P(Y(1) < Y(0) ) 기존 CATE가 양수라도, 이 확률이 크면 치료를 보류
예시 3 ) 치료 효과 Y(1) - Y(0)의 95% 신뢰 구간이 [1.5, 4.5]로 긍정적인 경우 치료,, 반면 [-2, 6]인 경우 불확실하니까 보류

다음은 인과 추론에서 평균 효과 (ATE)만 보고 판단하면 위험하다는 점을 보여주는 시뮬레이션 기반 예시이다.

데이터는 구조적 인과 모델 (오른쪽 식들)을 통해 생성됨

 

- P(Y[0]), P(Y[1]): 전체 인구에 대해 치료를 안받았을/받았을 경우 나타났을 결과 분포 (관측 데이터를 통해 추정 불가능, counterfactual outcome을 포함하여 추정)

- P(Y=y|A=0), P(Y=y|A=1) : factual outcome 분포

-  P(Y[1]=y|A=0), P(Y[0]=y|A=1) : counterfactual outcome 분포

그림 (b), (c)에서 두 점선 분포의 mixture인 실선 분포가 전체 인구에 대해 치료를 안받았을/받았을 경우 나타났을 결과에 대한 분포 (a)임

설명을 보면, 두 가지 치료에 대한 potential outcome의 평균과 분산이 같다. (평균 = 4.77, 분산 = 4.06) 그러나, 그래프를 보면 (a) intervention distributions에서 평균은 같지만 분포의 모양 (꼬리, 다봉성 등)이 달랐다. 

이런 경우 outcome(가령 혹의 크기)이 5보다 작을 때 성공적이라는 기준이 있을때, 치료/비치료에 대해 outcome이 5보다 작을 확률을 추정 한다고 했을 때 다음과 같은 결론을 얻을 수 있는데,
- P(Y[1] < 5) = 0.63
- P(Y[0] < 5) = 0.51
즉, 두 분포의 평균과 분산이 같더라도, 분포가 다르기 때문에 확률에 차이가 생겨 치료를 했을 때의 결과가 더 좋을 확률이 높다는 결론을 낼 수 있다.

분포를 고려하면 치료의 성공 확률 차이를 알 수 있다. 이는 임상 의사결정에 필수적이다.  

potential outcome의 aleatoric uncertainty
aleatoric uncertainty란 내재적 불확실성 또는 우연에 의한 불확실성을 뜻한다. 이것은 측정 오차, 데이터 자체의 자연스러운 변동성 등 본질적으로 제거할 수 없는 불확실성을 의미한다.  -> 같은 치료를 해도 환자 A는 90%의 성공 가능성이 있지만 B는 60%의 성공 가능성이 있는 경우 (Y(a)가 확률적인 결과이다)

결과 Y가 어떤 확률 분포를 따르는지 추정하는 문제인 Interventional Density Estimation 방법론 중에 Interventional Normalizing Flows (INF)가  있다. 

INF는 정규화된 밀도 함수를 생성하고, 직접 샘플링도 가능하며 대규모, 고차원 데이터에도 잘 작동하는 fully parametric 방법이다. Kennedy et al. (2023)의 이론을 기반으로 복잡한 조건 (편향 보정 moment condition)을 최적화 가능한 형태로 변환하고, 이를 딥러닝 기반 모델로 실제 학습 가능하게 만든다. 

다음은 Normalizing Flow에 대한 배경 지식을 시작으로 INF 논문에 기반한 내용을 정리하였다.

 Normalizing Flow

데이터 $x$가 존재할 때, 그 데이터의 확률 분포 $p_X(x)$를 알고 싶은 상황. 하지만 현실에서는 $p_X(x)$는 복잡하고 계산이 어렵기 때문에 간단한 분포 $p_Z(z)$를 잘 정의해두고, 그것을 적절한 함수 $f$를 통해 변환해서 $x$로 만든다. 

$$x = f(z), z \sim  P_Z(z)$$ 

이때 $f$는 가역 함수 (invertible function)이고, 미분 가능해야 한다.

확률 변수 $z$를 함수 $x = f(z)$를 통해 변환했을 때 결과 변수 $x$의 확률 분포 $p_X(x)$는 어떤 분포를 가질지 계산하려면 확률의 보존 원칙을 적용해야 한다. (그러니까 변환을 하더라도, 확률(질량)은 보존되어야 한다.)

이 이야기를 조금 쉽게 풀어보자면

1차원 예시)

z가 구간 [0, 1]에서 Uniform을 따른다고 하면 $z \sim Uniform(0, 1)$ 구간 [0, 1]에서 확률은 $p_Z(z) = 1$ 일 것이다.

만약 여기서 $x = 2z$ 변환을 한다고하면 $p_X(x)$는  z의 공간인 [0,1] 보다 2배 더 긴 [0, 2]에서 존재한다. 그래도 이 구간에서 확률의 합은 항상 1로 동일해야한다.

그러므로 $\int_{0}^{2}p_X(x)dx = \int_{0}^{1}p_Z(z)dz = 1$ 이다. 이걸 아주 작은 변화에 대한 확률로 다시 표현하면 $p_X(x)dx = p_Z(z)dz$다.

위 식에서 $p_X(x)$를 좌변에 남기고 넘기면  $p_X(x) = p_Z(z) \cdot \frac{dz}{dx}$이고, 이는 즉 $p_X(x) = p_Z(z) \cdot \frac{1}{2}$를 의미한다. 따라서 $x$의 공간(구간)이 2배 길어졌지만, x의 확률 밀도($p_X(x)$)는 z의 확률 밀도($p_Z(z)$) 보다 2배 작아졌다.

2차원 예시)

다변량으로 확장해서 2개의 입력, 2개의 출력으로 변환되는 상황을 가정한다. $[x_1, x_2] = f([z_1, z_2])$

$x_1 = 2z_1$, $x_2 = 2z_2$ 이렇게 변환되는 상황에서 $p_X(x)$를 구하고자 한다.
- 입력 $z=[z_1, z_2]$
- 변환 함수 :  $\begin{bmatrix} x_1 \\ x_2\end{bmatrix} = \begin{bmatrix} 2 & 0 \\ 0 & 2 \\ \end{bmatrix}\begin{bmatrix} z_1 \\ z_2\end{bmatrix} = \begin{bmatrix} 2z_1 \\ 2z_2\end{bmatrix}$

구간 [0, 1]에서 $z_1, z_2$가 존재한다고 하면, 면적이 1인 정사각형 영역에서 확률 밀도는 p_Z(z)로 표현할 수 있다.  
이를 변환 하면 $x_1, x_2$는 구간 [0, 2]에서 존재하게 되고, 변환된 공간의 면적(4)에서 확률 밀도 p_X(x)가 정의된다.

같은 양의 확률 질량이 존재해야 하므로 x의 밀도 p_X(x)는 p_Z(z)의 1/4 배여야 한다.

이러한 관계를 1차원 예시와 마찬가지로, $p_X(x) = p_Z(z) \cdot \frac{dz}{dx}$로 나타낼 수 있는데, 다변량에서는 $\frac{dz}{dx} $를 여러 입출력에 대해 미분량을 구해야 하므로, 편미분을 사용하는 jacobian 행렬의 행렬식을 이용한다.

$$ p_X(x) = p_Z(z) \cdot |det J|^{-1} $$

$$J_f(\mathbf{z}) = \frac{\partial f(\mathbf{z})}{\partial \mathbf{z}^\top} =
\begin{bmatrix}
\frac{\partial f_1}{\partial z_1} & \frac{\partial f_1}{\partial z_2} & \cdots & \frac{\partial f_1}{\partial z_n} \\
\frac{\partial f_2}{\partial z_1} & \frac{\partial f_2}{\partial z_2} & \cdots & \frac{\partial f_2}{\partial z_n} \\
\vdots & \vdots & \ddots & \vdots \\
\frac{\partial f_m}{\partial z_1} & \frac{\partial f_m}{\partial z_2} & \cdots & \frac{\partial f_m}{\partial z_n}
\end{bmatrix}$$

$\frac{dz}{dx} $ 는 x에 대한 미분이기 때문에 $ J_f(\mathbf{z}) $의 역수를 취해 행렬식을 구한다. 

 $$J = \frac{\partial x}{\partial z} = \begin{bmatrix}
\frac{\partial x_1}{\partial z_1} & 0 \\
0 & \frac{\partial x_2}{\partial z_2} \\
\end{bmatrix} = \begin{bmatrix}
2 & 0 \\
0 & 2 \\ \end{bmatrix}  $$

$$|det J| = |2 \times 2 - 0| = 4 $$

그러므로, $p_X(x) = p_Z(z) \cdot \frac{1}{4}$이다. 
 따라서 $x$의 공간(구간)이 4배 커졌지만, x의 확률 밀도($p_X(x)$)는 z의 확률 밀도($p_Z(z)$) 보다 4배 작아졌다

--

야코비안 행렬이 의미하는 것은 입력 $z_j$가 조금 변했을 때, 출력 $f_i$가 어떻게 변하는지 전부 기록해 놓은 것이다. 그래서 이것의 행렬식(det)를 구하면 입력 공간이 얼마나 늘어나거나 줄어드는지를 측정할 수 있다. 

$\begin{vmatrix} det(J_f(z))\end{vmatrix}$ = 공간이 얼마나 변형되었는가

- $\begin{vmatrix} det(J_f(z))\end{vmatrix} = 1$: 면적/부피 보존
 $\begin{vmatrix} det(J_f(z))\end{vmatrix} > 1$: 공간이 늘어남 $\rightarrow$ 밀도는 작아짐
- $\begin{vmatrix} det(J_f(z))\end{vmatrix} < 1$: 공간이 줄어듦 $\rightarrow$ 밀도는 커짐

--

다시 돌아와서 우리는 쉬운 분포 $z \sim N(0, I)$를 통해 $x$의 분포를 추정하고자 하는 과정에서 가역함수 $f$를 통해 $x$의 densitiy function $p_X(x)$를 다음과 같이 추정할 수 있었다.

$$ p_X(x) = p_Z(z) \cdot |\frac{dz}{dx}|$$

여기서 $\frac{dz}{dx}$는 역변환 $f^{-1}$의 jacobian이다. ($f^-1(x) = z$)

$$ p_X(x) = p_Z( f^-1(x) ) \cdot |\frac{d f^-1(x) }{dx}|$$

우리는 데이터 $x$가 주어졌을 때 이를 $f^{-1}$로 base space (쉬운 분포 z)로 되돌린 후, density p_X(x)를 계산하는 것이 목적이다. 이때 jacobian 보정이 필요한데 우리는 $f^(-1)$의 야코비안을 직접 모르고, $f$의 jacobian만 아는 경우가 많다.

그래서 Inverse Function Theorem을 도입한다.  Inverse Function Theorem은 어떤 점 근처에서 이 함수가 잘 정의되고, 되돌릴 수 있다면 (invertible) 그 역함수도 미분 가능하고, 그 변화율(jacobian)은 정방향의 jacobian의 역행렬이다.

$$\frac{dz}{dx}=\frac{f^{-1}}{dx}=(\frac{df}{dz})^{-1}$$ 

(머신러닝으로 비유하면 학습 데이터 x를 알고, base 분포 z도 알고 모델 f를 아는 상황 (f의 역함수만 모름)) $\rightarrow$ 아는 요소들로만 다시 재구성 하자

그래서 수식을 다시 쓰면

$$p_X(x) = p(z=f^{-1}(x)) \cdot |det(\frac{df}{dz}(z=f^{-1}(x)))|^{-1}$$

현실에서는 $x = f(z)$를 단 하나의 변환으로 구성하기 어렵기 때문에 다음처럼 여러 번의 변환 $f_1, f_2, ..., f_k$로 구성한다.

$$x = f_K \circ f_{K-1} \circ \cdots \circ f_1(z_0)$$ 

그럼 최종 밀도는 

$$p_X(x) = p_{Z_{0}}(z_0) \cdot \prod_{k=1}^{K}|det(\frac{df_k}{dz_{k-1}})|^{-1}$$

즉, base density에 모든 jacobian 보정 인자를 곱한 것의 역수이다.

계산을 용이하게 하기 위해서 log를 씌우면

$$logp_{Z}(z_0) - \sum_{k = 1}^{K}log|det \frac{df_i}{dz_{i-1}}|$$

이다. Transformation function $f_i$는 쉽게 invertible하고, 그것의 jacobian의 행렬식이 쉽게 계산될 수 있는 함수로 설정해야 한다.  

Interventional Density Estimation

우리는 observation data를 기반으로, treatment $A$를 $a$로 개입했을 때의 결과 $Y[a]$의 분포, 즉 interventional density $P(Y[a] = y)$를 추정하고자 한다. 

[Notation]

- $P(Z)$: 확률 변수 Z의 분포 전체
- $P(Z = z)$: 확률 질량 또는 밀도 함수
- $\pi_{a}(x) = P(A = a | X = x)$: propensity score, 환자 특성 $X$에 대해 치료 $A = a$를 받을 확률 
- $P(Y | X, A)$: conditional outcome distribution, 관찰된 치료 A, 공변량 X 하에서의 결과 분포
- $Y[a]$: potential outcome, "만약 A = a였더라면"의 결과
- $\widehat{P}(Y[a] = y)$: 우리가 추정하고 싶은 interventional distribution

관찰 데이터로부터 분포를 추정하고, 추정된 분포와 진짜 분포(정답 분포가 있어야 함)의 차이를 average log-likelihood, empirical wassertein distance 등을 통해 비교한다.

데이터셋 $D = \left\{(X_i, A_i, Y_i) \right\}^{n}_{i=1} \sim P(X, A, Y)$는 공변량 $X$, 치료 $A\in \left{ 0, 1\right\}$, 연속형 결과 $Y$의 삼쌍으로 이루어져있음. 여기서 X는 treatment와 outcome 모두에 영향을 줄 수 있는 변수이므로, confounder이고, 만약 $P(Y[a]) \neq P(Y | A = a)$ 라면 confounding bias가 존재한다는 것임.

(하다보니까 section4.1에서 수식에 propensity score를 활용하고 있는걸 발견 -> 우리의 상황에서는 쓰지 못함 우리는 representation에서 어떤 치료를 받을지에 대한 정보를 지워버렸기 때문에)

방향을 틀어서 그냥 어차피 CausalML로 potential outcome을 다 추정을 해놨으니까 이 데이터를 기반으로 NF 기법 아무거나 쉬운거 하나 적용해보자

Affine Coupling Layer -> D차원 입력이 필요해서 pass (우리의 경우 potential outcome은 1차원)

Planar flow

Planar flow는 NF의 한 종류로, 1D 또는 low-dimensional setting에서도 사용할 수 있다는 점에서 유용함

$$f(z) = z + u \cdot tanh(w^{T}z +b)$$

$u$, $w$는 학습 가능한 파라미터

$$|\frac{df}{dz}|=|1 + u \cdot (1 - tanh^2(wz+b)) \cdot w|$$

 

데이터 내에서 시간에 따른 분포 변화를 식별 -> Conditional Planar flow 필요

 

일단 빠르게 컨셉만 확인할겸 GPT에 물어봄