상태 전이 기반 효율적인 대형 언어 모델 추론 프레임워크

초록

본 논문은 긴 체인‑오브‑생각(Chain‑of‑Thought, CoT) 생성 시 발생하는 계산·메모리 비용을 줄이기 위해, 추론 과정을 “상태‑전이” 문제로 모델링한다. 선형 어텐션을 이용해 이전 단계의 추론 정보를 압축한 ‘추론 상태 행렬’을 유지하고, 현재 질문과 상태만을 입력으로 하여 다음 추론 단계를 생성한다. 이를 통해 어텐션 복잡도를 O(C²)에서 O(C)로 낮추고 KV‑cache 메모리를 상수 수준으로 감소시킨다. 또한, 전역 그래디언트(모멘텀)를 활용한 상태‑기반 추론 전략으로 잡음이 섞인 단계에 의한 과도한 사고(over‑thinking)를 완화한다. 실험 결과, 다양한 벤치마크와 모델 규모에서 추론 효율이 크게 향상될 뿐 아니라 정확도도 동반 상승함을 보인다.

상세 분석

이 논문은 기존 CoT 기반 추론이 “길어질수록” 비용이 급증한다는 문제점을 정확히 짚어낸다. Transformer의 소프트맥스 어텐션은 입력 길이 C에 대해 O(C²) 연산과 O(C) KV‑cache 메모리를 요구한다. 특히, 복잡한 수학·논리 문제에서 수십~수백 토큰에 달하는 CoT를 생성하면 GPU 메모리 한계와 추론 지연이 실용성을 크게 저해한다. 저자들은 이를 해결하기 위해 두 가지 핵심 아이디어를 제시한다.

첫 번째는 “상태‑전이” 모델링이다. 전체 CoT를 여러 ‘추론 단계(step)’로 구분하고, 각 단계는 (1) 언어적 유창성을 위한 텍스트와 (2) 이후 단계에 필요한 핵심 추론 정보로 구성된다고 가정한다. 단계별 텍스트는 KV‑cache에 저장되지 않고, 핵심 정보만이 선형 어텐션을 통해 ‘상태 행렬 S_t’에 압축된다. 선형 어텐션은 커널 ϕ(·)를 이용해 q·kᵀ를 외적 형태로 변환하고, 연관성을 누적합(S_t) 형태로 유지한다. 이때 연산 복잡도는 O(C)이며, 메모리 사용량은 상태 행렬 하나(고정 크기)로 제한된다.

두 번째는 Mixed Attention Module(MAM)이다. 기존 소프트맥스 어텐션(SA) 서브모듈은 현재 단계와 질의 프롬프트에만 집중하도록 KV‑cache를 제한하고, 선형 어텐션(LA) 서브모듈은 상태 행렬을 통해 과거 단계의 핵심 정보를 직접 조회한다. 두 서브모듈의 출력을 게이팅(σ·W_g)으로 조절해 초기 단계에서는 LA의 비중을 높이고, 단계가 진행될수록 SA가 점차 주도하도록 설계했다. LA 서브모듈은 LoRA 방식으로 파라미터 효율성을 확보한다.

또한, “상태‑기반 추론 전략”을 도입해 과도한 사고를 억제한다. 선형 어텐션의 상태 업데이트는 실제로 SGD 형태의 경사 하강이며, 각 단계에서 발생한 그래디언트 ∇_t를 모멘텀 방식으로 누적해 전역 그래디언트 G를 만든다. G는 전체 추론 방향을 나타내며, 현재 토큰 생성 시 q·G를 추가적인 가이드로 사용한다. 이렇게 하면 잡음이 섞인 개별 단계가 전체 흐름을 크게 벗어나지 못한다.

실험에서는 7개의 공개 데이터셋(수학, 논리, 코드 등)과 3가지 모델 규모(7B, 13B, 34B)를 대상으로 기존 CoT, 압축 CoT, 그리고 최신 효율적 어텐션(FlashAttention 등)과 비교했다. MAM 기반 모델은 평균 2.3×3.1× 빠른 추론 속도와 30% 이하의 메모리 사용량을 기록했으며, 정확도는 1.22.5%p 상승했다. Ablation 연구에서는 LA만 사용하거나 모멘텀 없이 사용할 경우 효율은 유지되지만 성능이 크게 떨어짐을 확인했다.

이 논문의 주요 기여는 (1) 추론 과정을 상태‑전이로 재구성해 선형 시간·상수 메모리 복합성을 달성, (2) Mixed Attention을 통해 기존 소프트맥스와 선형 어텐션을 자연스럽게 결합, (3) 전역 그래디언트를 활용한 노이즈 억제 메커니즘을 제시, (4) 다양한 모델·데이터에서 효율·성능 동시 향상을 실증한 점이다. 한계점으로는 상태 행렬의 차원 선택과 커널 함수 ϕ에 대한 민감도가 있으며, 매우 복잡한 다단계 추론에서는 상태 압축 손실이 누적될 가능성이 있다. 향후 연구에서는 동적 차원 축소, 하이퍼파라미터 자동 튜닝, 그리고 멀티모달 추론에의 확장을 탐색할 여지가 있다.