기술 요약: 확산 언어 모델의 적응형 캐싱을 위한 특이 프록시

문제 제기

확산 언어 모델 (DLM) 은 임의 순서 디코딩과 병렬 생성을 가능하게 함으로써 자기회귀 (AR) 모델에 대한 유연한 대안을 제공합니다. 그러나 비인과적 성질과 양방향 어텐션에 대한 의존성으로 인해 AR 효율성의 핵심인 표준 키 - 밸류 (KV) 캐싱과 호환되지 않습니다. DLM 에서는 진화하는 전역 의존성을 포착하기 위해 재계산 없이 은닉 상태를 단순히 재사용할 수 없어, 추론 지연 시간이 prohibitive 하게 증가하며 ($O(T \times N^2)$ 복잡도) 발생합니다.

이 문제를 캐싱을 통해 완화하려는 기존 시도들은 두 가지 주요 한계에 직면해 있습니다:

1. 비싼 식별 (Costly Identification): 현재 방법들은 어떤 토큰이 업데이트를 필요로 하는지 식별하기 위해 고차원 유사도 계산 (예: 밸류 상태 드리프트 모니터링) 에 의존합니다. 시퀀스 길이가 확장됨에 따라 이러한 식별 오버헤드는 희소 업데이트로 인한 이득을 상쇄하는 경우가 많습니다.
2. 경직된 예산 할당 (Rigid Budget Allocation): 기존 접근법은 일반적으로 모든 모델 레이어에 걸쳐 균일한 업데이트 예산을 적용합니다. 이는 일부 레이어는 안정적 (업데이트가 거의 필요 없음) 이고 다른 레이어는 변동성이 큼 (빈번한 업데이트 필요) 은 은닉 상태의 이질적인 역학을 고려하지 못해, 계산적 중복이나 성능 저하를 초래합니다.

방법론: SPA-Cache

저자들은 업데이트 식별과 예산 할당을 공동으로 최적화하는 프레임워크인 SPA-Cache(적응형 캐싱을 위한 특이 프록시) 를 제안합니다. 단일 트랜스포머 레이어의 워크플로우는 세 가지 단계로 나뉩니다:

1. 특이 프록시를 통한 업데이트 식별

고차원 공간 ($d$) 에서 유사도를 계산하는 대신, SPA-Cache 는 저차원 부분 공간 ($r \ll d$) 에서 업데이트가 중요한 토큰을 식별하기 위해 특이 프록시 (singular proxy) 를 도입합니다.

• 이론적 근거: 저자들은 밸류 (Value) 상태의 안정성이 어텐션 출력 및 후속 FFN 출력의 안정성을 위한 충분 조건임을 증명합니다. 따라서 밸류 상태를 모니터링하면 출력 안정성에 대한 보수적인 보장을 제공합니다.
• 차원 축소: 고차원 투영의 계산 병목 현상을 피하기 위해, 이 방법은 밸류 투영 행렬 $W$ 에 특이값 분해 (SVD) 를 적용합니다. 상위 $r$ 개의 특이 벡터와 특이값만 유지하여 잘라낸 투영 $W_r$ 을 생성합니다.
• 메커니즘: 입력 상태를 이 저차원 매니폴드로 투영합니다. 현재 특이 프록시와 캐시된 프록시 간에 코사인 유사도를 계산합니다. 유사도가 가장 낮은 (드리프트가 가장 큰) 토큰들이 재계산을 위해 선택됩니다.
• 효율성: 이로써 식별 복잡도가 $O(d)$ 에서 $O(r)$ 로, 투영 오버헤드가 $O(d^3)$ 에서 $O(rd^2)$ 로 감소합니다.

2. 적응형 예산 할당

은닉 상태 드리프트가 모델 레이어 간에 크게 변한다는 점을 인식하여, SPA-Cache 는 정적이며 균일한 업데이트 비율을 동적이며 레이어별 적응형 전략으로 대체합니다.

• 관찰: 드리프트는 초기 레이어 (임베딩 변환) 에서 낮고, 중간 레이어 (복잡한 의미 변환) 에서 정점을 찍으며, 후기 레이어에서 안정화됩니다.
• 구현: 레이어 $l$ 의 업데이트 비율 $\rho(l)$ 은 구간별 가우시안 함수를 사용하여 매개변수화됩니다. 이는 변동성이 큰 중간 레이어에는 더 높은 업데이트 예산을 할당하고, 안정된 초기 및 후기 레이어는 적극적으로 캐싱하여 품질을 희생하지 않으면서 처리량을 극대화합니다.

