Apple ML Research, 'Stochastic KV Routing' 공개… 깊이 축 캐시 공유로 KV 메모리 풋프린트 대폭 절감

Apple Machine Learning Research가 'Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing' 논문을 공개했다. 저자는 Anastasiia Filippova, David Grangier, Marco Cuturi, João Monteiro 4인이다.

논문은 트랜스포머 언어 모델을 높은 처리량으로 서빙하려면 자기회귀(autoregressive) 생성 과정의 중복 계산을 피하기 위해 Key-Value(KV)를 캐싱해야 한다는 점에서 출발한다. 그러나 KV 캐싱의 메모리 풋프린트는 상당히 크고 서빙 비용에 직접적인 영향을 준다는 것이 연구진의 문제 의식이다.

최근 연구들은 주로 시간 축(temporal axis)에서의 압축과 eviction을 통해 KV 캐시 축소를 다뤄왔다. 그러나 저자들은 깊이(depth) 차원이 시간 축에 직교적이면서도 견고한 최적화 경로를 제공한다고 주장한다.

선행 연구들은 모든 레이어에 대한 풀 캐시가 중복적이라는 점을 시사해왔다. 다만 cross-layer 캐시 공유의 실제 구현은 실용적 도전 과제로 남아 있었으며, 기존 방법들은 처리량(throughput)이 감소하거나 첫 토큰까지의 시간(time-to-first-token)이 증가하는 문제를 겪었다.

이번 논문은 한 레이어의 캐시를 드롭(dropping)하는 것이 정보 손실 없이 효율적인 최적화를 제공할 수 있음을 보였다. 이를 위해 'random cross-layer attention'이라는 단순한 훈련 접근법을 제안했다.

훈련 중 각 레이어는 자신의 KV state에 attend할지, 또는 선행하는(preceding) 어느 레이어의 state에 attend할지를 무작위로 선택한다. 이 확률적(stochastic) 프로세스를 통해 모델은 다양한 깊이 축 캐시 공유 전략에 견고하게 적응한다.

이러한 적응성은 배포 시점의 알 수 없는 하드웨어 제약에 대한 유연성을 보장한다는 것이 저자들의 설명이다. 배포 환경에 따라 어떤 깊이 축 공유 전략을 택하더라도 모델이 견고하게 작동할 수 있다는 의미다.

평가 결과 저자들은 이 기법을 사전훈련(pre-training) 또는 파인튜닝(fine-tuning) 단계에서 적용할 경우 다양한 모델 패밀리에 대해 깊이 축 캐시 공유가 가능해진다고 보고했다.

특히 데이터가 제한된(data-constrained) 환경의 더 큰 모델에서는 이 접근이 정규화(regularization)와 유사한 효과를 보였으며, 캐시의 메모리 풋프린트를 크게 줄이면서도 성능을 유지하거나 개선하는 사례가 빈번했다고 논문은 밝혔다.