📔

Diffusion Schrödinger Bridge를 적용하여 diffusion model을 개선할 수 있을까?

diffusion_schrodinger_bridge
JTT94

Generative Modeling by Estimating Gradients of the Data Distribution

Introduction

generative modeling techniques grouped based on how they represent probability distributions
1.
likelihood-based models
2.
implicit generative models
3.
score-based models
a.
(Stein) score function
i.
= the gradient of the log probability density function
b.
score matching
i.
directly learned

The score function, score-based models, and score matching

N개의 데이터를 가진 데이터셋을 가지고 있고, 이 데이터셋의 데이터 포인트들은 동일한 데이터 분포 p에서 각각 독립적으로 추출했다고 가정하자.
generative modeling 의 목적은 내가 만든 모델의 파라미터를 정해서 원본 데이터 분포 p와 유사하게 fit 하여, 내가 만든 모델에서 샘플링하는 것만으로도 진짜 추출한 데이터 포인트와 유사한 새로운 데이터 포인트를 마음껏 생성할 수 있게 하는 것이다.
그러기 위해서 가장 먼저 설정해야 할 것은 모델링하고자 하는 데이터 분포를 어떻게 표현할지인데, 이를 확률 분포로 표현하고자 한다.
likelihood-based model에서는 데이터가 연속적이냐 불연속적이냐에 따라 확률 분포를 확률 밀도 함수(pdf) 혹은 확률 질량 함수(pmf)로 표현한다. 여기서는 연속적인 데이터를 가정하여 확률 밀도 함수로 설명하자.
학습 가능한 파라미터들을 세타로 표현한다고 하고, 데이터 포인트들은 실수 값만을 가진다고 가정할 때, 단일 변수 x에 대해 실수 치역을 가지는 함수 Eθ(x)E_\theta(\mathbf{x})를 가정하자. (unconditional, single dependent variable)
부연 설명으로, 이 함수를 energy-based model 연구하시는 분들은 에너지라고 부른다. 누구는 그냥 energy-based model 이라고 부르기도 한다. 또 회귀 분석에서는 regression fuction 이라고 멋있게 부른다. 또 다른 이름으로는 unnormalized probabilistic model 라고 불리기도 한다. 이름이 정말 많은 친구인 걸 보면 꽤나 중요한 녀석임을 눈치챌 수 있다. 한 가지 덜 중요한 특징으로는 딱히 선형성을 가정한 적이 없으므로 비선형적이라고 볼 수 있다.
아무튼 위에서 정의한 이름이 참 많은 함수를 이용하여 아래와 같이 밀도 함수 p를 정의할 수 있다.
이를 멋있게 부르면, energy-based model 에 의해 정의된 밀도 함수 p 라고 부를 수 있다.
이때 Z는 normalizing constant 라고 부르는데, 모델의 파라미터에 따라 변하는 값으로 항상 양수이며 밀도 함수 p 전체를 적분했을 때 값이 1이 되도록 해주는 스케일 값이다. (pθ(x)dx=1\int p_\theta(\mathbf{x}) \textrm{d} \mathbf{x} = 1) 이 변수도 다른 멋있는 이름을 가지고 있는데, partition function이라고도 부른다. 왜 상수 (constant) 라고 부르냐면 x의 변화에 대해서는 불변하기 때문이다. 하지만 모델의 파라미터에 따라서는 변한다.
만약에 누군가 귀찮게 p를 구해달라고 하거나, (p를 구하기 위해) θ\theta에 대해 log p라도 미분해달라고 할 때 거절할 핑계를 찾아야 한다면 어떻게 해야할까? 당연히 모든 미적분은 원래 어려우니까 계산하기 어렵다고 말해도 된다. 이보다 더 그럴싸한 변명을 대야 한다면 모델의 파라미터 θ\theta에 대해 에너지(E)도 변하면서 동시에 normalizing constant Z를 정의하는 적분값 또한 변하기 때문에 미분하기 어렵다고 말할 수 있다.
이제 학습을 통해 찾고 싶은 것은 모델 파라미터 θ\theta이다. 이때 이 모델 파라미터의 이상적인 조건을 생각해본다면, 이 파라미터로 모델링 된 밀도 함수 p가 데이터셋에 대한 log-likelihood 값이 최대라면 좋을 것이다. likelihood 값을 최대로 하는 모델 파라미터 값을 구하는 것은 그 문장 그대로도 직관적으로 이해가 가는 부분이다. 말 그대로 likelihood가 높아지는 거니까.
굳이 log를 붙이는 이유는 log의 발명 이유와 마찬가지로 철저히 계산 상의 이유다. log를 이용해 곱셈을 덧셈으로 계산할 수 있다던가, 컴퓨터 계산 시 정밀도 자릿수 등등의 이유로 사용한다.
아무튼 이를 수식으로 적으면 아래 log-likelihood 값을 최대로 하는 파라미터 θ\theta를 학습을 통해 구한다고 할 수 있다. (이러한 행위를 멋있게 maximum likelihood training/estimation 이라고 부른다.)
그런데 아뿔싸! 위 식을 계산하려면 밀도 함수 p의 값을 계산해야 하는데, 방금 위에서 어렵다고 변명했던 짓을 해야 한다. 앞서 말했듯이 Z부터 계산이 어렵다. (이를 영어로 적을 때 hard 하다 말고 멋있는 문장으로 적으려면 due to the intractable normalizing constant Z 라고 표현하면 된다.)
이 부분에서 머리 좋으신 분들이 각자의 멋있는 아이디어를 내놓으셨는데, Z를 계산 가능하게 하기 위해 모델 구조에 제약을 걸거나 (e.g., autoregressive model의 causal convolution, normalizing flow model의 invertible networks) 아니면 (계산량이 무척 많이 실용적으로 활용하기에는 어렵지만) Z를 근사하는 방법들이 있다 (e.g., VAE의 variational inference, contrastive divergence에서 쓰이는 MCMC sampling).
밀도 함수 p를 모델링하기 어렵다면, 애초에 다른 함수를 모델링하는 건 어떨까? 밀도를 직접 계산하지 말고, 일종의 점수 값으로 비슷하게 표현할 수 있다면 충분하지 않을까? 그러면 (방금 배운 멋있는 표현으로 써보자면) intractable 한 normalizing constant의 문제를 빗겨 나갈 수 있지 않을까? 그래서 또 머리 좋으신 분들이 이를 밀도 함수 p의 score function 라고 부르고 아래와 같은 함수를 그 score function 으로 써보자고 주장하셨다.
직관적으로 생각할 때, 어떤 함수의 gradient 이니까 그 원본 함수가 어떻게 변화하는 지에 대한 경향성을 파악하는데는 충분할 것으로 생각된다. 또 이러한 score function의 모델 s 를 score-based model 이라고 부르기로 하셨다. 따라서 score-based model s 를 학습시킨다 함은 s 가 score fuction (=log p의 x에 대한 gradient) 와 최대한 유사하도록 s 파라미터 값을 찾아가는 것을 목적으로 한다고 할 수 있다. 이때는 직접 p를 계산할 필요가 없어지므로 다가가기 어려웠던 친구인 Z 와도 이별을 고할 수 있게 된다! 그러면 이제 (Z 없이) s 의 파라미터를 구하기 위한 식이 필요하다.
앞서 정의한 energy-based model에 의해 정의된 밀도 함수 p 를 위 score-based model s 로 표현하는 것도, 위 score function 에 p 식을 대입하면 된다. 대입하면 아래와 같다.
위 식에서 볼 수 있듯 normalizing constant Z 가 들어간 항은 무시된다! 따라서 더 이상 Z 를 신경 쓰지 않고 모델을 고를 수 있게 되는 것이다.
원본 글에는 GIF로 계속 두 그래프가 계속 바뀌면서 더 직관적으로 보여주고 있다.
위에서 우리가 조절할 모델은 score-based model로 정해졌다. 그렇다면 이제 학습을 위해 정답 값과 모델 출력 값의 차이를 계산하는 방법, 즉 loss를 어떻게 설정하는지만 결정하면 된다. 여기서는 단순히 값과 값을 비교하는 것이 아니라 분포와 분포, 정확히는 분포와 모델의 차이를 계산해야 한다. 아래와 같이 Fisher divergence 를 조금 변형하여 사용하여 아래와 같이 데이터 분포와 모델과의 차이를 계산한다.
참고로 원래 두 분포 p 와 q 에 대해 적용되는 Fisher divergence 의 식은 아래와 같다.
직관적으로 생각해보면, 위 (변형된) Fisher divergence 식은 score의 GT값과 score를 표현하는 모델의 l2 거리의 루트 제곱을 계산한다. 여기까지 읽으면 뭔가 이상하다는 생각이 들 것이다. 아니 score의 GT값을 바로 알 수 있으면 뭐하러 학습을 하나? score를 알아야 score의 모델에 대한 측정을 할 수 있다니!
다행히도 이런 상황에 써먹을 수 있게 score matching 이라는 기법을 이미 연구해놓으신 머리 좋으신 분들이 있었다. 많이 쓰이는 방법으로는 denoising score matching, sliced score matching 등이 있다고 한다. 이 기법을 이용해서 GT score 값을 모르고도 Fisher divergence 는 최소화할 수 있다.
score matching 의 목적 함수는 데이터셋에서 바로 계산이 가능하고, stochastic gradient descent로 최적화할 수 있다. 이런 특징은 likelihood-based model을 학습시키기 위해 normalizing constant 를 알고 있을 때 log-likelihood 목적 함수를 최적화 하는 것과 비슷하게 유추할 수 있다고 한다.
또 하나 주요한 특징으로는 adversarial optimization 없이도 score matching 목적 함수를 최소화하여 score-based model 을 학습시킬 수 있다는 점이다.
마지막으로 주요한 특징은 score matching 을 이용하면 score model (=score-based model) s 가 어떤 형태든
이하 작성 중

Langevin dynamics

Naive score-based generative modeling and its pitfalls

Score-based generative modeling with multiple noise perturbations

Score-based generative modeling with stochastic differential equations (SDEs)

Perturbing data with an SDE
Reversing the SDE for sample generation
Estimating the reverse SDE with score-based models and score matching
How to solve the reverse SDE
Probability flow ODE
Controllable generation for inverse problem solving

Connection to diffusion models and others

Concluding remarks

Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling

Score-based generative modeling

Schrödinger Bridge

Some advantages
Current limitation

Two dimensional examples

References

How to Train Your Energy-Based Models Y. Song, D.P. Kingma.arXiv preprint arXiv:2101.03288. 2021.