Nyströmformer가 Nyström 행렬 근사 방법을 자체-주의 메커니즘에 적용해 시간 복잡도를 O(n²)에서 O(n)으로 감소

Context

Transformer의 자체-주의 메커니즘은 입력 시퀀스의 모든 토큰 쌍 간 상호작용을 포착하지만, O(n²)의 시간 및 메모리 복잡도로 인해 긴 입력 시퀀스에 대한 학습이 비용이 많이 든다. 표준 자체-주의에서 Softmax 행렬 S = softmax(QK^T/√d)의 전체 계산을 피할 수 없다는 문제가 있다.

Technical Solution

• Nyström 점(landmarks)을 쿼리와 키에서 샘플링: 전체 Q, K 행렬 대신 m개의 쿼리 랜드마크 Q̃와 키 랜드마크 K̃를 선택해 계산량 감소
• Softmax 행렬을 세 개의 부분행렬 곱으로 근사: S = F̃Ã⁺B̃ = softmax(QK̃^T/√d) · softmax(Q̃K̃^T/√d)⁺ · softmax(Q̃K^T/√d)로 표현
• QK^T 전체 곱 연산 회피: 근사된 Softmax 행렬을 값(V) 행렬과 곱하여 O(n²) 복잡도 제거
• 세 개의 작은 행렬 곱셈으로 자체-주의 근사: F̃(n×m) · Ã⁺(m×m) · B̃(m×n) 순차 계산으로 선형 시간 달성

Key Takeaway

Nyström 방법은 전체 행렬을 계산하지 않고 선택된 행과 열만 샘플링하여 근사하는 수학적 기법인데, 자체-주의의 Softmax 행렬 특성을 고려해 쿼리/키 랜드마크 샘플링으로 적응시켜 Transformer의 근본적인 복잡도 문제를 해결한다. 이는 표준 알고리즘의 수학적 제약을 도메인별로 맞춤 변형하는 최적화 전략의 모범 사례이다.