Attention Mechanism

1. Overview

어텐션 메커니즘(Attention Mechanism)은 입력 시퀀스의 각 토큰이 다른 모든 토큰과의 관계를 계산하여 문맥에 의존하는 표현을 만드는 연산으로, Transformer의 핵심 구성 요소입니다. 각 토큰은 "어디를 얼마나 참고할지"를 가중치로 계산하고, 그 가중치로 다른 토큰의 정보를 섞어 자신의 새 표현을 만듭니다.

어텐션은 원래 RNN 기반 번역 모델에서 디코더가 인코더 출력의 관련 부분에 집중하도록 하려고 제안되었고(Bahdanau et al., 2014), Transformer에서 순환 구조를 대체하는 주 연산이 되었습니다. 이 글은 Scaled Dot-Product Attention 수식의 의미, 여러 헤드를 쓰는 이유, Self-Attention과 Cross-Attention의 차이, 그리고 시퀀스 길이에 대해 제곱으로 늘어나는 계산 복잡도와 그 완화 방법을 정리합니다.

2. Scope & Boundaries

In-Scope

• Scaled Dot-Product Attention: Q/K/V의 역할과 $\sqrt{d_k}$ 스케일링.
• Multi-Head Attention: 여러 헤드의 병렬 계산과 연결.
• Self-Attention vs. Cross-Attention: Q와 K/V의 출처에 따른 구분.
• 계산 복잡도: $O(n^2 \cdot d)$ 의 의미와 완화 기법.

Out-of-Scope

• Transformer 전체 구조(정규화, FFN, 위치 인코딩) → Transformer Architecture 영역.
• 토큰 임베딩의 학습 → LLM Embedding Spaces 영역.
• GPU 커널 수준의 구현 최적화: 메모리 계층을 고려한 구현 세부는 개요만 다룹니다.

Boundaries

• Attention vs. Transformer: 어텐션은 하나의 연산이고, Transformer는 어텐션을 FFN, 정규화, 잔차 연결과 함께 층으로 쌓은 아키텍처입니다.
• 어텐션 가중치 vs. 설명 가능성: 어텐션 가중치는 모델이 계산에 사용한 값이지만, 그것이 곧 예측의 이유라는 해석은 별도의 검증이 필요한 주장입니다.

3. Counterexample

• "어텐션 가중치 = 모델이 이 단어 때문에 그렇게 답했다"는 해석: 가중치가 높은 위치가 곧 예측의 원인이라고 단정하는 것. 여러 층과 헤드, 잔차 연결을 거치므로 한 층의 가중치만으로 인과를 설명할 수 없습니다.
• 긴 컨텍스트가 공짜라는 가정: 시퀀스 길이 $n$ 이 두 배가 되면 어텐션 행렬 계산은 네 배가 됩니다. 컨텍스트 윈도우를 늘리는 일은 메모리와 시간 비용이 함께 늘어나는 결정입니다.
• FlashAttention을 근사 기법으로 오해: FlashAttention은 어텐션 결과를 근사하지 않고 정확히 계산하면서, GPU 메모리 계층의 읽기·쓰기 횟수를 줄여 속도와 메모리를 개선하는 구현입니다. 연산량(FLOPs)은 여전히 $O(n^2)$이지만, $n \times n$ 어텐션 행렬 전체를 메모리에 두지 않아 메모리 사용량이 시퀀스 길이에 선형으로 줄어듭니다. 계산 결과 자체를 단순화하는 근사 기법(희소 어텐션 등)과는 구분해야 합니다.

4. Prerequisites

• 행렬 곱과 softmax (Basic): $QK^\top$ 의 형상과 softmax의 행 단위 정규화를 이해해야 합니다.
• 임베딩 벡터 (Basic): Q/K/V는 임베딩 벡터를 선형 변환한 것입니다. → LLM Embedding Spaces
• seq2seq 구조 (Recommended): Cross-Attention이 인코더-디코더 모델에서 하는 역할을 이해하는 데 필요합니다.

