SAS, 예측에 직접 기여하는 컨텍스트 순위로 희소 어텐션 개선

긴 컨텍스트를 처리하는 Transformer에서는 시퀀스가 길어질수록 어텐션의 누적 계산 비용이 이차적으로 증가합니다. 이를 줄이기 위해 각 쿼리가 참조하는 토큰이나 블록의 일부만 남기는 희소 어텐션이 연구되고 있습니다. 하지만 계산량을 줄이는 것만으로는 충분하지 않습니다. 고정된 어텐션 예산 안에서 어떤 컨텍스트를 보존할지 정확하게 순위를 매겨야 합니다.

기존 방식의 목표 불일치

일반적인 학습형 희소 어텐션은 가벼운 선택기가 컨텍스트 단위별 점수를 계산한 뒤 하드 Top-K를 적용합니다. 이 방식은 예산을 명확하게 통제할 수 있지만, 이산 선택 단계가 언어 모델링 손실에서 선택기로 이어지는 기울기를 차단할 수 있습니다. 그 결과 선택기는 최종 예측에 미치는 영향보다, 원래 밀집 모델의 계층별 어텐션 분포를 모방하는 방향으로 학습되는 경우가 많습니다.

밀집 어텐션 가중치를 따라가는 것은 합리적인 대안이지만, 제한된 예산에서 실제로 가장 유용한 컨텍스트를 직접 최적화하는 것과는 다릅니다. 밀집 모델에서 높은 가중치를 받은 단위가 다른 단위들이 제거된 이후에도 반드시 가장 큰 예측 이득을 주는 것은 아니기 때문입니다. 보존 가능한 컨텍스트가 적을수록 이러한 순위 오차는 더 큰 성능 손실로 이어질 수 있습니다.

SAS의 핵심 아이디어

Simple Attention Sparsification, 즉 SAS는 복잡한 선택 파이프라인을 추가하기보다 학습 경로 자체를 바꿉니다. 학습 중 선택기는 연속적인 점수를 만들고, 이 점수를 게이트로 사용해 어텐션 logits에 주입합니다. 특히 게이트를 로그 형태로 어텐션 Softmax 내부에 배치합니다. 그러면 언어 모델링 손실이 일반적인 역전파를 통해 선택기를 직접 업데이트할 수 있습니다.

논문이 강조하는 실용적 설계는 다음과 같습니다.

• Softmax 내부의 로그 게이트: 선택을 별도의 비미분 가지치기 단계로 두지 않고 어텐션 정규화와 통합합니다.
• 정규화된 Softmax 게이트: 과거 컨텍스트는 선택 대상이지만 현재 블록은 항상 유지됩니다. 정규화를 통해 두 정보원의 스케일 차이를 조정합니다.
• 연속 선택기 점수 유지: 단순한 유지·삭제 라벨 대신 후보 사이의 상대적 우선순위를 학습합니다. 예산이 달라져도 순위 정보를 활용하기 쉽습니다.

장기 시퀀스 학습을 지원하기 위해 연구진은 SAS를 FlashAttention 스타일의 계산에 통합하는 메모리 효율적인 Triton 커널도 구현했습니다. 이는 방법론뿐 아니라 실제 학습 시스템에서 발생하는 실행 비용까지 고려한 설계입니다.

의미와 한계

논문 설명에 따르면 SAS는 추론, 장기 컨텍스트 이해, 에이전트 작업에서 학습형 희소 어텐션 기준선보다 일관되게 나은 결과를 보였으며, 어텐션 예산이 매우 제한된 조건에서 특히 큰 이점을 보였습니다. 핵심 시사점은 희소 어텐션의 품질이 단순히 얼마나 많은 계산을 줄였는지뿐 아니라, 최종 예측에 도움이 되는 컨텍스트를 얼마나 잘 순위화하는지에 달려 있다는 것입니다.

이 관점은 장기 컨텍스트 추론, 컨텍스트 압축, 효율적인 Transformer 서빙 연구에 영향을 줄 수 있습니다. 밀집 모델이 어디에 어텐션을 두었는지를 그대로 모방하기보다, 희소화 이후에도 예측에 유용한 정보를 직접 보존하도록 선택기를 학습할 수 있기 때문입니다. 다만 제공된 소재에는 모델 규모, 데이터셋, 압축 비율, 절대 성능 수치가 포함되어 있지 않으므로, 개선 폭과 일반화 범위는 전체 논문의 실험 조건을 확인한 뒤 판단해야 합니다.

출처: Hugging Face Daily Papers