FlowSE 개선 아이디어 with knowledge distillation

 

 


[202504161110]

FlowSE 개선 아이디어, FlowSE에서 사용하는 ODE는 아래와 같다.

\frac{d\psi_t(x_1|y)}{dt}:=v_t(\psi_t(x_1|y)|y), \psi_1(x_1|y)=x_1

Loss는 아래와 같은 loss를 사용하죠.

\mathcal{L}_{CFM} = E \lVert v_\theta(x_t,y,t)-v_t(x_t|x_0,y)\rVert^2

위의 식에서 target value는 v_t(x_t|x_0,y) 이고 이것의 변수는 t,x_t, x_0,y 이다. 고런데, neural network v_\theta (x_t,y,t) 는 입력변수를 t,x_t, y 만 사용한다. 그리고 차피 inference 시에 x_1 \sim N(y,\sigma^2 I) 이므로 v_\theta(x_t,y,t) 는 y에 대한 정보를 사용하게 된다. v_\theta(x_t,y,t) 에서 두번째 변수 y 대신 x_0 를 사용할 수 있는 방법은 무엇일까? 아래와 같은 두개의 loss를 생각하자.

\mathcal{L}_{clean} = E \lVert v_\theta(x_t,x_0,t)-v_t(x_t|x_0,y)\rVert^2 \mathcal{L}_{KD} = E \lVert v_\theta(x_t,y,t)-v_\theta(x_t,x_0,t)\rVert^2

\mathcal{L}_{clean} 에서는 neural network변수에 x_0 를 넣었다. 아마 \mathcal{L}_{CFM} 보다 \mathcal{L}_{clean} 가 최적화가 잘 될것이다. inference 시에는 x_0 를 사용할 수 없다. 그렇지만 학습과정중에 x_0 에 관한 정보를 뉴럴네트워크에 지식전이할 수 있다. 그 loss가 \mathcal{L}_{KD} 이다. 이것을 종합하여 loss를 설계할 수 있겠다.

\mathcal{L}_{total} = \alpha_1 \mathcal{L}_{CFM} + \alpha_2 \mathcal{L}_{clean} + \alpha_3 \mathcal{L}_{KD}

여기서 \alpha_i 들은 0보다 크거나 같은 hyperparameter이다. 부족한 내용은 차차 채워 넣겠다.

 

Leave a Comment