5. Learning Map

연결 노트

• Next-Token Prediction — Attention으로 계산된 표현이 다음 토큰 예측에 사용
• LLM Embedding Spaces — Q/K/V는 임베딩 벡터를 선형 변환한 것
• LLM Scaling Laws — 헤드 수·d_model 크기가 스케일링의 핵심 변수

6. Learning Topics

Basic

Core Topic 01: Scaled Dot-Product Attention

• Why to Learn: 모든 Transformer 층의 핵심 계산이므로 이 수식을 정확히 읽을 수 있어야 이후의 변형(마스크, 멀티헤드, KV 캐시)을 이해할 수 있습니다.
• What to Learn:
  • Concepts: Q(Query)는 현재 토큰이 무엇을 찾는가, K(Key)는 각 토큰이 어떤 정보를 보유하는가, V(Value)는 실제 전달할 정보입니다.
  • Skills: $QK^\top$ 로 모든 쌍의 유사도 점수를 만들고, softmax로 행 단위 가중치를 만든 뒤 $V$ 의 가중 평균을 계산.
  • Tools: NumPy, PyTorch.
  • Trade-offs: √d_k 스케일링은 내적 값이 커질수록 softmax가 포화되어 기울기가 작아지는 것을 방지합니다.

$\mathrm{Attention}(Q,K,V) = \mathrm{softmax}\!\left(\dfrac{QK^{\top}}{\sqrt{d_k}}\right)V$

• How to Learn:
  • 1단계: 토큰 3~4개의 작은 예제에서 $QK^\top$ 행렬과 softmax 결과(각 행의 합이 1)를 직접 계산합니다.
  • 2단계: $d_k$ 를 키우면서 스케일링 유무에 따른 softmax 출력의 최댓값 변화를 관찰합니다.
• Implement: 배치와 마스크를 받는 scaled_dot_product_attention(Q, K, V, mask) 함수.

Recommended

Core Topic 02: Multi-Head Attention

• Why to Learn: 한 번의 가중 평균으로는 한 가지 관계만 반영되므로, 여러 부분 공간에서 동시에 어텐션을 계산하는 이유와 구성을 알아야 합니다.
• What to Learn:
  • Concepts: 동일 입력에 대해 여러 Attention을 병렬 계산한 뒤 연결합니다. 각 헤드가 서로 다른 의미 관계(통사, 의미, 거리 등)를 포착할 수 있습니다.
  • Skills: 헤드별 투영 행렬, Concat, 출력 투영 $W^O$.
  • Tools: PyTorch nn.MultiheadAttention.
  • Trade-offs: 헤드마다 어떤 역할을 맡는지는 학습 결과에 따라 달라지며, 모든 헤드가 해석 가능한 역할을 갖는 것은 아닙니다.

$\mathrm{MultiHead}(Q,K,V) = \mathrm{Concat}(\mathrm{head}_1,\dots,\mathrm{head}_h)\,W_O$

• How to Learn:
  • 1단계: 사전 학습 모델의 한 층에서 헤드별 어텐션 맵을 그려 서로 다른 패턴이 있는지 관찰합니다.
  • 2단계: $d_k = d_{model}/h$ 로 헤드를 나눴을 때 전체 계산량이 단일 헤드와 비슷한 이유를 설명합니다.
• Implement: 단일 헤드 함수를 이용해 $h$ 개 헤드를 병렬로 계산하고 $W_O$ 로 합치는 모듈.

Practical

Core Topic 03: Self-Attention과 Cross-Attention

• Why to Learn: 같은 어텐션 연산이 Q와 K/V를 어디서 가져오느냐에 따라 전혀 다른 역할을 하므로, 모델 구조를 읽을 때 두 유형을 구분할 수 있어야 합니다.
• What to Learn:
  • Concepts: Self-Attention은 한 시퀀스 안에서, Cross-Attention은 서로 다른 두 시퀀스 사이에서 정보를 가져옵니다.
  • Skills: 아래 표에서 Q와 K/V의 출처를 구분하고, 디코더의 인과 마스크를 이해.
  • Tools: 어텐션 맵 시각화 도구.
  • Trade-offs: Cross-Attention은 조건 정보를 명시적으로 주입할 수 있지만 두 시퀀스 길이의 곱에 비례하는 비용이 듭니다.

