일관성 모델

생성과 제어와 인지를 연결하는 관점

자연과학이란 대체로, 조건을 상정하고 그러한 조건을 만족하는 함수를 찾는 게임이다. 가령 아래와 같은 방정식을 보자.

\begin{equation} i\hbar\frac{\partial}{\partial t}\Psi(\mathbf{r},t)=\hat{H}\Psi(\mathbf{r},t) \end{equation}

$\mathbf{r}$ 은 입자의 위치, $t$ 는 시간이다. 어떤 입자의 상태가 변화하는 속도는 그 입자의 에너지에 비례한다는 상식적인 서술이다. 이러한 조건을 만족하는 $\Psi(\mathbf{r}, t)$ 를 찾는다면 입자의 거동을 설명하고 예측할 수 있다.

수치적으로는 어떠한 위치와 시간에 대해서도 양변을 같게 만드는 함수 $\Psi(\mathbf{r},t)$ 를 찾으라는 말이다. 당연하게도 아무 함수나 이런 관계를 만족하지는 않으며 그런 함수를 찾는 과정은 대체로 쉽지 않은 일이다.

딥러닝도 마찬가지다. 조건을 상정하고 그런 조건을 만족하는 함수를 찾는다. 조건은 우리가 풀고자 하는 문제에 따라 달라진다. 가령 간단한 어떤 $\mathbf{x}$ 를 입력으로 받아 처리하는 신경망 $\Psi$ 를 생각해보자. 그렇다면 대개 이런 스타일의 조건을 상정한다.

\begin{equation} \Psi(\mathbf{x}) = \mathbf{y}({\mathbf{x}}) \end{equation}

$\mathbf{y}(\mathbf{x})$ 는 $\mathbf{x}$ 에 상응하는 어떤 Ground Truth라고 하자. 가령 언어 모델이라면 $\mathbf{x}$ 는 어떤 자연어 뭉치, $\mathbf{y}(\mathbf{x})$ 는 다음에 나올 단어가 되겠다. 이미지 분류 문제라면 $\mathbf{x}$ 는 이미지, $\mathbf{y}(\mathbf{x})$ 는 이미지 $\mathbf{x}$ 에 담겨 있는 대상들의 카테고리다. 이런 조건을 만족하는 함수 $\Psi$ 를 찾는다면 마찬가지로 요긴하게 사용할 수 있다.

대칭과 시각

