수달이네 기술 블로그

12. Attention 연산 구현 본문

AI공부/자연어처리

12. Attention 연산 구현

슬픈 수달이 2026. 3. 5. 22:13

구현

import numpy as np

# 전체 출력 형식을 소수점 이하 네 자리로 설정
np.set_printoptions(precision=4, suppress=True)

# 단어와 해당 임베딩 벡터를 딕셔너리로 정의합니다.
embedding_dict = {
    '<sos>': np.random.rand(512),
    '<eos>': np.random.rand(512),
    '커피': np.random.rand(512),
    '한잔': np.random.rand(512),
    '어때': np.random.rand(512),
    '오늘': np.random.rand(512),
    '날씨': np.random.rand(512),
    '좋네': np.random.rand(512),
    '옷이': np.random.rand(512),
    '어울려요': np.random.rand(512),
    'PAD': np.zeros(512)  # 패딩 벡터는 0으로 채웁니다.
}

# 입력 문장
sentences = [
    ['<sos>', '커피', '한잔', '어때', '<eos>'],
    ['<sos>', '오늘', '날씨', '좋네', '<eos>'],
    ['<sos>', '옷이', '어울려요', '<eos>', 'PAD']
]

# 토큰을 임베딩 벡터로 변환
embeddings = np.array([[embedding_dict[token] for token in sentence] for sentence in sentences])
print("임베딩 행렬의 형태:", embeddings.shape)

# 임베딩 행렬의 형태: (3, 5, 512)
  • 단어 리스트를 생성해 임베딩 벡터 딕셔너리로 정의한다. (512차원)
    • 3개의 문장 x 문장의 길이 5 x 512차원으로 임베딩 행렬이 표현된다.
    • 아직까지는 난수로 채워넣음
# 쿼리, 키, 밸류 행렬 초기화
num_heads = 8
head_dim = 512 // num_heads  # 각 헤드의 차원
heads = np.split(embeddings, num_heads, axis=2)  # 512차원 임베딩 벡터를 8개의 헤드로 분할하여 heads에 저장
queries = heads.copy()
keys = [head.transpose(0, 2, 1) for head in heads]  # 키 행렬을 각 헤드의 전치를 통해 초기화 (첫 번째 축: 배치 크기, 두 번째 축: 문장 길이, 세 번째 축: 헤드 차원)
values = heads.copy()

print("쿼리 행렬의 형태:", queries[0].shape)  
print("키 행렬의 형태:", keys[0].shape)  
print("밸류 행렬의 형태:", values[0].shape)

# 쿼리 행렬의 형태: (3, 5, 64)
# 키 행렬의 형태: (3, 64, 5)
# 밸류 행렬의 형태: (3, 5, 64)
  • 헤드를 8개로 나누어 줄 줄 것이다.
    • 이로 만든 쿼리, 키, 밸류의 행렬은 위와 같다.(키는 전치행렬)
# 특정 토큰 (커피, 한잔, 어때)의 인덱스
tokens_of_interest = ['커피', '한잔', '어때']
indices_of_interest = [sentences[0].index(token) for token in tokens_of_interest]

# 어텐션 이전의 임베딩 테이블 중 특정 토큰들의 평균 값 계산
print("어텐션 이전의 임베딩 테이블 중 '커피', '한잔', '어때' 토큰의 평균 값:")
initial_avg = np.mean(embeddings[0, indices_of_interest, :], axis=1)
print(initial_avg)

# 어텐션 이전의 임베딩 테이블 중 '커피', '한잔', '어때' 토큰의 평균 값:
# [0.4919 0.4783 0.5085]
  • 변화 이전의 커피, 한잔, 어때의 임베딩 벡터의 평균을 구해본다.

어텐션 연산

# 스케일링 및 어텐션 스코어 계산
attention_scores = np.matmul(queries[0], keys[0])
scaling_factor = np.sqrt(head_dim)
scaled_attention_scores = attention_scores / scaling_factor

attention_scores: 쿼리와 키의 내적을 계산(유사도 행렬)

  • 위의 값을 스케일링 scaling_factor(헤드 차원의 제곱근)을 나눠줌.
# 패딩 처리
mask = np.array([[token == 'PAD' for token in sentence] for sentence in sentences])
mask = mask[:, np.newaxis, :]  # 차원을 맞추기 위해 확장
scaled_attention_scores = np.where(mask, -np.inf, scaled_attention_scores)
  • 문장의 길이를 패딩으로 맞춰준다. -np.inf즉 -∞을 패딩으로 사용해 소프트맥스하면 0으로 변환
# 소프트맥스 적용 함수
def softmax(x):
    exp_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
  • 소프트맥스 적용함수 정의
# 복원된 헤드를 저장할 리스트
restored_heads = []

for i in range(num_heads):
    query = queries[i]
    key = keys[i]
    value = values[i]
    
    # 내적 계산 후 스케일링
    attention_scores = np.matmul(query, key) / scaling_factor
    
    # 패딩 처리
    mask = np.array([[token == 'PAD' for token in sentence] for sentence in sentences])
    mask = mask[:, np.newaxis, :]  # 차원을 맞추기 위해 확장
    attention_scores = np.where(mask, -np.inf, attention_scores)
    
    # 소프트맥스 적용
    attention_weights = softmax(attention_scores)
    
    # 밸류와의 곱셈
    restored_head = np.matmul(attention_weights, value)
    restored_heads.append(restored_head)
  • 각 헤드별로 내적 계산, 수케일링
  • 이후 패딩처리
  • 소프트맥스 적용
  • 밸류 행렬과 곱해줌.
# 모든 헤드를 결합하여 원래 차원으로 복원
final_output = np.concatenate(restored_heads, axis=2)
  • 헤드 결합: 다른 관점에서 계산한 결과를 합쳐 풍부한 표현 학습
# 어텐션 이후의 결과 중 특정 토큰들의 평균 값 계산
print("어텐션 이후의 결과 중 '커피', '한잔', '어때' 토큰의 평균 값:")
final_avg = np.mean(final_output[0, indices_of_interest, :], axis=1)
print(final_avg)

# 어텐션 이후의 결과 중 '커피', '한잔', '어때' 토큰의 평균 값:
# [0.4932 0.4905 0.4952]
  • 결과 출력, 즉, 위의 값이 문맥을 반영한 어텐션 결과이다.