이산 확산 모델·디퓨전 언어모델 원리 정리: MIT 6.S184 마지막 강의로 보는 이산 플로우 매칭

핵심 메시지

• '디퓨전 언어모델'이라는 이름과 달리 그 안에는 확산 방정식도 미분방정식도 없고, 연속시간 마르코프 연쇄(CTMC)라는 다른 수학이 들어 있다.
• 연속 공간의 '벡터장' 자리에 이산 공간에서는 '전이율 행렬'이 들어가며, 방향으로 움직이는 대신 토큰을 다른 토큰으로 바꾸는 속도를 지정한다.
• 한 번에 토큰 하나만 바뀔 수 있다는 제약(희소성 조건)을 걸어야 조합적으로 폭발하는 상태 공간을 실제로 학습 가능한 크기로 줄일 수 있다.
• 연속 플로우 매칭의 학습이 결국 회귀 문제였다면, 이산 플로우 매칭의 학습은 각 자리의 원래 토큰을 맞히는 분류 문제로 귀결된다.
• 장점은 여러 토큰을 병렬로, 임의의 순서로 생성할 수 있다는 점이고, 단점은 KV 캐싱 같은 기존 최적화 기법을 그대로 쓰기 어렵다는 점이다.

쉽게 이해하기

MIT의 생성 모델 강의 6.S184의 마지막 시간은 이미지·영상 생성에서 다뤄 온 확산 모델의 원리를 문장이나 단백질 서열 같은 이산 데이터로 확장하는 방법을 다룬다. 강연자는 최근 화제가 된 이른바 '디퓨전 언어모델'을 먼저 언급하며, 구글의 텍스트 확산 모델처럼 글자를 왼쪽부터 순서대로 쓰지 않고 문장 전체를 조금씩 채워 나가는 모델들이 등장하고 있다고 소개한다.

다만 강연자는 이름이 오해를 부른다고 못 박는다. 이 모델들에는 확산 과정도, 확률미분방정식도, 상미분방정식도 들어 있지 않다는 것이다. 대신 바탕에 깔린 수학은 연속시간 마르코프 연쇄로, 상태가 시간이 흐르며 다른 상태로 '점프'하는 과정을 다룬다. 연속 공간에서 벡터장이 어느 방향으로 얼마나 움직일지 알려 줬다면, 이산 공간에서는 알파벳 Y에서 알파벳 X로 바뀌는 '속도'를 지정하는 전이율 행렬이 그 역할을 대신한다.

실용성의 열쇠는 희소성 제약이다. 길이 D인 문장이 가질 수 있는 모든 경우의 수는 어휘 크기의 D제곱이라 지수적으로 커지므로, 그 거대한 행렬을 통째로 저장하는 것은 불가능하다. 그래서 두 문장이 두 자리 이상 다르면 전이율을 0으로 두어, 한순간에는 한 자리만 바뀔 수 있게 제한한다. 그러면 신경망은 각 위치마다 어휘 크기만큼의 값만 내놓으면 되고, 크기가 차원에 선형으로만 늘어난다.

학습 방법도 연속 버전의 구조를 그대로 따라간다. 스케줄러 함수로 데이터를 서서히 지우고 잡음을 채우는 확률 경로를 정의하고, 데이터 하나에 조건부로 붙은 전이율 행렬을 데이터 분포 전체에 대해 평균 내는 '주변화' 단계를 거친다. 이때 정답 문장을 모른다는 사실만 반영하면 되는데, 그 자리에 들어가는 것이 '지금 망가진 문장을 보고 원래 이 자리에 무슨 토큰이 있었을지'에 대한 사후확률이다. 즉 학습해야 할 대상은 분류기이고, 손실 함수는 익숙한 교차 엔트로피다.

실제 구현으로 가장 흔한 형태는 마스크 확산 모델이다. 어휘에 '마스크'라는 새 토큰을 추가하고 전부 마스크로 채운 상태에서 시작해, 매 단계 각 위치의 전이율에 따라 마스크를 실제 토큰으로 하나씩 벗겨 낸다. 강의에서는 대형 마스크 확산 모델에 셰익스피어 문체로 보스턴에 대한 글을 요청해, 마스크가 왼쪽부터 순서대로가 아니라 여기저기서 풀리며 문장이 완성되는 과정을 시연한다.

주요 인사이트

• 기존 확산 모델과의 차이를 한 문장으로 압축하면 '질량을 옮기는 대신 확률의 무게를 갈아 끼운다'는 것이다. 연속 공간에서는 확률 질량을 공간상에서 이동시킬 수 있지만, 이산 공간에는 방향이 없으므로 잡음 분포를 서서히 줄이고 데이터 분포를 서서히 키우는 방식으로 경로를 만든다.
• 조건부 전이율 공식에 등장하는 계수는 시간이 끝에 가까워질수록 무한대로 발산한다. 강연자는 이를 '끝나는 시점에는 반드시 제자리에 가 있어야 하므로 마지막에 확실히 밀어붙이는 항'이라고 직관적으로 설명한다.
• 샘플링 절차는 전이 확률을 몰라도 된다. 전이율 행렬만 알고 있으면 시간 간격을 잘게 쪼갠 1차 테일러 근사로 각 단계의 확률분포를 만들 수 있고, 이는 위치마다 범주형 분포에서 뽑는 간단한 연산이 된다.
• 임의 순서 생성은 코드 파일 중간을 고치는 편집 작업이나, 단백질·DNA처럼 왼쪽에서 오른쪽이라는 자연스러운 순서가 없는 데이터에서 특히 쓸모가 있다. 반대로 말하기나 글쓰기처럼 순서 자체에 의미가 실리는 데이터에서는 기존 방식의 이점이 남는다.
• 강연자는 이 모든 것이 우연이 아니라고 덧붙인다. 미분방정식이든 마르코프 연쇄든, 생성 과정을 규정하는 '생성자'만 있으면 같은 학습 원리를 적용할 수 있다는 더 일반적인 틀이 존재한다는 것이다.

자주 묻는 질문

관련 AI 소식