두 이미지 $\mathbf{x}$ 과 $\mathbf{x'}$ 가 같은 컨텐츠를 담고 있다고 해 보자. $\mathbf{y}(\mathbf{x}) = \mathbf{y}(\mathbf{x'})$ 라는 말이다. 그렇다면 조건 $2$ 로부터 신경망 $\Psi$ 는 $\Psi(\mathbf{x}) = \Psi(\mathbf{x'})$ 을 만족해야 한다. 이를 새로운 조건으로 상정하고 싶다.

\begin{equation} \Psi(\mathbf{x}) = \Psi(\mathbf{x'}) \end{equation}

간단한 발상이다. 비슷한 입력이 들어온다면 비슷한 출력을 내라는 말이다. 더 이상 $\mathbf{y}(\mathbf{x})$ 의 구체적인 값은 중요하지 않기 때문에, $\Psi$ 의 입장에서는 이런 조건을 만족하기가 더 쉬울 것이다.

그렇다면 임의의 이미지 $\mathbf{x}$ 과 $\mathbf{x'}$ 가 같은 컨텐츠를 담고 있다는 사실은 어떻게 아는가? 알 수 없다. 다만 이미지 $\mathbf{x}$ 가 주어졌을 때 $\mathbf{x}$ 와 같은 컨텐츠를 담고 있는 또 다른 이미지를 생각해볼 수 있다. 가령 픽셀 몇 개의 값이 살짝 바뀌었다고 하여 컨텐츠가 달라지지는 않을 것이다. 이미지의 일부분을 가리거나 잘라 내었을 때도 그렇고, 밝기나 색상을 바꾸었을 때도 그렇다. 모두 컨텐츠를 보존하는 변환이다. 왜 그런 변환이 컨텐츠를 보존하는지 (opens in a new tab)는 흥미로운 질문이나 일단은 차치하자. 그런 변환에 보존되는 무언가를 우리가 컨텐츠로 인식한다고 해야 옳을 것이다.

이 다음부터는 간단하다. 이미지 $\mathbf{x}$ 를 약간 뒤틀어 $\mathbf{x'}$ 를 만든다. 최대한 $\Psi(\mathbf{x}) = \Psi(\mathbf{x'})$ 를 만족하게끔 $\Psi$ 를 조정한다. 완벽히 같게 만드는 것은 대체로 불가능하다. $\Psi(\mathbf{x})$ 와 $\Psi(\mathbf{x'})$ 의 간극을 최대한 좁히는 것이 최선이다. 더 많은 이미지와 더 극단적인 변환에 대해 이런 관계를 만족시키는 $\Psi$ 를 찾는다. 찾는다기보다는 만든다는 표현이 더 어울리는 것도 같다. 이렇게 만든 $\Psi$ 는 유용하게 사용할 수 있는데, 경우에 따라 조건 $2$ 를 상정했을 때보다 더 흥미로운 결과 (opens in a new tab)를 보여주기도 한다.

다만 위의 논리에는 큰 구멍이 있다. 가령 모든 이미지 $\mathbf{x}$ 에 대해 $\Psi(\mathbf{x}) = \mathbf{0}$ 이라면 어떤가? 모든 이미지를 받아 같은 출력을 내니 조건 $3$ 을 충족한다. 그러나 쓸모가 없다. $\Psi$ 가 쓸모 있으려면 다른 컨텐츠를 담고 있는 이미지가 들어왔을 때 다른 출력을 내어야 한다. 이미지 $\mathbf{x}$ 와 $\mathbf{x''}$ 가 서로 다른 컨텐츠를 담고 있다고 하자. 그렇다면 $\Psi$ 는 다음 조건 또한 만족해야 한다.

\begin{equation} \Psi(\mathbf{x}) \not = \Psi(\mathbf{x''}) \end{equation}

SimCLR (opens in a new tab)과 SwAV (opens in a new tab) 등의 방법론은 이런 조건을 직접 겨냥한다. 조건 $3$ 과 조건 $4$ 를 동시에 만족하는 $\Psi$ 를 찾는다. 논리적인 접근이다. 그러나 다른 컨텐츠를 담는 $\mathbf{x}$ 와 $\mathbf{x''}$ 를 정의하는 것은 전통적으로 번거로운 과제였다. 사실상 $\mathbf{x}$ 를 변형한 이미지가 아니라면 모두 $\mathbf{x''}$ 의 후보가 될 수 있기에 약간의 모호함이 생긴다.

흥미로운 발견은 반드시 조건 $4$ 를 직접 상정할 필요가 없다는 사실이다. BYoL (opens in a new tab)은 오직 조건 $3$ 만을 상정하며 앞서 말한 방법론들을 상회하는 성능을 보여준다. 이는 일견 이해되지 않는 결과이다. 왜 $\Psi$ 는 모든 입력에 대해 $\mathbf{0}$ 을 뱉는 간단한 해답을 거부하는가? Bootstrapping이 그 이유이다.

Bootstrapping

영미권에서 Bootstrap이란 자수성가를 일컫는 관용구이다. 장화 (Boot)의 끈 (Strap)을 묶고 비장하게 일터로 나서는 일꾼의 모습에서 유래된 말인 줄로 짐작했으나, 찾아보니 장화의 끈을 당겨 공중부양을 시도하는 바보의 모습에서 유래된 말이라고 한다. 아무튼, 스스로의 능력으로 성공을 노린다는 말이다.

조건 $3$ 을 다시 보자. 신경망 $\Psi$ 는 유사한 입력을 받아 일관된 출력을 내어야 한다. 즉, 우리는 일관성을 추구할 뿐이며, 어떤 구체적인 값을 원하는 것은 아니다. 앞서 말했듯 일관성을 달성하는 것은 쉽다. 모든 파라미터를 $\mathbf{0}$ 으로 설정해 버리면 된다. 고장난 시계에도 일관성은 있다.

현재로서 이런 고장난 시계를 만들지 않는 방법의 표준은 이렇다. $\Psi$ 로 하여금 이미지 $\mathbf{x}$ 에 대한 자신의 판단과, 이미지 $\mathbf{x'}$ 에 대한 과거 자신의 판단을 일치하게끔 만드는 것이다. $\Psi$ 의 과거란 무엇인가? 간단하다. 과거의 $\Psi$ 를 표현하는 또 다른 신경망을 생각하면 된다.

신경망 $\Psi$ 를 구성하는 파라미터를 생각하자. 과거와 현재를 표현해야 하므로 두 벌이 필요하다. 현재 $\Psi$ 를 표현하는 파라미터를 $\theta$ , 과거 $\Psi$ 를 표현하는 파라미터를 $\xi$ 라 하겠다. 우리가 원하는 조건을 수식으로 표현하자면 이렇다.

\begin{equation} \Psi_\theta(\mathbf{x}) = \Psi_\xi(\mathbf{x'}) \end{equation}

말했듯이 현재의 판단이 과거의 판단과 일치하도록 만들어 주자는 것이다. 수치적으로 $\xi$ 는 $\theta$ 의 이동 평균으로 정의한다. 과거 모든 $\theta$ 의 모습을 $\xi$ 에 기록하겠다는 의도이다. 시작 시점에는 $\xi$ 와 $\theta$ 를 동일하게 설정하는데, 이후로 $\xi$ 는 $\theta$ 를 천천히 따라가게 된다.

\begin{equation} \xi \leftarrow \tau \xi + (1-\tau)\theta \end{equation}

조건 $5$ 를 만족하도록 $\theta$ 를 조정한다. $\theta$ 를 한번 바꾸었다면, 식 $6$ 에 따라 $\xi$ 도 갱신한다. 그런 과정이 적절히 수렴한다면 $\theta$ 와 $\xi$ 도 대체로 유사해질 것이며, 결과적으로 조건 $3$ 을 달성하는 셈이다. 바로 이를 Bootstrapping이라 부른다.

왜 이런 기교가 $\Psi(\mathbf{x}) = \mathbf{0}$ 과 같이 자명한 해답으로 빠지는 사태를 방지해 주는가? 가령 $\theta$ 와 $\xi$ 가 모두 $\mathbf{0}$ 이 되어버릴 가능성은 없는가? 사실은 Bootstrapping이 이러한 붕괴를 방지하는 기전 역시 뚜렷하게 밝혀지지 않았다. 그러나 직관적으로 이해는 간다. 마냥 이리저리 휘둘리기만 하는 사람보다는, 과거의 자신을 믿는 뚝심있는 사람이 더 좋은 결과에 도달한다는 말이니. 인간은 스스로를 가르치며 성장한다 (opens in a new tab)는 말도 있지 않은가.

그러나 과거의 내가 신뢰할만한 대상이라는 보장은 어디에 있는가? 어찌되었든 백지에서 시작하는데, 단순히 과거의 나를 신뢰했더니 영리해졌다는 결과는 쉽게 납득하기 어렵다. 때문에 이를 Dark Knowledge로 부르기도 하는데, 실은 신경망의 구조적 특성 (opens in a new tab)이 물밑에서 작용한 결과로 보아야 한다.

어쨌든 이유는 중요하지 않다. Bootstrapping은 조건 $3$ 이 암시하는 일관성을 효과적으로 달성하기 위한 한 가지 기교에 불과하다. 이런 일관성을 이용하면 우리는 보는 기계를 만들 수 있다. 식 $1$ 이 입자의 거동을 지배하는 법칙인 것처럼, 식 $3$ 은 시각의 작용을 지배하는 법칙인 것이다. 인류는 진화를 통해 이 방정식을 풀었지만, $\Psi$ 는 계산을 통해 푼다는 차이가 있을 뿐이다.

이러한 일관성의 예시는 무수히 많이 들 수 있는데, 사실상 모든 자연 법칙은 이런 형태로 기술되기 마련이다. 우리는 자연과학자가 아니니 조금 더 현실적인 예시를 생각해보자.

선택과 보상

죽기 전까지 돈을 얼마나 벌 수 있을까? 가끔 생각해보는 질문이다. 언제 죽을지도 모르는 인생, 그 수입을 예측하기란 더욱 어려운 일이다. 그러나 앞으로 벌 돈에는 희한한 속성이 있다. 오늘부터 벌 돈은, 오늘 벌 돈에 내일부터 벌 돈을 더한 금액과 같다는 것이다. 무슨 당연한 말인가 싶지만 어쨌든 사실이다.

이로부터 앞으로 벌 돈이 만족해야 하는 법칙을 생각해 볼 수 있다. $t$ 시점 이후로 벌 돈을 $\Psi(t)$ , $t$ 시점과 $t+1$ 시점 사이에 번 돈을 $r_t$ 라 하자. 그렇다면 아래와 같은 관계가 성립한다.

\begin{equation} \Psi(t) = r_t + \Psi(t+1) \end{equation}

마찬가지로 $\Psi$ 의 판단에는 일관성이 있어야 한다는 말이다. 하루에 벌 수 있는 돈에는 한계가 있는 법이다. 그러므로 $\Psi$ 는 $t$ 와 $t+1$ 시점에서 크게 다르지 않은 판단을 내어야 한다. 조건 $7$ 을 만족하는 모든 함수가 앞으로 벌 돈을 의미하는가? 이는 생각해 볼 만한 질문이다. 그러나 앞으로 벌 돈을 의미하는 함수라면 분명 조건 $7$ 을 만족해야 한다.

조건 $7$ 은 타당하나, 현실적으로 응용하기는 어렵다. $\Psi$ 는 시점만 보고도 앞으로 벌 돈을 알려주어야 하니, 사실상 당신의 인생에 대한 모든 정보를 알고 있어야 한다. $\Psi$ 는 신적 존재가 되어야만 하는 셈이다. 너무 큰 부담을 지우지는 말자. 나의 상황을 객관적으로 판단할 수 있는 최소한의 근거를 던져주는 것이 합리적이다.

\begin{equation} \Psi(\mathbf{s}_t) = r_t + \Psi(\mathbf{s}_{t+1}) \end{equation}

$\mathbf{s}_t$ 는 $t$ 시점에서 당신의 상황에 대한 정보이다. 학력, 건강 등 미래 수입을 추론하기에 충분한 정보여야 이상적일 것이다. $\Psi$ 는 당신의 상황을 기반으로 미래 수입을 판단하는 함수가 된다. $\Psi$ 의 판단이 정확하다면, 내일의 상황 $\mathbf{s}_{t+1}$ 을 기반으로 판단한 $\Psi(\mathbf{s}_{t+1})$ 과 $\Psi(\mathbf{s}_t)$ 사이에는 상기 관계가 성립해야 한다.

조건 $3$ 이 떠오르지 않는가? 따라서 이 조건을 만족하는 $\Psi$ 또한 신경망으로 표현해보고 싶은 욕망이 든다. 이 경우에는 $\Psi(\mathbf{s})=\mathbf{0}$ 처럼 자명한 해는 존재하지 않는 듯 보이니 상황은 더 낫다. 이는 아주 (opens in a new tab) 오래 (opens in a new tab) 전 (opens in a new tab)부터 있었던 시도였으나 녹록지 않은 문제였다. 많은 해석과 시도가 있었으나, 역시 지금 시점에서 가장 무난한 풀이는 Bootstrapping을 사용하는 것이다.

\begin{equation} \Psi_\theta(\mathbf{s}_t) = r_t + \Psi_\xi(\mathbf{s}_{t+1}) \end{equation}

$\theta$ 와 $\xi$ 의 관계는 식 $6$ 과 같다. 적당한 가정과 문제 상황 하 (opens in a new tab)에서, 조건 $9$ 는 우리가 원하는 $\Psi$ 를 찾을 수 있는 레시피이다. 그러나 이러한 서술 (opens in a new tab)에는 선택이 배제되어 있다. 당신은 흘러가는대로 사는 존재가 아니며, 매 순간 선택을 거듭하며 더 나은 보상을 추구하는 존재다. 따라서 선택이라는 개념을 도입해 문제를 확장하는 것이 자연스러우며 또 유용하다.

$\Psi$ 가 행동과 상황을 동시에 고려하길 바란다. 이런 상황에서 저런 행동을 했을때 생애소득이 어떻게 결정될지 알고 싶다는 말이다. 이런 $\Psi$ 가 만족해야 하는 일관성은 다음과 같다.

\begin{equation} \Psi(\mathbf{s}_t, \mathbf{a}_t) = r_t + \Psi(\mathbf{s}_{t+1}, \mathbf{a}_{t+1}) \end{equation}

$\mathbf{a}_t$ 는 $t$ 시점에서 당신의 선택이다. 오늘의 상황에서 어떤 선택을 했는지는 당일 수입과 내일의 상황을 결정할 것이다. 물론 당신은 내일도 선택을 할 것이다. 이 다섯 가지 변수에 대해 $\Psi$ 는 조건 $10$ 에 따른 일관성을 만족해야 한다. 거꾸로 이러한 일관성을 만족하는 $\Psi$ 를 찾는다면 당신의 생애소득을 판단할 수 있을 것이다. 적어도 이상적으로는 그렇다.

그래서 그게 어쨌다는 것인가? 단순히 생애소득을 아는 것으로는 큰 의미가 없다. 어떤 선택을 해야 생애소득을 더 올릴 수 있을지, 즉 제어를 해낼 궁리를 하는 것이 자연스럽다. 여기서 재미있는 발상을 해볼 수 있는데, 당신이 오직 최선의 선택만을 거듭하는 현자라고 가정하는 것이다. 그런 가정 하에서, $\Psi$ 가 어떤 일관성을 만족해야 하는지 생각해보자.

$\Psi(\mathbf{s}_t, \mathbf{a}_t)$ 이 어떤 양과 일관되어야 하는가? 오늘의 상황에서 어떤 선택을 했는지가 주어졌으니, 당일 수입과 내일의 상황까지는 자동으로 결정된다. 그러나 내일의 선택부터는 달라진다. 현자의 의지가 개입한다. 현자는 반드시 내일부터 벌 돈을 가장 크게 만드는 선택을 해야 한다. 이외의 선택을 한다면 현자의 자격이 없다.

\begin{equation} \Psi(\mathbf{s}_t, \mathbf{a}_t) = r_t + \Psi(\mathbf{s}_{t+1}, \mathbf{a}_{t+1}^{*}) \end{equation}

\begin{equation} \mathbf{a}_{t+1}^{*} = \argmax_{\mathbf{a}} \Psi(\mathbf{s}_{t+1}, \mathbf{a}) \end{equation}

더 간결하게도 쓸 수 있는데, 사실은 위의 기술이 조금 더 직관적이기는 하다. Bootstrapping까지 고려해 쓰면 이렇다.

\begin{equation} \Psi_{\theta}(\mathbf{s}_t, \mathbf{a}_t) = r_t + \max_{\mathbf{a}}\Psi_{\xi}(\mathbf{s}_{t+1}, \mathbf{a}) \end{equation}

이러한 일관성을 만족하는 $\Psi$ 를 찾아 나가는 과정을 Q-Learning, 특히 신경망으로 풀고자 한다면 Deep Q-Learning이라 부른다. 이제는 10년이 다 되어가는 알파고 hype을 터트렸던 방정식이다.

조건 $13$ 은 행위자가 현자라는 가정을 다른 말로 표현한 것에 지나지 않는다. 그러한 기술에서도 일관성을 확인할 수 있다는 점이 재밌다. 이 방정식을 푸는 것은, 즉 이러한 일관성을 만족하는 $\Psi$ 를 찾는 것은 곧 현자를 찾는 것과 다름없는데, 일견 직관적이지 않을 수 있으니 조금 더 생각해보자.

가령 당신이 어떤 상황 $\mathbf{s}$ 에 처했다고 하자. 당신에게 조건 $13$ 을 만족하는 $\Psi$ 가 주어졌다. 그런 상황에서 최선의 선택을 하고 싶다면 $\Psi(\mathbf{s}, )$ 를 최대로 만드는 두 번째 인자를 골라 행하면 된다. $\Psi$ 는 현자의 생애소득을 알려주는 지표이고, 현자는 본인의 생애소득을 가장 키우는 선택을 할 것이기 때문에 그렇게 선택한 행동이 곧 현자의 선택이다. 논리가 약간 미묘하지만 곰곰히 생각해보면 이해가 갈 것이다.

이러한 현자의 방정식 (opens in a new tab)을 풀어 나가는 분야가 강화 학습이다. 다만 위와 같은 기술은 $\Psi$ 로 하여금 아주 먼 미래의 선택까지 고려할 것을 강요하는 측면이 있다. 때문에 개인적으로는 신경망으로 $\Psi$ 를 모델링하는 것에 회의적이다. 정말 단 한번 $\Psi(\mathbf{s}, \mathbf{a})$ 를 계산하는 것 만으로 먼 미래의 소득까지 계산한다는 것이, 그런 신경망이 존재한다는 것이 가당키나 한 일인가? 이는 우리 인간에게도 여러 단계의 추론과 논리를 요구하는 영역이며 이는 작금의 딥러닝으로 해결할 수 있는 스타일의 문제가 아니다. 때문에 MCTS 따위의 트리 서치와 함께 사용하는 모습을 흔히 볼 수 있다.

혹은 조건 $13$ 을 완전히 무시하고 흑마법의 영역으로 끌고 들어가려는 시도 (opens in a new tab)가 (다시 데리고 나오려는 시도 (opens in a new tab) 또한) 보이는데, 이 방향이 몹시 흥미롭다. 무시한다기보다는 조건 $13$ 의 완전히 다른 모델링이라고 보아야 할 것이다. 이는 폭발적으로 성장하는 (opens in a new tab) Foundation 모델 (opens in a new tab)과 엮기 용이하며 또 필연적으로 그렇게 되리라 본다. 기회가 되면 더 이야기 해 보자.

확산과 생성

언어 모델과 더불어, 딥러닝을 향한 세간의 관심을 견인하는 주체는 이미지 생성 모델이다. 시각을 원한다면 조건 $3$ 을, 제어를 원한다면 조건 $13$ 을 풀면 되었다. 그렇다면 이미지 생성을 위해서는 어떤 방정식을 풀어야 하는가? 이는 이미지의 생성 과정을 어떻게 서술할 것인지에 따라 다르다. 가장 인기있는 방식은 단연 확산 모델 (Diffusion Models) (opens in a new tab)이다.

확산 모델의 아이디어는 간단하다. 물에 풀어놓은 잉크가 확산되듯이, 어떤 이미지 $\mathbf{x}$ 가 확산되어 형태를 알아볼 수 없는 노이즈 $\mathbf{z}$ 로 변모하는 과정을 상정하는 것이다. 이 과정을 되돌리는 신경망 $\Psi$ 를 만들어, 노이즈 $\mathbf{z}$ 로부터 $\mathbf{x}$ 를 복원하자는 발상이다.

확산 모델의 가장 최근 모습은 일관성 모델 (Consistency Models) (opens in a new tab)이다. 이름에서 감이 오지 않는가? 지난 포스팅에서 보았던 방법론들과 궤를 같이한다. 일관성 모델은 확산 과정의 어느 시점에서도, 원본 이미지 $\mathbf{x}$ 를 복원해주는 신경망 $\Psi$ 을 일컫는다.

확산의 초기 시점을 $0$ , 마지막 시점을 $T$ , $t$ 시점의 이미지를 $\mathbf{x}_t$ 라 하자. 그렇다면 $\mathbf{x}_0=\mathbf{x}$ 이고 $\mathbf{x}_T=\mathbf{z}$ 이겠다. 말했듯이 일관성 모델은 어떤 $t$ 에 대해서도, 원본 이미지 $\mathbf{x}$ 를 출력하는 $\Psi$ 를 일컫는다. 즉 $\Psi$ 를 기술하는 방정식은 아래와 같다는 말이다.

\begin{equation} \Psi(\mathbf{x}_t) = \Psi(\mathbf{x}_{t'}) = \mathbf{x}_0 \end{equation}

이를 왜 일관성 모델이라 부르는지 이해할 수 있을 것이다. $\Psi$ 로 하여금 ( $0$ 과 $T$ 사이의) 그 어떤 시점 $t$ 와 $t'$ 에 대해서도 일관된 출력을 낼 것을 요구하기 때문이다. 앞서 유도했던 조건들과의 형태적 유사성이 느껴진다.

단순히 이미지 $\mathbf{x}_t$ 만을 제공하기보다는, 정확히 어떤 시점의 스냅샷인지를 알려주어야 $\Psi$ 가 편할 것이다. 시점 $t$ 를 알려주자.

\begin{equation} \Psi(\mathbf{x}_t, t) = \Psi(\mathbf{x}_{t'}, t') =\mathbf{x}_0 \end{equation}

시작 시점 $t=0$ 에서도 마찬가지다. $\Psi(\mathbf{x}_0, 0)=\mathbf{x}_0$ 라는 말인데, 이는 곧 $\Psi(\mathbf{x},0)$ 가 항등함수라는 조건이 된다. 이 조건은 약간의 기교를 부려 쉽게 달성할 수 있다. 아래 코드를 보자.

from diffusers import UNet2DModel
from torch import nn
import torch
 
T = torch.tensor([80.])
unet = UNet2DModel(sample_size=(32, 32))
 
class Consistency(nn.Module):
    def __init__(self, unet: UNet2DModel):
        super().__init__()
        self.unet = unet
 
    def forward(self, x: torch.Tensor, t: torch.Tensor):
        skip = torch.einsum("b, b c h w -> b c h w", 1 - t/T, x)
        out = torch.einsum("b, b c h w -> b c h w", t/T, self.unet(x, t).sample)
        return skip + out
 
model = Consistency(unet=unet)

$T$ 는 $80$ 으로, cifar10을 염두에 두고 있기에 해상도는 $32 \times 32$ 로 잡았다. 입력과 출력의 shape가 같아야 하니 UNet이 무난한 선택이다. diffusers의 UNet2DModel (opens in a new tab)을 사용하자. 출력 타입은 UNet2DOutput (opens in a new tab)이다.

크게 복잡할 것은 없다. $\Psi(\mathbf{x},0)$ 를 항등함수로 만들기 위해 UNet의 출력과 입력 $\mathbf{x}$ 를 내삽한다. 시간의 영향을 받는 Skip Connection이라 생각하면 쉬울 것이다. $t=0$ 일 때는 입력을 그대로 뱉으며, $t=T$ 일 때는 UNet의 출력을 그대로 뱉는다. 정말 그런가?

images = torch.randn(16, 3, 32, 32)
times = torch.zeros(16)
torch.allclose(images, model(images, times))  # True

정말 그렇다. 다만 실제로 내삽을 하진 않는데, 이해를 돕기 위해 상황을 단순화했다. 어쨌든 발상은 동일하다. 실제로 times는 $0$ 과 $T$ 사이의 값을 가지는 텐서가 될 것이다.

이제 조건 $15$ 를 만족하도록 $\Psi$ 를 조정해보자. 조건 $15$ 를 $\Psi(\mathbf{x_t},t)=\mathbf{x}_0$ 의 관점에서만 바라본다면 단순한 Denoising Autoencoder (opens in a new tab)다. 노이즈가 뿌려진 이미지 $\mathbf{x}_t$ 를 받아 노이즈를 제거하라는 문제를 푸는 것이다. 작금의 확산 모델은 대개 이런 목표를 직접 겨냥하는데, 사실 이는 이미지 생성이라는 거창한 과업을 풀기에는 지나치게 무심한 목표라고 본다. Denoising Autoencoder라는 컨셉은 오랫동안 알려져 있었으나 이를 선뜻 이미지 생성 모델로 사용하지 못했던 이유가 있다. 이러한 이유 때문인지 확산 모델로 품질이 높은 이미지를 생성하기 위해서는 수십, 수백번의 보정이 필요하다.

일관성 모델은 $\Psi(\mathbf{x}_t, t) = \Psi(\mathbf{x}_{t'}, t')$ 에 주목한다. 노이즈가 없는 이미지를 만들어 내는 것이 아니라, 확산 과정의 두 스냅샷 $\mathbf{x}_t$ 와 $\mathbf{x}_{t'}$ 에 대해 신경망 $\Psi$ 가 같은 출력을 내게끔 만들어 주기만 하면 된다. $t<t'$ 라 할 때, 시간이 흐르면 $\mathbf{x}_t$ 가 $\mathbf{x}_{t'}$ 로 변하게 되어 있으며, $\Psi$ 는 그렇게 시간으로 엮인 두 지점을 같은 도착지로 보내주는 역할을 해야 한다는 것이다. 조건 $15$ 를 일관성의 관점에서 바라보는 셈이다.

그렇다면 시간상으로 연결되어 있는 $\mathbf{x}_t$ 와 $\mathbf{x}_{t'}$ 을 어떻게 표현할 수 있는가? 이는 생각만큼 간단하지 않다. 가령 물에 잉크를 풀어 확산시켰을 때, 10초와 11초에 찍은 스냅샷 간 관계를 표현할 수 있겠는가? 글쎄다. 정확한 표현을 위해서는 확산 과정이라는 물리적 현상에 대한 최소한의 모델링이 필요하다.

물리 시간은 아니므로 대강 결론만 보자. $t$ 와 $t'$ 가 시간상으로 아주 가깝다고 치면 간단한 근사를 할 수 있다. $t'=t + \Delta t$ 라 표기하면 아래와 같다.

\begin{equation} \mathbf{x}_t \approx \mathbf{x}_0 + t\mathbf{z} \end{equation}

\begin{equation} \mathbf{x}_{t+\Delta t} \approx \mathbf{x}_0+(t+\Delta t)\mathbf{z} \end{equation}

$\mathbf{z}$ 는 표준 정규분포에서 추출한 노이즈이다. $\mathbf{x}_t$ 는 원본 이미지 $\mathbf{x}_0$ 와 노이즈 $\mathbf{z}$ 의 합으로 표현되는데, 확산 과정이 $\Delta t$ 만큼 더 진행되면 노이즈가 차지하는 비중이 그만큼 커진다는 말이다. 동일한 노이즈 $\mathbf{z}$ 를 이용해 두 이미지가 시간상으로 엮여 있음을 표현한 셈이다. 상당히 거친 근사인데, 어쨌든 현실적으로는 다음과 같은 조건을 만족하는 $\Psi$ 를 찾으면 된다. 조건 $2$ 를 다시 쓴 것이다.

\begin{equation} \Psi(\mathbf{x}_0 + t\mathbf{z}, t) = \Psi(\mathbf{x}_0 + (t+\Delta t)\mathbf{z} , t + \Delta t) \end{equation}

실제 학습 로직은 대강 아래와 같다. $t=0$ 과 $t=T$ 사이에서 적당한 시간을 골라 쓰면 된다. $t+\Delta t$ 가 $T$ 를 넘으면 안되니 정확히는 $t=T-\Delta t$ 까지겠다.

delta = 0.5
for images in dataloader:
    t = torch.rand(32) * (T - delta)
    z = torch.randn(32, 3, 224, 224)
 
    noise_at_t = torch.einsum("b, b c h w -> b c h w", t, z)
    noise_at_delta_t = torch.einsum("b, b c h w -> b c h w", t + delta, z)
 
    left = model(images + noise_at_t, t)
    right = model(images + noise_at_delta_t, t + delta)
 
    loss = F.l1_loss(left, right)
 
    loss.backward()
    optimizer.step()

여기서 dataloader는 cifar10 이미지를 32장씩 던져주며, optimizer는 model의 파라미터를 바라본다. 배치 단위로 처리하는 코드이므로 헷갈릴 수 있으나, 조건 $18$ 을 그대로 옮긴 것이니 비교를 해 가며 읽어보길 바란다.

물론 이런 상황에서는 $\Psi$ 가 $0$ 으로 빠져버릴 가능성이 몹시 높다. left와 right의 차이를 좁히는 가장 간단한 방법이 아닌가? 이렇게 $\Psi$ 에게 일관성을 요구하는 상황에서는 Bootstrapping을 사용하는게 표준이라고 했으니, 여기서도 본능적으로 그런 생각이 들어야 한다.

import copy
 
model_past = copy.deepcopy(model)
tau = 0.999
 
@torch.no_grad()
def update_past(model, model_past):
    for p, p_past in zip(model.parameters(), model_past.parameters()):
        p_past = tau * p_past + (1 - tau) * p
 
for images in dataloader:
    t = torch.rand(32) * (T - delta)
    z = torch.randn(32, 3, 224, 224)
 
    noise_at_t = torch.einsum("b, b c h w -> b c h w", t, z)
    noise_at_delta_t = torch.einsum("b, b c h w -> b c h w", t + delta, z)
 
    with torch.no_grad():
        left = model_past(images + noise_at_t, t)
 
    right = model(images + noise_at_delta_t, t + delta)
 
    loss = F.l1_loss(left, right)
 
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    update_past(model, model_past)

model의 학습이 끝나고 나면 어떤가? 이상적으로 model은 확산 과정의 한 스냅샷을 원본 이미지로 보내 주어야 한다. 확산 과정의 마지막 스냅샷 $\mathbf{x}_T = \mathbf{x}_0 + T\mathbf{z}$ 을 생각해보자. 마지막 시점 즈음이 되면 $T$ 가 충분히 커져 사실상 노이즈가 된다. $\mathbf{x}_T \approx T\mathbf{z}$ 로 보아도 된다는 말이다. 정확히는 그렇게 되게끔 확산 과정과 $T$ 를 설정한다고 해야 옳겠지만. 그러니 노이즈 $\mathbf{z}$ 를 뽑아 $\Psi(T\mathbf{z}, T)$ 를 계산하면 이미지를 생성할 수 있다는 논리다.

noise = torch.randn(1, 3, 224, 224) * T
model(noise, T)  # Generated Image 🏞

$\Psi(T\mathbf{z}, T)$ 은 순수한 노이즈로부터 만든 이미지이니 품질 기준에 미치지 못할 가능성이 높다. $\Psi$ 의 입장에서 생각해보자. 순수한 노이즈에서 이미지를 만드는 것 보다는, 어느 정도 형태가 잡혀 있는 이미지를 복원하는게 더 쉽게 느껴지지 않겠는가? 그래서 몇 번의 계산을 더 감당할 수 있다면 생성된 이미지의 품질 개선이 가능하다.

계산한 $\Psi(T\mathbf{z}, T)$ 를 새로운 원본 이미지 $\mathbf{x}_0$ 으로 간주하는 것이다. 말하자면 $\Psi$ 의 첫 번째 시도이다. 이를 다시 확산시키고, $\Psi$ 를 통해 복원한다. 확산 시간을 점점 줄여가며 이러한 과정을 반복한다면 더 고품질의 이미지를 생성할 수 있다. 계산과 품질을 교환할 수 있는 여지가 있는 것이다. 그러나 한 번의 $\Psi(T\mathbf{z}, T)$ 계산 만으로도 충분히 괜찮은 이미지를 뽑을 수 있다는 것이 일관성 모델의 매력이 아닐까.

times = [60.0, 30.0, 15.0, 7.0, 3.0]
 
noise = torch.randn(1, 3, 224, 224) * T
image = model(noise, T)
 
for time in times:
    image = image + torch.randn(1, 3, 224, 224) * time
    image = model(image, torch.tensor([time]))  # Gets better and better!

일관성 모델은 상대적으로 단순한 관점으로 이미지 생성을 풀어낼 수 있다. 다만 이 일관성이라는 특징은 일관성 모델의 고유한 성질이 아닌, 우리가 살펴보았듯 다양한 맥락에서 나타난다는 사실을 이해할 수 있다. 이는 저자들의 관찰이기도 하다.

In addition, consistency models share striking similarities with techniques employed in other fields, including deep Q-learning (Mnih et al., 2015) and momentum-based contrastive learning (Grill et al., 2020; He et al., 2020). This offers exciting prospects for cross-pollination of ideas and methods among these diverse fields.

일관성은 다시 말해 불변성이며 우리가 찾는 해 공간에 강한 대칭 구조가 존재한다는 사실을 암시한다. 이러한 접근이 늘 유효한 것은 아니겠지만, 일관성을 찾아내 Bootstrapping으로 해결하는 방식은 우아하며 항상 놀라움을 준다.

Source code (opens in a new tab)