• How to Learn:
  • 1단계: 번역 예시에서 디코더의 한 위치가 인코더의 어느 토큰에 어텐션을 주는지 맵으로 확인합니다.
  • 2단계: 디코더 전용 모델에서 Cross-Attention이 사라지고 Self-Attention만 남는 구조를 비교합니다.
• Implement: 같은 어텐션 함수에 입력만 바꿔 Self와 Cross를 모두 수행하는 예제.

Advanced

Core Topic 04: 계산 복잡도와 완화 기법 (Complexity)

• Why to Learn: 컨텍스트 길이 확장의 병목이 어텐션의 이차 비용에서 오므로, 어떤 완화 방법이 무엇을 바꾸는지 구분해야 합니다.
• What to Learn:
  • Concepts: 어텐션의 계산 복잡도는 $O(n^2 \cdot d)$ 로 시퀀스 길이 $n$ 에 제곱 비례하며, 긴 컨텍스트의 병목 원인입니다.
  • Skills: 정확한 어텐션을 유지하며 메모리 접근을 줄이는 방법(FlashAttention)과, 계산 자체를 줄이는 근사·희소 방법(sparse attention 등)의 구분.
  • Tools: PyTorch 내장 SDPA, FlashAttention 구현체.
  • Trade-offs: 정확한 구현은 결과가 같지만 하드웨어에 의존하고, 근사 방법은 더 긴 길이를 다룰 수 있으나 품질 영향을 평가해야 합니다.
• How to Learn:
  • 1단계: 시퀀스 길이를 1K, 4K, 16K로 늘렸을 때 어텐션 행렬의 원소 수를 계산합니다.
  • 2단계: FlashAttention 논문에서 연산량이 아닌 메모리 읽기·쓰기가 병목이라는 분석을 읽습니다.
• Implement: 길이별 어텐션 연산 시간과 메모리를 측정해 그래프로 그리는 벤치마크.

7. Terminology

8. References

Primary References

• [Vaswani et al., 2017. "Attention Is All You Need." arXiv:1706.03762] — Scaled Dot-Product Attention, Multi-Head Attention, 계산 복잡도 비교.
• [Bahdanau, Cho & Bengio, 2014. "Neural Machine Translation by Jointly Learning to Align and Translate." arXiv:1409.0473] — seq2seq 번역에서의 어텐션 원형.

Secondary References

• [Dao et al., 2022. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." arXiv:2205.14135] — 정확한 어텐션을 IO 관점에서 가속하는 방법.
• [Child et al., 2019. "Generating Long Sequences with Sparse Transformers." arXiv:1904.10509] — 희소 어텐션으로 복잡도를 줄이는 접근.

Industry References

• [PyTorch Documentation: scaled_dot_product_attention] — 어텐션 연산의 표준 구현과 사용법.

9. Final Checklist

Primary Checklist

• $\text{softmax}(QK^\top/\sqrt{d_k})V$ 에서 각 항의 형상과 의미를 설명할 수 있는가?
• 스케일링이 없을 때 softmax가 포화되는 이유를 설명할 수 있는가?

Secondary Checklist

• Self-Attention과 Cross-Attention에서 Q와 K/V의 출처를 구분할 수 있는가?
• Multi-Head가 단일 헤드 대비 무엇을 추가하는지, 그리고 헤드의 역할을 단정하면 안 되는 이유를 말할 수 있는가?

Industry Checklist

• 컨텍스트 길이를 늘릴 때 어텐션 비용이 $n^2$ 으로 증가함을 감안해 메모리와 지연 시간을 측정했는가?
• 정확한 어텐션 구현(FlashAttention)과 근사 어텐션을 구분해 품질 영향을 평가했는가?