DLite, 에이전트 AI를 위한 히스토리 없는 추측 디코딩

에이전트 작업에서 드러나는 추측 디코딩의 한계

추측 디코딩은 작은 드래프트 모델이 여러 토큰을 먼저 제안하고, 큰 타깃 모델이 이를 한 번에 검증하는 방식이다. 제안된 토큰 중 상당수가 승인되면 타깃 모델이 토큰을 하나씩 생성해야 하는 횟수가 줄어들어, 출력의 검증 가능성을 유지하면서 추론 지연을 낮출 수 있다.

하지만 에이전트 시스템의 작업 흐름은 일반적인 대화보다 훨씬 동적이다. 여러 차례의 사용자 대화에 도구 호출 결과, 외부 환경의 피드백, 수정된 계획과 새로운 상태가 계속 추가된다. 컨텍스트는 길어질 뿐 아니라 매 턴마다 바뀐다. 논문은 이런 환경에서 기존 드래프트 모델의 수용 길이가 최대 80%까지 감소할 수 있다고 설명한다. 드래프트 모델이 과거의 KV 캐시를 계속 유지하는 방식 역시 메모리와 시스템 관리 부담을 키운다.

타깃 모델의 은닉 상태를 재사용하는 DLite

DLite는 두 문제의 공통 원인을 계속 증가하는 과거 컨텍스트로 본다. 드래프트 모델이 전체 상호작용 기록을 다시 저장하고 해석하도록 하는 대신, 타깃 모델이 추론 과정에서 이미 계산한 컨텍스트 은닉 상태를 직접 활용한다. 이를 바탕으로 로컬한 추측 예측을 수행하므로, 드래프트 측에서 과거 컨텍스트에 대응하는 별도의 KV 캐시를 명시적으로 유지하지 않는다.

여기서 ‘history-free’는 과거 정보가 필요 없다는 뜻이 아니다. 과거 정보는 타깃 모델이 이미 인코딩한 표현을 통해 간접적으로 사용된다. 따라서 예측 문제는 계속 커지는 대화 전체를 드래프트 모델이 매번 모델링하는 방식에서, 현재의 컨텍스트 표현을 이용해 가까운 범위의 다음 토큰을 예측하는 방식으로 바뀐다.

DLite는 여기에 가벼운 sequential head를 추가한다. 추측 블록의 여러 위치를 병렬로 예측하면 속도는 얻을 수 있지만, 블록 안의 뒤쪽 토큰은 앞쪽 토큰의 예측에 의존해야 한다. 순차 헤드는 이런 블록 내부의 인과 관계를 낮은 추가 비용으로 보완한다. 결과적으로 병렬 블록 예측의 이점을 유지하면서 위치 간 의존성을 더 자연스럽게 반영하려는 설계다.

학습 방식과 실험 결과

DLite는 타깃 모델의 은닉 상태, 로컬 예측 모듈, 순차 헤드를 함께 사용하는 이질적인 구조다. 저자들은 일반적인 독립 드래프트 모델의 학습 방식을 그대로 적용하지 않고, 이 구조에 맞춘 2단계 학습 전략을 제안한다. 제공된 자료에는 세부 학습 절차가 모두 공개되어 있지는 않지만, 히스토리를 명시적으로 유지하지 않는 예측과 블록 내부의 인과 구조를 함께 맞추는 것이 핵심 목표로 제시된다.

일반 목적 작업, 장문 컨텍스트 작업, 에이전트 작업을 대상으로 한 실험에서 DLite는 2~6의 수용 길이를 달성했다. 또한 최신 추측 디코딩 방법인 DSpark와 비교해 최대 70%의 성능 향상을 보였다고 보고한다. 이는 과거 상태의 중복을 없애는 일이 단순한 메모리 절감에 그치지 않고, 컨텍스트가 계속 변하는 상황에서 드래프트 예측의 안정성에도 영향을 줄 수 있음을 시사한다.

추론 시스템에 주는 의미

DLite의 핵심 기여는 추측 디코딩에서 컨텍스트를 어디에 보관해야 하는지를 다시 생각하게 만든다는 점이다. 긴 에이전트 세션에서 드래프트 모델이 전체 상호작용의 두 번째 복사본을 유지할 필요가 없을 수 있다. 타깃 모델의 중간 표현을 재사용하면 긴 세션에서 캐시 압력을 줄이고, 대화 기록의 길이에 덜 민감한 드래프트 비용을 설계할 가능성이 생긴다.

다만 제공된 요약만으로는 다양한 하드웨어, 모델 규모, 배치 조건에서의 실제 엔드투엔드 효과를 판단하기 어렵다. 수용 길이는 전체 처리량이나 지연과 동일하지 않으며, 은닉 상태 이동, 타깃 검증, 배칭, 동시성, 메모리 관리가 실제 배포 성능에 함께 영향을 준다. 그럼에도 DLite는 에이전트 추론을 위한 분명한 방향을 제시한다. 계속 커지는 과거 상태를 중복 저장하기보다, 타깃 모델이 이미 계산한 정보를 바탕으로 로컬 추측을 수행하는 것이다.

출처: Hugging Face Daily Papers