3. 희소 실행

선택된 토큰 (인덱스 $I$) 은 희소 쿼리, 키, 밸류를 생성하는 데 사용됩니다. 어텐션 메커니즘은 부분적으로 업데이트된 KV 캐시에 대해 이러한 희소 쿼리를 사용하여 출력을 계산합니다. 피드포워드 네트워크 (FFN) 는 희소 어텐션 출력만 처리하며, 최종 레이어 출력은 선택된 인덱스에서만 업데이트되고 다른 위치는 캐시된 기능을 재사용합니다.

주요 기여

1. 밸류 프록시에 대한 이론적 정당성: 이 논문은 밸류 상태 유사도가 어텐션 및 FFN 출력의 발산을 제한함을 보여주는 공식적 증명 (정리 3.1 및 3.2) 을 제공하여, 업데이트 선택을 위한 원칙적인 식별자로서 밸류 상태의 사용을 검증합니다.
2. 저오버헤드 식별을 위한 특이 프록시: 잘라낸 SVD 를 활용하여 저차원 부분 공간에서 중요한 토큰을 식별하는 방법을 도입함으로써, 식별 과정 자체의 계산 비용을 크게 줄였습니다.
3. 적응형 예산 할당: 레이어별 이질성과 계산 자원을 정렬하는 동적 업데이트 전략을 도입하여, 균일한 할당의 비효율성을 피했습니다.
4. 병렬 디코딩과의 직교성: SPA-Cache 는 시퀀스 수준 최적화와 직교하도록 설계되어 병렬 디코딩 프레임워크와 통합하여 처리량 이득을 증폭시킬 수 있습니다.

실험 결과

저자들은 LLaDA-8B 와 Dream-v0-Instruct-7B 두 가지 DLM 에 대해 GSM8K, MMLU-PRO, MBPP, HumanEval 을 포함한 일곱 가지 벤치마크에서 SPA-Cache 를 평가했습니다.

• 처리량 개선: SPA-Cache 는 기본 디코딩 대비 최대 8 배의 처리량 개선을 달성했으며, 기존 캐싱 베이스라인 (dLLM-Cache, Fast-dLLM) 대비 2~4 배의 속도 향상을 보였습니다.
• 구체적 벤치마크:
  • LLaDA-8B 를 사용한 MBPP 에서 SPA-Cache 는 8.0 배의 속도 향상을 달성했습니다.
  • MMLU-PRO 에서 6.0 배(LLaDA-8B) 및 6.2 배(Dream-7B) 의 속도 향상을 달성했습니다.
• 병렬 디코딩과의 통합: 병렬 디코딩 방법과 결합했을 때, SPA-Cache 는 BBH 벤치마크에서 기본 디코딩 대비 총 28 배의 속도 향상을 보여주었으며, 동등한 정확도를 유지했습니다.
• 애블레이션 연구:
  • 전체 차원 밸류 식별자를 특이 프록시 ($r=128$) 로 대체한 결과, 정확도 손실 없이 초당 토큰 수 (TPS) 가 약 165 에서 약 179 로 향상되었습니다.
  • 적응형 예산 할당은 평균 업데이트 비율을 25% 에서 16% 로 줄여 TPS 를 약 189 로 추가로 증가시켰으며, 반면 균일한 16% 비율은 정확도 저하를 초래했습니다.

중요성과 주장

이 논문은 SPA-Cache 가 확산 언어 모델에 대한 새로운 효율성 기준을 수립하여, 고처리량 실용 응용을 위한 계산적으로 실현 가능한 대안으로 위치시킨다고 주장합니다. 원칙적인 이론적 분석과 적응형 자원 할당을 통해 추론 병목 현상을 근본적으로 해결함으로써, 이 방법은 생성 충실도를 훼손하지 않고 DLM 이 효과적으로 확장되도록 합니다. 저자들은 그들의 접근 방식이 학습이 필요 없으며 기존 병렬 디코딩 전략과 호환된다고 강조하며, 이는 이산 확산 모델의 실용적 배포를 위한 길을 연다고 결론지었습니다.