확산 모델 샘플링 속도의 수학, 고차 ODE 솔버 수렴 분석 연구

핵심 메시지

• 확산 모델은 그림에 잡음을 조금씩 더해 완전한 잡음으로 만든 뒤, 그 과정을 거꾸로 되돌려 새 이미지를 만든다.
• 되돌리는 과정을 확률 미분방정식 대신 상미분방정식으로 풀면 수천 번 걸리던 단계를 열 번 안팎으로 줄일 수 있다.
• 같은 정보를 쓰더라도 갱신하는 방식만 바꾸면 훨씬 정확한 해를 얻는다는 것이 고차 솔버의 핵심이다.
• 연구진은 생성 결과와 실제 분포의 차이를 초기화 오차, 점수 함수 오차, 이산화 오차 세 항으로 나눠 정리했다.
• 기존 이론과 달리 이 결과는 학습 과정에서 실제로 얻을 수 있는 값만 사용하며, 오차가 시간에 따라 지수적으로 폭증하지 않는다.

쉽게 이해하기

마이크로소프트 리서치 세미나에 초청된 MIT 수학과 연구자는 고양이 사진 여러 장으로 새로운 고양이 사진을 만들어내는 문제에서 이야기를 시작했다. 이상적으로는 가우시안 잡음을 실제 이미지 분포로 옮겨주는 함수를 하나 찾으면 되지만, 이미지나 음성, 영상은 차원이 지나치게 높고 그 분포가 매끄러운 형태를 띠지 않아 그런 함수를 곧바로 배우는 것은 불가능에 가깝다.

확산 모델은 이 문제를 우회한다. 먼저 원본 이미지에 잡음을 조금씩 더해 가우시안 잡음으로 만드는 정방향 과정을 정의하고, 그 궤적을 거꾸로 따라가는 역방향 과정을 학습한다. 이때 역방향은 확률 미분방정식으로도, 확률 흐름을 기술하는 상미분방정식으로도 표현할 수 있는데, 발표는 후자에 집중한다. 확률 미분방정식 기반 방식은 보통 100번에서 1만 번까지 단계를 밟아야 해서 실시간 서비스에는 너무 느리기 때문이다.

발표자는 미분방정식을 수치적으로 푸는 고전적인 방법을 예로 들어 고차 솔버의 원리를 설명했다. 오일러 방법은 테일러 전개의 1차 항만 남기므로 한 걸음의 오차가 간격의 제곱에 비례한다. 반면 2차 방법은 같은 두 지점의 정보를 쓰면서도 오차를 간격의 세제곱 수준으로 줄인다. 청중이 계산량이 늘어나는 것 아니냐고 묻자, 발표자는 두 방법이 쓰는 정보의 개수는 같고 갱신 공식만 다르다는 점을 강조했다. 정방향 과정에서 이미 수집한 같은 데이터로 훨씬 정확한 곡선을 그릴 수 있다는 것이다.

핵심 성과는 이렇게 만든 빠른 샘플러가 얼마나 정확한지를 수학적으로 보장한 것이다. 연구진은 학습된 분포와 실제 분포 사이의 총변동 거리를 세 항의 합으로 나눴다. 첫째는 잡음을 충분히 오래 더하면 초기 분포가 가우시안에 지수적으로 가까워지므로 매우 작다. 둘째는 신경망이 점수 함수를 얼마나 잘 배웠는지를 나타내는 항이고, 셋째는 이산적인 단계를 밟으며 쌓이는 오차다. 세 항 모두 학습 과정에서 실제로 측정할 수 있는 양으로 표현된다.

이론적 난관도 있었다. 고차 솔버의 오차 분석은 보통 함수가 매우 매끄럽다는 가정 위에서 성립하는데, 실제로 흐름을 만드는 것은 학습된 신경망이라 그런 매끄러움을 보장할 수 없다. 연구진은 참값 점수 함수로 만든 가상의 흐름을 중간에 끼워 넣어, 매끄러움이 필요한 부분을 신경망이 아니라 참값 흐름 쪽으로 옮기는 방식으로 이 문제를 풀었다. 또한 신경망의 2계 도함수가 실제로 다루기 좋은 크기인지 MNIST와 패션 MNIST 데이터셋에서 직접 계산해 확인했다.

주요 인사이트

• 성능 향상의 출처가 더 많은 계산이 아니라 더 나은 갱신 공식이라는 점이 이 발표의 가장 실용적인 메시지다. 같은 관측값으로 더 정확한 곡선을 그릴 수 있다면 단계 수를 줄여도 품질이 유지된다.
• 발표자는 이 아이디어가 확산 모델에만 갇힐 이유가 없다고 말했다. 어떤 흐름을 근사하는 구조라면 어디든 고차 방법을 시도해 학습이나 추론 단계를 줄일 수 있다는 것이다.
• 기존 연구들은 학습된 점수와 참 점수의 도함수까지 가깝다는 추가 가정을 요구했지만, 이 연구는 훈련 손실에서 실제로 얻어지는 제곱 오차만으로 결과를 세웠다.
• 선행 연구에서 흔히 등장하던 시간에 대한 지수적 항이 사라지고 선형보다 완만한 형태로 정리된 점도 실무적으로 의미가 크다. 단계를 늘려도 오차 보장이 급격히 무너지지 않는다는 뜻이기 때문이다.
• MNIST에서 헤시안 값 자체는 수백 규모였지만 잡음 크기에 따른 가중치를 곱하면 1~2 수준의 작은 상수로 정리됐다. 가정이 이론에만 사는 것이 아님을 수치로 확인한 셈이다.

자주 묻는 질문

관련 AI 소식