StreamKL: 어텐션 증류 (Attention Distillation) 성능 향상을 위한 빠르고 메모리 효율적인 KL 발산 (KL

어텐션 증류 (Attention distillation)는 하나의 어텐션 분포가 다른 분포와 일치하도록 Kullback-Leibler (KL) 발산 (divergence)을 최소화하며 학습하는 기법으로, 지식 증류 (knowledge distillation), 모델 압축 (model compression), 지속 학습 (continual learning), 그리고 희소 어텐션 (sparse-attention) LLM 학습에 널리 사용됩니다. 그러나 기존 방식들은 KL 감소 (KL reduction)를 계산하기 전에 두 어텐션 분포를 모두 실체화 (materialize)하며, 이로 인해 긴 컨텍스트 길이(long context lengths)에서는 감당하기 어려운 $O(N_QN_K)$의 메모리 및 IO 비용이 발생합니다. 본 논문에서는 이러한 이차적 실체화 (quadratic materialization)를 제거하는 어텐션 KL 발산 (attention KL divergence)을 위한 최초의 융합된 GPU 프리미티브 (fused GPU primitive)인 StreamKL을 제시합니다. StreamKL은 결합된 두 분포의 KL 감소 (KL reduction)를 위한 새로운 온라인 공식 (online formulation)을 도출하여, 쿼리-키 타일 (query-key tiles)을 온칩 SRAM (on-chip SRAM)을 통해 스트리밍하는 단일 패스 순전파 커널 (single one-pass forward kernel)을 가능하게 합니다. 역전파 (backward pass)의 경우, StreamKL은 어텐션 확률 (attention probabilities)을 타일 단위로 재계산하여 이차적 중간값 (quadratic intermediates)의 저장을 방지합니다. 나아가 우리는 전용 최적화가 적용된 효율적인 GPU 커널을 설계하고 구현했습니다. 실험 결과, StreamKL은 베이스라인 방법 대비 순전파 (forward pass)에서 최대 $43 imes$, 역전파 (backward pass)에서 최대 $14 imes$의 속도 향상을 제공함을 보여줍니다. 가장 중요한 점은, StreamKL이 어텐션 증류 (attention distillation)의 추가적인 HBM 점유 공간 (HBM footprint)을 $O(N_QN_K)$에서 $O(1)$로 줄여, 단일 GPU에서 긴 컨텍스트 증류 (long-context distillation)를 가능하게 한다는 것입니다.