재미있는 삽질 기록

Multi-Head Attention 직접 구현하기

#pytorch#transformer#attention#multi-head

멀티헤드 어텐션을 직접 구현하면서 헷갈렸던 지점들을 정리한다.

헤드의 차원은 embedding_dim을 헤드 수로 나눈 것#

먼저 헤드를 빼고 보면, 입력은 (B, L, D) = (batch_size, sequence_length, embedding_dim)이고, 는 각각 Linear(embedding_dim, embedding_dim)이다. 입력을 이 레이어에 통과시키면 Q, K, V가 모두 (B, L, D)로 나온다. 멀티헤드에서는 이 D를 헤드 수 H로 나눠 각 헤드가 d_k = D / H 차원을 맡는다.

같은 입력인데 왜 서로 다른 레이어를 통과시키나#

forward의 인자로 들어오는 Q, K, V는 self-attention에서는 모두 동일한 src 입력이다. 동일한 src 세 개를 Q, K, V라는 서로 다른 임무의 텐서로 만들어야 한다. 그러려면 서로 다른 가중치 행렬을 가진 레이어를 각각 통과시켜야 한다. 레이어를 지난다는 건 학습 가능한 파라미터를 거친다는 뜻이고, 그래서 같은 입력에서도 Q와 K와 V가 서로 다른 표현으로 갈라진다. 헤드마다 학습하는 관점이 다른 것도 같은 이유다.

(B, L, d)에서 (B, H, L, d_k)로 갈 때 왜 바로 view를 못 쓰나#

헤드로 쪼갤 때 (B, L, d)를 곧장 (B, H, L, d_k)로 view하면 안 된다. 순서를 나눠서 봐야 한다.

  • 먼저 d를 헤드 수와 헤드 차원으로 쪼갠다: (B, L, d) → (B, L, H, d_k)
  • 그다음 순서를 바꾼다: (B, L, H, d_k) → (B, H, L, d_k)

이유는 헤드와 헤드 차원(H, d_k)이 d라는 메모리 공간을 쪼갠 것이기 때문이다. d를 (H, d_k)로 나누는 건 같은 자리에서의 분할이라 view로 맞지만, (B, L, d)에서 H를 곧장 L 옆으로 빼버리면 H가 L의 데이터로부터 값을 끌어오게 되어 엉뚱하게 섞인다. 그래서 view로 (B, L, H, d_k)를 만든 뒤, transpose로 L과 H의 자리를 바꿔 (B, H, L, d_k)로 가야 한다. view가 메모리를 연속된 순서로 다시 끊어 읽을 뿐이라는 점은 2. View의 작동 원리 이해하기에 정리해 두었다.

scaled dot-product attention#

def scaled_dot_product_attention(Q, K, V, scale, mask=None):
    # 1) 스코어 계산 (normalize까지)
    attention_score = (Q @ K.permute(0, 1, 3, 2)) / scale

    if mask is not None:
        attention_score = attention_score.masked_fill(mask, -1e10)

    # 2) 소프트맥스 확률분포
    attention_dist = torch.softmax(attention_score, dim=-1)

    # 3) 컨텍스트 계산
    attention = attention_dist @ V

    # 4) 결과 반환
    return attention, attention_dist

Q와 K를 곱할 때 K의 마지막 두 축을 바꿔(permute) 내적이 되게 하고, softmax는 dim=-1로 각 query 행마다 확률이 되도록 잡는 점에 유의한다. masking된 자리는 softmax 전에 아주 작은 값(-1e10)으로 채워 확률이 0에 가깝게 죽인다.

scale은 head_dim이어야 한다#

스케일 값에서 실수하기 쉽다. d_k는 멀티헤드에서 헤드의 차원이므로, 전체 embedding_dim이 아니라 self.head_dim이 들어가야 한다.

self.scale = torch.sqrt(torch.tensor(self.head_dim))
self.register_buffer('scale', self.scale)

register_buffer를 쓰는 이유#

self.scale처럼 모델의 파라미터로 취급되어야 하지만 학습 대상은 아닌 상수는 register_buffer로 등록하는 게 가장 깔끔하다. register_buffer로 등록하면

  • 모델의 state_dict에 포함되어 저장하고 불러올 때 함께 관리된다.
  • .to(device) 호출 시 다른 파라미터와 함께 해당 장치로 자동 이동한다.
  • optimizer가 업데이트하지 않는, 즉 학습되지 않는 파라미터임을 명확히 한다.

Add & Norm#

마지막으로 각 서브레이어를 감싸는 Add & Norm은 한 줄로 LayerNorm(x + dropout(sublayer(x)))이다. 입력 x를 서브레이어 출력에 잔차로 더한 뒤 LayerNorm을 건다.

정리#

  • 헤드 차원 d_k = embedding_dim / 헤드 수. 같은 src를 서로 다른 가중치 레이어로 통과시켜 Q와 K와 V로 가른다.
  • 헤드로 쪼갤 때는 (B, L, d)를 view로 (B, L, H, d_k)로 만든 뒤 transpose로 (B, H, L, d_k)로 간다. 곧장 view하면 H가 L의 데이터를 끌어와 섞인다.
  • scale은 √head_dim이고, 학습 대상이 아닌 상수라 register_buffer로 등록한다. Add & Norm은 LayerNorm(x + dropout(sublayer(x)))이다.

댓글