확장 플로우(Expanding Flows): 길이가 변하는 데이터를 몇 단계 만에 생성하는 새 프레임워크

핵심 메시지

• 기존 플로우 매칭과 플로우 맵은 픽셀 수나 시퀀스 길이가 이미 정해진 '고정 캔버스' 위에서만 동작해, 생성 도중에 길이를 늘릴 수 없다는 한계가 있었다.
• 연구진은 생성 과정을 새 좌표를 끼워 넣는 '확장(expand) 연산'과 기존 방식대로 흘려보내는 '이동(transport) 연산'으로 분해해, 차원이 커지는 흐름을 정식화했다.
• 새로 삽입된 좌표에는 삽입 시점부터 흐르기 시작하는 '지역 시계'를 부여하고, 아직 삽입되지 않은 좌표의 속도는 0으로 두어 학습을 성립시켰다.
• 텍스트 같은 이산 데이터에서는 삽입 개수를 푸아송 대신 이항분포로 예측해, 큰 폭으로 건너뛸 때 시퀀스 길이가 폭발하지 않도록 막았다.
• 분자 구조 생성, 분자 그래프 생성, 가변 길이 언어 모델링 세 과제에서 단계 수를 크게 줄이고도 기존 확산·플로우 모델과 비슷하거나 더 나은 품질을 보였다.

쉽게 이해하기

이 강연은 펜실베이니아대학교 소속으로 하버드대에서 방문 연구 중인 소피아 탕(Sophia Tang) 연구자가 마이크로소프트 리서치에서 발표한 내용이다. 주제는 이미지·분자·텍스트를 만들어내는 생성 모델의 기본 틀인 플로우 매칭(flow matching)과 플로우 맵(flow map)을 '크기가 변하는 데이터'로 확장하는 방법이다. 발표자는 논문의 아카이브 공개를 앞두고 핵심 아이디어와 실험 결과를 공유했다.

플로우 매칭은 잡음 분포에서 뽑은 점과 실제 데이터에서 뽑은 점 사이를 잇는 연속적인 경로를 정의하고, 그 경로를 따라가는 속도장을 학습하는 방식이다. 정확한 샘플을 얻으려면 경로를 잘게 쪼개 여러 번 적분해야 해서 느리다는 문제가 있었고, 이를 해결하려고 나온 것이 한두 번의 큰 점프로 끝내는 플로우 맵이다. 다만 두 방식 모두 시작부터 끝까지 데이터의 차원이 고정돼 있다는 전제를 깔고 있다.

발표자가 지적한 문제의식은 여기서 출발한다. 자기회귀 언어 모델은 토큰을 계속 이어 붙이면 되므로 길이가 자유롭지만, 플로우 기반 모델은 생성 도중 시퀀스를 늘릴 방법이 없다는 것이다. 연구진의 해법은 하나의 흐름을 두 단계로 쪼개는 것이다. 먼저 확장 연산이 잡음 좌표를 새로 만들어 상태 공간의 차원을 키우고, 그다음 이동 연산이 늘어난 공간 위에서 기존 플로우 맵과 똑같이 동작한다.

새 좌표를 어디에 넣을지는 과제에 따라 달라진다. 시퀀스 끝에 이어 붙이는 방식, 중간 위치에 끼워 넣는 방식, 하나의 노드를 여러 자식 노드로 펼치는 방식 등이 가능하다. 좌표마다 삽입 시점이 다르기 때문에 연구진은 각 좌표에 삽입 순간부터 0에서 시작하는 지역 시계를 부여했고, 삽입 전 좌표의 속도는 0으로 정의해 전체 흐름이 잘 정의되도록 했다. 삽입 자체는 불연속적인 점프이므로 전체 과정은 조각별 결정적 마르코프 과정으로 정식화했다.

실험은 세 갈래로 진행됐다. 무거운 원자만 있는 거친 구조에서 수소 원자까지 채워 넣는 분자 형태 생성에서는 20단계 안팎, 증류를 거치면 한두 단계만으로도 확산 모델급 품질이 나왔다. 원자 개수를 미리 모르는 분자 그래프 생성에서는 삽입 개수까지 학습해 기존 이산 플로우 모델을 앞섰고, 언어 모델링 실험에서는 길이가 변하는 더 어려운 문제를 풀면서도 고정 길이 모델과 비슷한 생성 품질을 유지했다.

주요 인사이트

• 확장 연산을 항등함수로 두면 기존의 고정 차원 플로우 맵이 이 프레임워크의 특수한 경우로 그대로 복원된다. 새 방법이 기존 방식을 대체하는 것이 아니라 감싸 안는 구조라는 뜻이다.
• 작은 점프 두 번과 큰 점프 한 번의 삽입 개수가 어긋나면 차원이 맞지 않는데, 연구진은 부족한 좌표의 속도를 0으로 채워 손실을 계산한다. 청중석에서는 이 방식이 원점에서 먼 좌표를 더 크게 벌주는 편향을 만들 수 있다는 지적이 나왔고, 발표자는 학습 초반에 고정 차원 플로우 맵을 먼저 익히게 하는 단계적 학습으로 대응한다고 답했다.
• 차원을 늘리는 것만 다루고 줄이는 삭제 연산은 아직 포함되지 않았다. 발표자는 삭제를 함께 허용하는 것을 향후 과제로 꼽았다.
• '최대 길이를 크게 잡고 빈 자리를 널(null) 상태로 두면 되지 않느냐'는 질문에 발표자는, 학습 때는 실제로 패딩과 최대 길이를 쓰지만 확장 연산은 어디에 몇 개를 넣을지를 학습 가능한 분포로 다룰 수 있게 해 설계 공간을 넓혀준다고 설명했다.
• 단백질 적용 여부에 대해서는, 서열 생성이라면 지금 기법을 거의 그대로 쓸 수 있지만 구조 생성은 제약 조건이 많아 추가 실험이 필요하다는 신중한 답이 돌아왔다.

자주 묻는 질문

관련 AI 소식