[논문 리뷰] Laughing Hyena Distillery: Extracting Compact Recurrences From Convolutions

Stefano Massaroli, Michael Poli외 12명|VU Research Portal|2023. 10. 28.

Advanced Neural Network Applications인용 수 4

한 줄 요약

이 논문은 사전에 훈련된 장거리 컨볼루션 시퀀스 모델(LCSMs)에서 압축되고 반복적인 상태공간 모델(SSMs)을 추출하는 디스틸레이션 방법인 Laughing Hyena를 소개한다. 이는 자동귀납적 생성 시 토큰당 O(1) 계산량과 메모리 사용을 가능하게 하며, 1.3B 파라미터에서 트랜스포머보다 10배 높은 처리량을 달성하고, Hyena보다 1.5배 높으며, 디스틸레이션 후 품질에 손실가지 않는다.

ABSTRACT

Recent advances in attention-free sequence models rely on convolutions as alternatives to the attention operator at the core of Transformers. In particular, long convolution sequence models have achieved state-of-the-art performance in many domains, but incur a significant cost during auto-regressive inference workloads -- naively requiring a full pass (or caching of activations) over the input sequence for each generated token -- similarly to attention-based models. In this paper, we seek to enable $\mathcal O(1)$ compute and memory cost per token in any pre-trained long convolution architecture to reduce memory footprint and increase throughput during generation. Concretely, our methods consist in extracting low-dimensional linear state-space models from each convolution layer, building upon rational interpolation and model-order reduction techniques. We further introduce architectural improvements to convolution-based layers such as Hyena: by weight-tying the filters across channels into heads, we achieve higher pre-training quality and reduce the number of filters to be distilled. The resulting model achieves 10x higher throughput than Transformers and 1.5x higher than Hyena at 1.3B parameters, without any loss in quality after distillation.

연구 동기 및 목표

• 표준 추론에서 O(K) 메모리와 O(K²) 계산량의 병목 현상을 해결하여 장거리 컨볼루션 시퀀스 모델(LCSMs)에서 상수 시간, 상수 메모리 자동귀납적 생성을 가능하게 하기 위해.
• 모델 품질을 유지하면서 사전에 훈련된 컨볼루션 레이어에서 저차원이고 안정적인 상태공간 모델(SSMs)을 추출하는 디스틸레이션 프레임워크를 개발하기 위해.
• 채널 간 가중치 묶음을 통해 Hyena 아키텍처를 재설계하여 사전 훈련 품질과 디스틸레이션 효율을 향상시키기 위해.
• 메모리 사용량을 줄이고 LCSMs에서 반복적 추론을 가능하게 하여 고처리량, 대용량 배치 추론을 가능하게 하기 위해.

제안 방법

• 사전에 훈련된 LCSMs의 컨볼루션 필터에서 압축된 SSMs를 추출하기 위해 유리 보간법과 모델 차원 축소 기법을 적용한다.
• 바리센트릭 및 프로니 유사 방법을 영감으로 삼아 안정성 향상과 수치적 문제 방지를 위한 인자 분해 모드 매개변수화를 도입한다.
• 근사화 목표로 컨볼루션 필터의 이질성 지표를 사용하여 다양한 후행 작업 간 호환성을 보장한다.
• 최적의 상태 차원 d를 결정하기 위해 헨켈 연산자 스펙트럼 분석을 활용한다.
• 효율적 필터 차원 증가와 디스틸레이션 복잡도 감소를 위해 채널 간 가중치 묶음을 통해 Hyena 블록을 재설계한다.
• 균형 절단과 모드 절단을 SSM 모델 축소에 활용하며, 그 오차 한계는 그람형 고유값에서 유도된다.

실험 결과

연구 질문

• RQ1 사전에 훈련된 장거리 컨볼루션 모델에서 압축되고 반복적인 SSMs를 추출하여 토큰당 O(1) 추론 비용을 달성할 수 있는가?
• RQ2 SSM 디스틸레이션에서 근사 오차와 모델 효율성 간 최적의 트레이드오프를 이룰 수 있는 상태 차원 d는 무엇인가?
• RQ3 SSM 매개변수화 방식을 어떻게 개선하여 수치적 불안정성과 디스틸레이션 중 수렴 문제를 방지할 수 있는가?
• RQ4 Hyena 블록의 아키텍처 수정이 사전 훈련 품질 향상과 디스틸레이션 비용 감소에 기여할 수 있는가?
• RQ5 LCSMs를 SSMs로 디스틸레이션하는 과정에서 후행 작업 성능이 유지되며 고처리량 추론이 가능해지는가?

주요 결과

• Laughing Hyena는 1.3B 파라미터에서 트랜스포머보다 10배 높은 최대 처리량을 달성했고, Hyena보다 1.5배 더 높으며, 디스틸레이션 후 품질에 손실가지 않는다.
• 1.3B 파라미터에서 Laughing Hyena는 동일한 메모리 제약 조건 하에서 트랜스포머보다 3배 적은 메모리로 512개 토큰을 생성할 수 있다.
• K토큰 생성에 대해 기존 kv-캐시 트랜스포머의 O(K) 메모리와 O(K²) 시간 대비, Laughing Hyena는 상수 O(d) 메모리와 O(dK) 시간을 제공한다.
• Hyena 블록의 채널 간 가중치 묶음은 사전 훈련 퍼플렉서티를 향상시키고, 디스틸레이션 대상 필터 수를 감소시킨다.
• 균형 절단과 모드 절단 방법은 일부 레이어에서 오차 감소가 비단조화적임을 보여주며, 수치적 안정성에 민감함을 시사한다.
• 근사화 목표로 필터 이질성 지표를 사용함으로써 다양한 후행 작업 간 강건성이 확보된다.