TiDAR 설명: 확산·자기회귀 하이브리드로 LLM 추론을 4~6배 가속하는 원리

핵심 메시지

• 자기회귀 LLM 추론은 메모리 병목 탓에 GPU 연산이 놀고 있는 시간이 많은데, TiDAR는 그 남는 연산으로 다음 토큰의 초안을 미리 만든다.
• TiDAR는 확산으로 여러 토큰을 한꺼번에 생성해 초안을 만들고, 자기회귀 모델로 그 초안을 단 한 번의 순전파에서 검증해 속도를 높인다.
• 거부 표집(rejection sampling) 덕분에 최종 출력은 자기회귀 표집과 수학적으로 동일해, 품질을 잃지 않고 속도만 얻는다.
• 하나의 순전파 안에서 현재 초안의 검증과 다음 단계 초안의 생성을 동시에 처리하므로 가속이 '거의 공짜'에 가깝다.
• 1.5B 규모에서 자기회귀 모델과의 품질 격차를 사실상 좁히면서 초당 처리 토큰 수를 4~6배로 늘렸다.

쉽게 이해하기

언어모델이 답을 만드는 방식은 크게 두 가지다. GPT류의 자기회귀 모델은 앞선 토큰 전체를 조건으로 다음 토큰 하나를 뽑고, 그 결과를 다시 조건에 넣어 그다음 토큰을 뽑는 식으로 한 번에 하나씩 생성한다. 원리적으로 품질이 좋지만 순차적이라 느리다. 반대로 확산 언어모델은 미래의 여러 토큰을 한꺼번에 만들어 점진적으로 다듬는다. 빠르지만 각 토큰을 주변부(marginal) 분포에서 따로 표집하기 때문에 토큰 사이의 상호작용을 반영하지 못해 품질이 떨어지는 경향이 있다.

속도를 위한 기존 기법으로 추측 디코딩(speculative decoding)이 있다. 작고 빠른 모델이 여러 토큰을 미리 제안하면, 큰 모델이 그 제안을 한 번의 순전파로 병렬 검증한다. 인과적 어텐션 덕분에 검증은 병렬로 가능하다. 다만 별도의 작은 모델을 돌리는 비용이 있어, 제안이 자주 틀리면 오히려 손해가 될 수 있다.

TiDAR의 핵심 통찰은 자기회귀 디코딩 중 GPU가 메모리 병목으로 놀고 있다는 점이다. 그 남는 연산 용량을 써서 확산 방식으로 다음 단계의 초안을 만들면, 추측 디코딩의 초안 생성 비용을 사실상 공짜로 얻을 수 있다. 즉 하나의 하이브리드 모델이 확산으로 초안을 만들고 자기회귀로 검증한다.

동작을 보면 각 생성 단계에서 토큰을 접두부, 직전 단계에서 제안된 토큰, 다음 단계용으로 미리 뽑은 초안의 세 구획으로 나눈다. 접두부의 KV 캐시는 재사용하고, 제안된 토큰은 거부 표집으로 자기회귀 결과와 맞는지 검증한다. 초안 일부가 거부될 수 있으므로 남는 연산으로 '가능한 모든 수용 결과'에 대한 다음 초안을 어텐션 마스킹으로 동시에 계산해 둔다. 그래서 어느 쪽으로 검증되든 다음 단계에 쓸 유효한 초안이 항상 준비돼 있다.

결과적으로 TiDAR는 1.5B 규모에서 자기회귀 모델과 품질 격차를 좁히면서 추측 디코딩보다 높은 처리량을 냈고, Dream·LLaDA 같은 확산 모델보다 효율과 품질 모두 앞섰다. 최종 표집이 자기회귀와 수학적으로 동등하기 때문에 유일한 대가는 같은 모델을 확산과 자기회귀 두 목적으로 함께 학습해야 한다는 점뿐이다.

주요 인사이트

• 인과적 어텐션은 학습을 병렬화해 효율을 높이지만, 중간 처리 단계에서도 항상 과거만 보게 강제해 이론적으로 가능한 어텐션 패턴을 크게 제약한다는 트레이드오프가 숨어 있다.
• 확산 모델의 품질 저하는 각 토큰을 주변부 분포에서 독립적으로 표집해 토큰 간 상호작용을 무시하는 데서 온다. TiDAR가 자기회귀 검증을 붙이는 이유가 바로 이 결합(joint) 품질을 되살리기 위함이다.
• 가속의 열쇠는 '남는 연산은 공짜'라는 관찰이다. 플래시 어텐션처럼 메모리를 연산으로 맞바꾸는 기법과 결합하면, 메모리 병목으로 비어 있던 연산을 초안 계산에 돌려쓸 수 있다.
• 한때 생성 모델링에서 밀려났던 BERT식 마스킹(마스크 언어모델링)이 확산 초안 생성이라는 형태로 다시 쓰이는 점도 흥미롭다.

자주 묻는 질문

관련 AI 소식