GN⁺: Consistency LLM: LLM을 병렬 디코더로 전환해 추론 속도 3.5배 향상

(hao-ai-lab.github.io)

Consistency Large Language Models (CLLM)

LLM들은 전통적으로 하나의 토큰을 연속적으로 디코딩하는 방식으로 동작하는 시퀀셜 디코더로 여겨져 왔음
본 블로그에서는 사전학습된 LLM들이 효율적인 병렬 디코더로 쉽게 전환될 수 있음을 보여줌
Consistency Large Language Models(CLLM)은 추론 단계마다 n개 토큰 시퀀스를 효율적으로 디코딩하여 추론 지연시간을 감소시킬 수 있는 새로운 병렬 디코더 계열임
CLLM은 사람이 단어를 하나씩 말하기 전에 머릿속으로 완전한 문장을 형성하는 인지 과정을 모방하며, 사전학습된 LLM을 미세조정(fine-tuning)하는 것만으로 효과적으로 학습될 수 있음
CLLM은 무작위로 초기화된 n개 토큰 시퀀스를 가능한 적은 단계로 자기회귀(autoregressive, AR) 디코딩 결과와 동일하게 매핑하도록 병렬 디코딩을 수행하는 방식으로 학습됨
실험 결과에 따르면, CLLM은 기존 AR 디코더 대비 2.4배에서 3.4배의 생성 속도 향상을 보였으며, Medusa2나 Eagle과 같은 빠른 추론 기법들과 견줄만하거나 더 우수한 성능을 보였음
CLLM은 추가적인 메모리 비용 없이도 이러한 성능 향상을 달성할 수 있음

LLM은 AR 디코딩 방식으로 토큰을 하나씩 생성하기 때문에 긴 응답을 위해서는 높은 지연시간이 발생함
Jacobi 디코딩은 Jacobi 및 Gauss-Seidel 비선형방정식 해법에서 유래되었으며, greedy 샘플링을 사용한 AR 생성과 동일함이 증명되었음
Jacobi 디코딩은 순차적 생성 과정을 Jacobi 반복에 기반한 n개 변수를 가진 n개 비선형방정식 시스템으로 재구성하여 병렬 처리가 가능하도록 함
각 반복 단계에서는 1개 이상의 올바른 토큰을 예측할 수 있어 AR 디코딩을 잠재적으로 가속화 시킬 수 있음
그러나 실제로는 AR로 학습된 LLM은 선행 토큰에 오류가 있을 때 올바른 토큰을 거의 생성할 수 없기 때문에 대부분의 Jacobi 반복에서 n개 토큰 시퀀스에 대해 하나의 수정만 얻게 되어 더 긴 Jacobi 궤적(trajectory)을 생성하게 됨
Lookahead 디코딩이나 speculative 디코딩은 이러한 Jacobi 디코딩의 비효율성을 완화하려 하지만 추론시 추가 메모리 비용이 발생하는 반면, CLLM은 그렇지 않음

CLLM 학습은 크게 Jacobi 궤적 준비와 consistency 및 AR 손실 최적화 두 부분으로 구성됨
Jacobi 궤적 준비 단계에서는 전체 응답 시퀀스 l개 토큰이 생성될 때까지 n개씩 잘라서 순차적으로 Jacobi 디코딩을 수행하며, 각 궤적에서 생성된 시퀀스를 하나의 데이터 항목으로 간주함
학습시에는 consistency 손실과 AR 손실을 함께 최적화하는데, consistency 손실은 여러 토큰을 한번에 예측하도록 보장하며 AR 손실은 CLLM이 대상 LLM에서 벗어나지 않도록 하여 생성 품질을 유지하도록 함
Global consistency(GC) 손실은 Jacobi 궤적의 임의 지점과 고정점 사이의 거리를 최소화하여 CLLM이 Jacobi 궤적의 어느 지점에서든 고정점을 예측하도록 장려함
Local consistency(LC) 손실은 Jacobi 궤적의 인접 상태가 동일한 출력을 내도록 유도함
AR 손실은 대상 LLM의 생성 결과에 기반하여 전통적인 AR 손실을 포함하며, 목적은 CLLM이 대상 LLM의 분포에서 벗어나지 않도록 하는 것임

