Provably adaptive sampling with uniform and remasking discrete diffusion models

인공지능의 세계에서는 일관된 이야기를 쓰는 것부터 실제적인 단백질 구조를 생성하는 것에 이르기까지, 컴퓨터가 새로운 것을 창조하도록 가르치기 위한 끊임없는 경주가 이어지고 있습니다. 수년 동안 텍스트나 데이터 시퀀스를 처리하는 데 있어 지배적인 방법은 텍스트의 앞부분에 나온 모든 단어를 바탕으로 다음 단어를 예측하는 단계별 접근 방식이었으며, 이는 마치 사람이 문장을 한 단어씩 읽는 것과 유사합니다. 이 순차적 방식은 효과적이긴 하지만, 문장의 여러 부분을 동시에 작업할 수 없기 때문에 속도가 느립니다. 이보다 더 빠르고 새로운 대안인 '이산 확산(discrete diffusion)'이 등장했습니다. 이 방식은 처음부터 시퀀스를 구축하는 대신, 무작위로 뒤섞인 데이터 뭉치에서 시작하여 노이즈를 점진적으로 정제하고 깨끗하게 만들어 의미 있는 패턴을 도출해 냅니다. 이 접근 방식의 묘미는 데이터의 많은 부분을 동시에 업데이트할 수 있다는 점이며, 이는 훨씬 빠른 생성 경로를 제공합니다. 그러나 이 방법이 실세계에서 유용해지려면 효율적이어야 합니다. 만약 노이즈를 제거하는 과정에 너무 많은 단계가 소요된다면, 속도 측면의 이점은 사라지고 모델은 대규모 작업을 수행하기에 비실용적인 것이 됩니다.

이러한 확산 모델의 핵심 과제는 데이터에 도입되는 '노이즈'를 어떻게 다루느냐에 달려 있습니다. 명확한 문장을 가져와 일부 단어를 엉뚱한 말로 교체하거나 가리는(masking) 시스템을 상상해 보십시오. 새로운 텍스트를 생성하기 위해 모델은 이 과정을 역으로 수행하는 법, 즉 오염된 단어로부터 원래의 단어를 추측하는 법을 배워야 합니다. 오랫동안 연구자들은 이 역과정의 속도가 시스템 내의 총 단어 수나 기호의 수, 즉 '차원(dimension)'에 크게 의존한다고 믿었습니다. 만약 문장에 천 개의 위치가 있다면, 기존 이론에 따르면 실제 문장이 얼마나 단순하거나 복잡한지와 관계없이 모델은 이를 정제하기 위해 대략 천 번의 단계를 거쳐야 했습니다. 크기에 따른 이러한 선형적 의존성은 매우 구조적이고 예측 가능한 데이터일지라도 컴퓨터가 완전히 무작위적인 노이즈를 다룰 때와 똑같이 힘들게 작업해야 함을 의미했으며, 이는 병렬 처리의 이점을 사실상 무효화했습니다.

펜실베이니아 대학교의 연구팀은 이제 이 가설에 도전하며, 느린 속도가 균일한 확산 방법 자체의 근본적인 결함이 아니라, 정제 과정이 수행되는 방식의 결과임을 증명했습니다. 그들은 모델이 초기의 잠재적으로 잘못된 결정에 갇혀 있지 않고, 진행 과정에서 자신의 실수를 스스로 수정할 수 있도록 하는 새로운 샘플링 전략을 개발했습니다. 그들의 연구는 샘플을 생성하는 데 필요한 단계 수가 단순히 어휘의 크기나 시퀀스의 길이에 의해 결정되는 것이 아니라, 생성되는 데이터의 내부 구조에 의해 결정된다는 것을 보여줍니다. 만약 데이터가 서로 의존하는 부분이 있는 단순하고 예측 가능한 패턴을 가지고 있다면, 모델은 이전에 생각했던 것보다 훨씬 적은 단계만으로도 데이터를 생성할 수 있습니다.

연구진은 두 가지 특정 노이즈 프로세스에 집중했습니다. 하나는 토큰을 다른 유효한 토큰으로 무작위로 균일하게 교체하는 것이고, 다른 하나는 토큰을 가리고 모델이 확신이 없을 때 다시 가리거나 해제할 수 있게 하는 것입니다. 과거에는 이러한 과정을 역전시키기 위해 널리 채택된 '타우-리핑(tau-leaping)'과 같은 표준 알고리즘이 균일한 프로세스에 대해 비효율적이라는 것이 밝혀졌습니다. 이러한 구식 방법들은 종-종 데이터에 대한 단 한 번의 패스(pass)를 수행하며, 변경 사항이 나머지 시퀀스와 일치하는지 확인하지 않은 채 여러 위치를 동시에 업데이트하곤 했습니다. 만약 모델이 초기에 오류를 범한다면, 그 실수는 지속되어 이후의 모든 단계에 영향을 미치게 되며, 이는 이를 수정하기 위해 더 많은 단계를 필요로 하는 높은 오류율로 이어졌습니다. 이 논문에서 소개된 새로운 접근 방식은 '리브-원-아웃(leave-one-out, 하나를 제외함)' 전략을 사용합니다. 모델은 단일 토큰을 예측하기 위해 전체 시퀀스를 보는 대신, 특정 토큰이 제거되었을 때 나머지 시퀀스가 어떻게 보이는지를 고려합니다. 이를 통해 모델은 각 위치에 대해 더 정보에 입각하고 독립적인 업데이트를 병렬로 수행할 수 있으며, 결정적으로, 이후의 업데이트를 통해 이전의 예측이 틀렸음이 드러날 경우 선택을 수정할 수 있게 됩니다.

