- 긴 응답을 만들 때 토큰을 하나씩 생성하는 자기회귀(AR) 디코딩이 지연의 핵심 병목이며, CLLM은 이를 n개 토큰 단위의 병렬 디코딩으로 줄이려는 접근임
- Consistency Large Language Models(CLLMs) 는 사전학습 LLM을 미세조정해 임의의 n토큰 상태를 AR greedy 결과와 같은 고정점으로 빠르게 매핑하도록 학습함
- Jacobi decoding은 이론적으로 AR greedy 생성과 같은 결과에 수렴하지만, 기존 LLM에서는 평균 1.05배 수준의 제한적 속도 향상에 그쳐 실용성이 낮았음
- CLLM은 Spider, Human-Eval, GSM8k, MT-bench 실험에서 2.4배~3.4배 생성 속도 향상을 보였고 Medusa2, Eagle 같은 빠른 추론 기법과 비슷하거나 더 나은 수준으로 평가됨
- 추론 시 보조 모델 컴포넌트나 대상 모델 아키텍처 변경이 필요 없어, 속도 개선과 메모리 효율을 함께 노릴 수 있음
CLLM이 겨냥하는 디코딩 병목
- LLM은 프로그래밍, 법률, 건강 조언 등 다양한 영역에서 쓰이지만, 추론 단계에서는 보통 자기회귀(AR) 디코딩으로 토큰을 하나씩 생성함
- 응답이 길어질수록 토큰별 생성 방식 때문에 지연이 커지고, 사용자에게 체감되는 대기 시간이 늘어남
- 기존의 빠른 추론 방식은 여러 토큰을 한 번에 만들기 위해 아키텍처 수정, 보조 컴포넌트, 초안 모델을 요구하는 경우가 많음
Jacobi decoding의 작동 방식과 한계
- Jacobi decoding은 Jacobi 및 Gauss-Seidel 고정점 반복에서 출발한 방식이며, greedy 디코딩에서는 AR 생성과 동일함이 증명됨
- 순차 생성을 n개 변수와 n개 비선형 방정식의 시스템으로 바꾸고, Jacobi 반복으로 병렬 계산할 수 있게 만듦
- 구체적인 흐름은 다음과 같음
- 입력 프롬프트에서 다음 n개 토큰을 무작위로 추측함
- 프롬프트와 n토큰 시퀀스를 LLM에 넣어 반복적으로 업데이트함
- 시퀀스가 더 이상 변하지 않으면 고정점(fixed point) 에 도달함
- 최종 n토큰 시퀀스는 greedy 전략에서 AR 디코딩이 생성했을 출력으로 수렴함
- 초기 무작위 추측에서 최종 AR 생성 결과까지 이어지는 경로를 Jacobi 궤적(Jacobi trajectory) 이라고 부름
- 기본 Jacobi decoding은 실제 LLM에서 AR 디코딩 대비 평균 1.05배 정도의 제한적 속도 향상만 보임
- AR로 학습된 LLM은 앞 토큰에 오류가 있으면 뒤 토큰을 올바르게 예측하기 어려움
- 대부분의 Jacobi 반복이 n토큰 시퀀스에서 하나의 토큰만 교정해 궤적이 길어짐
- Lookahead decoding과 speculative decoding은 Jacobi decoding 및 기존 AR 디코딩의 비효율을 줄이려 하지만, 추론 시 추가 메모리 비용이 생김
Consistency LLM의 학습 목표
- CLLM은 사전학습 LLM이 Jacobi 궤적의 임의 지점에서 고정점으로 일관되게 이동하도록 적응시키는 방식임
- 이 목표는 확산 모델 가속 기법인 consistency models의 목표와 유사함
- 대상 모델에서 수집한 Jacobi 궤적을 이용해, Jacobi 반복 중 단일 단계 수렴을 장려하는 손실로 모델을 학습함
- 각 대상 모델을 CLLM으로 바꾸는 학습은 두 부분으로 구성됨
- Jacobi 궤적 준비
- 각 프롬프트에 대해 n토큰 단위의 모든 절단 구간마다 Jacobi decoding을 순차 수행함
- 전체 응답 시퀀스는 연속된 고정점들을 이어 붙인 형태가 됨
- 궤적 중 생성된 각 시퀀스가 하나의 학습 데이터 항목으로 계산됨
- 긴 응답에서 전체 긴 입력을 평가하는 속도 저하를 피하기 위해 n토큰 절단을 사용함
- Consistency loss와 AR loss 공동 최적화
- consistency loss는 여러 토큰을 한 번에 예측하도록 유도함
- AR loss는 CLLM이 대상 LLM의 분포에서 벗어나지 않게 해 생성 품질 유지를 돕음
- Jacobi 궤적 준비
손실 함수 구성
- 대상 LLM을
p, CLLM을qθ로 두며,qθ는p의 파라미터로 초기화됨 - Global consistency(GC) loss는 Jacobi 궤적의 임의 상태
y를 입력으로 넣었을 때 CLLM이 고정점y*를 출력하도록 유도함 - Local consistency(LC) loss는 Jacobi 궤적에서 인접한 상태
y(j)와y(j+1)이 같은 출력을 내도록 맞춤 - 분포 간 거리
D(·||·)는 GKD 방법에서 논의된 선택지를 따르며, 본 실험에서는 주로 forward KL을 사용함 - AR loss는 대상 LLM
p의 생성 결과l을 기반으로 전통적인 자기회귀 손실을 적용해 대상 LLM 분포에서 벗어나는 것을 막음 - 전체 학습 손실은 consistency loss와 가중치
w가 붙은 AR loss의 합으로 구성됨
실험 설정과 결과
- 실험은 세 가지 특화 도메인 작업과 하나의 오픈 도메인 대화 벤치마크를 포함함
- Spider: text-to-SQL
- Human-Eval: Python 코드 완성
- GSM8k: 수학
- MT-bench: 오픈 도메인 대화
- 작업에 따라 대상 모델로 fine-tuned coder LLM, Deepseek-coder-7B-instruct, LLaMA-2-7B, ABEL-7B-001을 사용함
- 학습과 평가는 모두 NVIDIA A100 40GB 서버에서 수행됨
- 특화 도메인에서는 원래 대상 모델, Medusa2, speculative decoding 등 기준선과 비교했을 때 CLLM이 가장 큰 속도 향상을 보임
- MT-bench에서는 ShareGPT 데이터셋으로 LLaMA2-7B에서 학습한 CLLM이 lookahead decoding과 결합했을 때 Medusa2와 비슷한 속도 향상을 달성함
- MT-bench 점수도 비교 가능한 수준임
- CLLM은 대상 모델의 원래 아키텍처 수정이 필요 없음
- 보조 컴포넌트가 없어 메모리 효율성이 높음
학습 비용
- CLLM 미세조정 비용은 중간 수준으로 제시됨
- LLaMA-7B는 Spider 데이터셋에서 약 100만 토큰만 통과해 3.4배 속도 향상을 달성함
- CodeSearchNet-Python처럼 데이터셋이 큰 경우, 데이터셋의 10% 만 Jacobi 궤적 생성에 사용해 약 2.5배 속도 향상을 얻음
- 총 토큰 수는 다음 식으로 추정함
- 평균 프롬프트당 궤적 수 × 평균 궤적 길이 × 프롬프트 수
- 데이터셋별 추정 학습 비용은 다음과 같음
- Spider: 200만 토큰, 사전학습 비용의
< 0.01% - CodeSearchNet-Python: 1억 토큰, 사전학습 비용의
~0.1% - GSM8K: 1,000만 토큰, 사전학습 비용의
~0.01% - ShareGPT: 2억 토큰, 사전학습 비용의
~0.2%
- Spider: 200만 토큰, 사전학습 비용의
빠른 전진과 고정 토큰
- 대상 LLM은 일반적으로 한 번의 Jacobi 반복에서 올바른 토큰을 하나만 생성함
- CLLM에서는 한 번의 Jacobi 반복에서 여러 연속 토큰을 올바르게 예측하는 빠른 전진(fast forwarding) 현상이 나타남
- 대상 LLM에서는 앞서 올바르게 생성된 토큰도 이후 반복에서 잘못된 토큰으로 바뀌는 경우가 있음
- CLLM은 앞 토큰에 오류가 있어도 올바른 토큰을 미리 예측하고, 그 토큰을 이후에도 유지하는 능력을 보임
- 이런 토큰을 고정 토큰(stationary tokens) 이라고 부름
- 빠른 전진과 고정 토큰은 CLLM의 Jacobi decoding이 더 빨리 수렴하게 만들어 생성 속도 향상에 기여함
언어적 패턴 학습
- CLLM은 학습을 통해 연어(collocations) 라는 언어 개념을 습득하는 것으로 관찰됨
- 연어는 무작위 우연보다 더 자주 함께 등장하는 단어나 용어의 연속을 뜻함
- 자연어와 코드 모두에서 연어 예시가 존재함
- 자연어:
talk to,remind … of … - 동사+명사 구조:
make a decision,catch a cold - 도메인별 문법 구조:
SELECT … FROM …,if … else
- 자연어:
- consistency generation 목표는 CLLM이 Jacobi 궤적의 어느 지점에서도 이런 구조를 추론하게 해, 반복 단계를 줄이기 위해 여러 단어를 동시에 예측하도록 장려함
자료와 코드
- 자세한 내용은 논문에서 확인할 수 있음
- 구현은 코드베이스로 제공됨
- CLLM checkpoints도 공개됨