실험에는 Spider(text-to-SQL), Human-Eval(파이썬 코드 완성), GSM8k(수학) 등의 특화된 도메인 작업과 MT-bench와 같은 폭넓은 개방형 대화 챌린지가 포함되었음
CLLM은 대상 모델 대비 가장 큰 속도 향상을 보였으며, 추론시 추가 비용 없이 Medusa2와 유사하거나 더 나은 속도 향상을 달성했음
MT-bench에서 CLLM은 Medusa2와 결합했을 때와 거의 동일한 속도 향상을 달성했지만, 더 높은 적응성과 메모리 효율성을 제공함
CLLM의 미세조정 비용은 적당한 수준이며, 데이터셋 크기가 큰 경우에도 Jacobi 궤적 생성에 데이터셋의 10% 정도만 사용해도 2.5배 정도의 속도 향상을 얻을 수 있었음
CLLM은 fast forwarding 현상을 통해 여러 연속 토큰을 한 번의 Jacobi 반복에서 올바르게 예측할 수 있음
CLLM은 선행 토큰에 오류가 있어도 올바른 토큰을 미리 예측하고 변경되지 않도록 유지하는 stationary 토큰 능력을 보여줌
CLLM은 학습을 통해 연어(collocation)와 같은 핵심 언어 개념을 습득하며, 이를 통해 Jacobi 궤적의 어느 지점에서도 구조를 유추하고 반복 단계를 최소화하기 위해 여러 단어를 동시에 예측할 수 있게 됨

CLLM은 기존 LLM의 AR 디코딩 방식이 가진 긴 지연 시간 문제를 Jacobi 디코딩을 활용해 효과적으로 해결한 것으로 보임. 특히 추가 메모리 비용 없이 병렬화된 디코딩을 통해 속도 향상을 달성한 점이 인상적임
CLLM의 학습 방법은 기존 LLM을 consistency 손실을 통해 미세조정하는 것으로 비교적 간단해 보이지만, 이를 통해 언어의 중요한 특성 중 하나인 연어(collocation)을 학습하여 병렬 디코딩 성능을 크게 향상시켰다는 점에서 의미가 있어 보임
다만 CLLM이 greedy 샘플링을 가정하고 있어서 더 다양한 decoding 전략에서도 잘 동작할 수 있을지는 추가 연구가 필요해 보임. 또한 현재는 영어에 한정된 실험 결과라 다양한 언어로의 일반화 가능성 역시 검증이 필요할 것 같음
CLLM은 LLM의 응답 속도를 빠르게 만드는 방법으로 실용적인 접근이라고 생각됨. 웹 검색이나 챗봇 등 실시간성이 요구되는 태스크에 잘 적용될 수 있을 것 같음
개인적으로 CLLM의 consistency 학습 방식이 GPT 등의 LLM 외에도 이미지 생성 모델이나 음성 합성 모델 등 다른 생성 모델에도 적용될 수 있지 않을까 기대됨. CLLM의 아이디어가 앞으로 다양한 생성 모델의 효율성 향상에 기여할 수 있기를 바람

자유로운 드로잉 수업에서 시간 제한을 두고 반복적으로 그리는 연습을 하면 그림 실력이 크게 향상됨. 급하게 그리는 연습이 오히려 비율과 윤곽을 정확히 잡는 능력을 키워줌.
Jacobi decoding은 greedy autoregressive decoding과 동일하지만, 실제로는 반복을 피하고 지나치게 일반적인 응답을 피하기 위해 샘플링 온도를 0 이상으로 설정해야 함.
곧 우리는 모델 학습이 필요 없다는 것을 깨닫게 될 것임. 좋은 인덱싱과 샘플링만 필요함. LLM은 기본적으로 훌륭한 NLP 인터페이스를 갖춘 데이터셋의 DB와 동등함.
CLLM 논문의 inference 성능 향상이 주목할 만함. Fine-tuning 비용이 합리적이고(사전 학습 비용의 약 0.01%) 성능 향상도 꽤 일관됨.
Jacobi trajectories는 높은 온도 경로를 배제하므로 데이터 검색에는 긍정적이지만 창의성 극대화에는 부정적일 수 있음.
LLM이 동일한 프롬프트에 대해 결정론적으로 응답하지 않는 이유에 대해 "AI 전문가에게 물어보기" 기능이 있으면 좋겠음.
Groq(https://groq.com/)의 속도 향상이 이 방법 때문인지 다른 방법을 사용하는지 궁금함.