이 정교한 방법을 사용하여 연구진은 샘플을 생성하는 데 드는 계산 비용이 데이터의 서로 다른 부분들이 서로 얼마나 의존적인지에 대한 척도에 의해 지배된다는 것을 보여주었습니다. 기술적인 용어로, 그들은 이 효율성을 전체 시퀀스에 걸친 공유 정보의 양을 정량화하는 '이중 총 상관관계(dual total correlation)'라는 개념과 연결했습니다. 명확한 문법을 가진 문장이나 특정 접힘 패턴을 가진 단백질과 같이 고도로 구조화된 데이터셋의 경우, 부분 간의 관계가 긴밀하게 제약되어 있기 때문에 이 척치는 작습니다. 새로운 분석은 이러한 데이터에 대해 샘플을 생성하는 데 필요한 단계 수가 전체 위치의 수가 아니라, 이러한 구조적 복잡성에 따라 스케일링된다는 것을 증명합니다. 이는 엄격한 문법 규칙을 따르는 길고 복잡한 문장의 경우, 기반 구조가 단순하다면 모델이 짧은 문장만큼이나 빠르게 생성할 수 있음을 의미합니다. 이 논문은 이러한 효율성 향상이 단순히 운 좋은 관찰이 아니라 실제적인 수학적 증명임을 입증하며, 이전의 제한 사항들이 확산 프로세스 자체가 아니라 정제 알고리즘의 선택 때문이었음을 확립했습니다.

이러적인 이론적 발견을 검증하기 위해 연구진은 실제 세계의 구조를 모방하도록 설계된 합성 데이터에 대한 수치 실험을 수행했습니다. 그들은 이진 시퀀스가 마르코프 체인 패턴(다음 비트가 이전 비트에 의존하는 형태)을 따르는 환경에서 새로운 샘플러를 기존의 표준 방법들과 비교 테스트했습니다. 이 테스트에서 새로운 방법은 전통적인 접근 방식들을 일관되게 능가하였으며, 단계 수를 매우 낮게 유지하더라도 낮은 오류율을 유지했습니다. 결과에 따르면 기존 방법들은 데이터의 차원이 증가함에 따라 어려움을 겪었지만, 새로운 방법은 데이터의 크기가 아닌 데이터 자체의 내재적 예측 가능성에 따라 성능이 유지되었습니다. 또한 그들은 데이터가 특정 패턴의 제한된 집합에서 나오는 시나리오인 이진 문자열 혼합물에 대해서도 이 방법을 테스트했습니다. 여기서도 새로운 샘플러는 밑바닥 분포의 저차원적 특성에 적응할 수 있음을 보여주었으며, 기존 이론이 예측한 최악의 경우보다 훨씬 적은 계산 단계만으로 높은 정확도를 달 Achieve 했습니다.

이 작업의 함의는 단순히 더 빠른 알고리즘을 넘어, 이산 확산 모델의 한계를 이해하는 방식을 근본적으로 바꾼다는 데 있습니다. 불리한 차원 의존성이 확산 프로세스의 본질적인 장벽이 아니라 해결 가능한 알고리즘 설계의 문제임을 보여줌으로써, 연구진은 더 효율적인 대규모 생성 모델의 문을 열었습니다. 이는 데이터가 고차원이면서도 매우 구조화된 자연어 처리 및 단백질 설계와 같은 응용 분야에 특히 중요합니다. 토큰의 수에 얽매이지 않고 병렬로 복잡한 시퀀스를 생성할 수 있는 능력은, 이산 확산이 조만가 자기회귀(autoregressive) 모델과 속도 및 품질 면에서 대등하거나 이를 능가할 수 있음을 시사합니다. 또한 이 연구는 모델이 중간 결정을 수정할 수 있도록 허용하는 것의 중요성을 강조하는데, 이는 인간이 글을 쓰거나 생각할 때 사용하는 반복적인 정제 과정과 유사한 특징입니다.

궁극적으로 이 연구는 생성형 AI의 효율성을 개선하기 위한 명확한 경로를 제공합니다. 이는 이산 확산이 데이터를 병렬로 생성할 수 있는 잠재력이 단순한 이론적 약속이 아니라, 올바른 도구를 사용하여 노이즈를 헤쳐 나갈 때 실현 가능한 현실임을 확인시켜 줍니다. 이 작업은 프로세스의 수학적 근사로 인해 발생하는 오류와 모델의 학습으로 인해 발생하는 오류를 분리하여, 전자가 데이터의 구조 자체에 의해 엄격하게 제어될 수 있음을 보여주었습니다. 분야가 더 크고 복잡한 모델로 나아감에 따라, 이러한 통찰력은 계산 비용이 문제의 크기에 따라 통제 불능으로 커지지 않도록 보장하는 데 필수적일 것입니다. 이 연구 결과는 이산 생성의 미래가 무차별적인 계산이 아니라, 데이터 내부의 자연스러운 질서와 의존성을 활용하는 더 스마트하고 적응적인 전략에 달려 있음을 시사합니다.