Cascade Token Selection for Transformer Attention Acceleration

기술 요약: 트랜스포머 어텐션 가속화를 위한 캐스케이드 토큰 선택

문제 제기
트랜스포머 모델은 자기 어텐션 메커니즘에서 상당한 계산 병목 현상에 직면해 있으며, 여기서 $QK^\top$ 곱을 계산하고 후속 소프트맥스 가중 합을 수행하는 비용은 시퀀스 길이에 따라 이차적으로 증가합니다 ($O(T^2 d_h)$). 최근 FlashAttention 과 같은 방법들이 메모리 대역폭을 최적화했음에도 불구하고, 알고리즘적 복잡성은 긴 시퀀스에 대한 장벽으로 남아 있습니다.

동행 논문인 활성화 비상관 어텐션 (Activation Decorrelation Attention, ADA) 은 각 레이어에서 $r \ll T$ 개의 작은 "대표" 토큰 집합을 선택함으로써 이 복잡성을 $O(T^2)$ 에서 $O(r^2)$ 로 줄이는 솔루션을 제안했습니다. ADA 는 그람 임계값을 사용하여 이러한 토큰을 식별합니다. 즉, 모든 선행 토큰과의 최대 코사인 상관관계가 임계값 $1 - \tau^2$ 미만인 경우 해당 토큰을 대표 토큰으로 간주합니다. 그러나 선택 과정 자체는 모든 레이어에서 완전한 $T \times T$ 그람 행렬을 계산해야 하므로 $O(T^2 d)$ 의 비용을 발생시킵니다. 중간 길이의 시퀀스 (예: $T=2048$) 의 경우, 이 선택 오버헤드는 줄이려는 어텐션 비용과 비교할 수 있을 정도로 커서 잠재적인 속도 향상을 무효화합니다. 본 논문이 다루는 핵심 문제는 ADA 의 대표 토큰 선택 단계의 높은 계산 비용입니다.

방법론: 캐스케이드 메커니즘
본 논문은 대표 토큰 집합의 깊이 일관성 (depth coherence) 을 활용하는 "캐스케이드" 메커니즘을 소개합니다. 저자들은 레이어 $l$ 의 정보성 토큰 집합이 레이어 $l+1$ 의 집합을 강력하게 예측할 수 있음을 관찰했습니다. 이는 이론적으로 트랜스포머 블록의 리프시츠 연속성 (Lipschitz continuity) 에 근거합니다. 한 레이어에서 거의 중복 (중복성) 이거나 비상관 (대표성) 인 토큰들은 컨텍스트 정보가 추가됨에 따라 약간의 교란만 발생하면서 다음 레이어에서도 그 관계를 유지하는 경향이 있습니다.

각 레이어마다 완전한 $T \times T$ 그람 행렬을 재계산하는 대신, 캐스케이드 메커니즘은 다음과 같이 작동합니다:

1. 이전 레이어에서 대표 집합 $R^{(l)}$ 을 상속합니다.
2. 두 개의 더 작은 그람 곱을 계산하여 상속된 집합을 검증합니다:
   • 상속된 대표들 간의 $r \times r$ 곱을 계산하여 중복된 항목이 있는지 확인합니다.
   • 비대표 항목과 유효한 상속 대표 간의 $(T-r) \times r$ 교차 곱을 계산하여 비대표 항목 중 비상관 (즉, 새로운 대표) 이 된 항목이 있는지 확인합니다.
3. 이러한 검사를 기반으로 소수의 추가 및 제거를 통해 집합을 업데이트합니다.

이 접근 방식은 레이어당 선택 복잡성을 $O(T^2 d)$ 에서 $O(Trd)$로 줄여$T/r$ 배의 속도 향상을 제공합니다.

주요 기여

• 알고리즘적 혁신: 완전한 그람 행렬 계산을 점진적 검증 및 업데이트 단계로 대체하는 캐스케이드 토큰 선택 알고리즘 제안.
• 이론적 정당성: 트랜스포머 블록의 리프시츠 연속성과 잔차 연결의 가산성으로 인해 대표 집합의 깊이 일관성이 발생함을 증명하는 증명 (제안 3).
• 보수적 보장: 캐스케이드 방법이 보수적임을 보여주는 증명 (제안 4). 이는 독립적인 대표 집합의 초집합을 생성합니다. 진정한 대표를 절대 놓치지 않으므로 어텐션 품질이 저하되지 않지만, 대표 수 ($r$) 가 약간 증가할 수 있습니다.
• 직교 가속: 본 논문은 이 방법을 선형 레이어를 깊이 일관성 기반을 상속받아 가속화하는 게이트드 서브스페이스 추론 (Gated Subspace Inference, GSI) 의 토큰 차원 유사체로 위치시킵니다. 두 방법은 트랜스포머 순전파 전체 ($d$ 축 및 $T$ 축 모두) 를 커버합니다.

실험 결과
이 방법은 AMD MI300X 하드웨어를 사용하여 세 가지 모델 계열 (GPT-2 124M, GPT-J 6B, OPT 6.7B) 에서 시퀀스 길이 $T=512$ 및 그람 임계값 $\tau=0.30$ 으로 검증되었습니다.

• 깊이 일관성: 깊은 레이어에서 연속된 레이어의 대표 집합 간 자카드 지수는 높습니다 (0.83~0.94). 이는 토큰 집합의 안정성을 확인합니다.
• 그람 연산 절감: 캐스케이드 메커니즘은 총 그람 연산을 다음과 같이 감소시켰습니다:
  • GPT-2(12 레이어): 22%
  • GPT-J(28 레이어): 52%
  • OPT(32 레이어): 63%
    모델 깊이가 깊어질수록 절감폭이 커지는데, 이는 레이어 0 에서의 초기 완전 그람 비용이 더 많은 레이어에 걸쳐 분산되기 때문입니다.
• 교체율: 깊은 레이어에서 교체율 (수정된 집합의 비율) 은 3~18% 로 떨어지며, 이는 상속된 집합이 최소한의 조정이 필요함을 나타냅니다.
• 확장 예측: 저자들은 $r$ 이 부분적으로 선형적으로 증가하는 더 긴 시퀀스 (예: $T=8192$) 의 경우, 선택 속도 향상이 $18\times$ 에 달할 것으로 예상하며, 선택 비용을 지배적인 항에서 무시할 수 있는 수준으로 낮출 것이라고 전망합니다.

의의 및 주장
본 논문은 캐스케이드 메커니즘이 훈련된 트랜스포머의 구조적 속성을 드러낸다고 주장합니다. 즉, 정보성 토큰 집합은 네트워크 깊이를 통해 전파되는 일관된 구조적 속성입니다. 이 일관성을 활용함으로써 이 방법은 ADA 에 대한 선택 오버헤드를 실용적으로 만들어, 어텐션 압축이 선택 단계의 병목 현상 없이 이론적 속도 향상을 달성할 수 있게 합니다.

저자들은 이 방법이 메모리 대역폭을 최적화하는 FlashAttention 이나 고정 패턴을 사용하는 희소 어텐션과 같은 기존 방법들을 보완하는 가속화 기술임을 강조합니다. 토큰 가지치기나 병합 방법처럼 표현을 영구적으로 변경하는 것과 달리, 캐스케이드 ADA 는 구조 변경이나 재학습 없이 입력 구조에 기반하여 대표를 동적으로 선택합니다. GSI 와 캐스케이드 ADA 의 결합은 추론 비용을 곱셈적으로 감소시켜, 순전파의 복잡성을 명목상 모델 크기가 아닌 입력의 "실제 복잡성"으로 효과적으로 줄여줍니다.