라마(LLaMA) 구조 해설: KV 캐시, 회전 위치 임베딩, 그룹 쿼리 어텐션 정리

핵심 메시지

• 라마는 인코더와 디코더를 함께 두지 않고 다음 토큰 예측만 하며, 정규화를 각 블록 뒤가 아니라 앞에 배치한다.
• RMS 정규화는 평균을 쓰지 않고 제곱평균제곱근만으로 스케일을 맞춘다. 층 정규화의 효과가 재중심화보다 재스케일링에서 온다는 가설에 기반해 통계량을 하나만 계산한다.
• 회전 위치 임베딩은 쿼리와 키에만, 그리고 가중치 행렬을 곱한 뒤에 적용되며 두 토큰의 거리가 멀어질수록 내적의 상한이 줄어드는 장기 감쇠 성질을 갖는다.
• KV 캐시는 추론에서만 쓰는 기법으로, 이미 계산한 키와 값을 저장해 매 단계 반복되던 내적을 없애고 마지막 행만 계산하게 한다.
• KV 캐시가 병목을 계산에서 메모리 접근으로 옮기자, 키·값의 헤드를 하나로 줄인 다중 쿼리 어텐션과 그 절충안인 그룹 쿼리 어텐션이 등장했다.

쉽게 이해하기

영상은 원조 트랜스포머 그림과 라마 구조를 나란히 놓고 차이를 짚는 데서 시작한다. 라마에는 인코더·디코더 구분이 없고 다음 토큰 예측을 위한 자기어텐션만 남으며, 정규화 층이 각 블록 앞으로 옮겨졌다. 위치 정보는 임베딩에 더하는 고정 벡터가 아니라 회전 위치 임베딩으로 바뀌어 쿼리와 키에만 적용되고, 어텐션에는 KV 캐시와 그룹 다중 쿼리 방식이 붙고, 피드포워드의 활성 함수는 ReLU에서 SwiGLU로 교체된다. 라마 1은 67억부터 650억까지 네 가지 크기로 나왔고 임베딩 차원은 4096이며, 라마 2에서는 문맥 길이와 학습 토큰 수가 대체로 두 배가 됐다.

정규화 설명은 왜 필요한지부터 풀어낸다. 발표자는 신경망을 하드웨어 팀과 소프트웨어 팀이 있는 휴대폰 공장에 비유한다. 손실이라는 최고경영자가 두 팀의 전략을 한꺼번에 크게 바꾸면, 소프트웨어 팀은 이전에 본 적 없는 하드웨어를 받아 결과가 크게 흔들린다. 층의 출력 분포가 이렇게 요동치는 현상이 내부 공변량 이동이고, 정규화는 각 항목을 평균과 분산으로 표준화한 뒤 학습 가능한 감마·베타로 다시 조절해 이 흔들림을 줄인다.

라마가 쓰는 RMS 정규화는 여기서 한 단계 덜어낸다. 원 논문은 층 정규화의 성공이 재중심화가 아니라 재스케일링 불변성에서 온다고 가정하고, 평균에 의존하지 않는 제곱평균제곱근 통계량으로 나눈 뒤 학습 가능한 감마만 곱한다. 평균과 분산을 모두 구하지 않으니 계산이 줄고, 실전에서도 잘 작동한다는 것이 채택 이유다.

위치 인코딩은 절대와 상대의 차이를 먼저 정리한다. 절대 방식은 토큰 하나에 고정 벡터를 더해 지도 위 위도·경도처럼 위치를 표시하고, 상대 방식은 어텐션 계산에 두 토큰의 거리를 나타내는 벡터를 끼워 넣는다. 회전 위치 임베딩은 쿼리와 키의 내적이 두 임베딩과 상대 거리만으로 결정되도록 하는 함수를 찾는 데서 출발하며, 복소수 공간의 식을 오일러 공식으로 바꾸면 회전 행렬이 나타난다. 고차원에서 이 행렬은 희소해 낭비가 크므로, 실제로는 두 번째 차원의 부호를 뒤집은 벡터를 활용해 더 효율적인 형태로 계산한다. 학습할 값이 없어 절대 방식처럼 한 번 계산해 재사용한다.

후반부는 추론 비용 이야기다. 다음 토큰 예측을 한 토큰씩 이어붙이며 반복하면 이미 계산한 내적을 매번 다시 구하게 되는데, 모델은 인과적 마스크를 쓰고 우리가 필요한 것은 마지막 행뿐이므로 키와 값을 캐시하면 낭비를 없앨 수 있다. 그런데 GPU는 연산보다 데이터 전송이 훨씬 느려서, 캐시 도입 뒤에는 메모리 접근이 새로운 병목이 된다. 그래서 키·값의 헤드를 하나만 두는 다중 쿼리 어텐션이 제안됐고, 영어-독일어 번역 BLEU 점수가 26.7에서 26.5로 소폭 떨어지는 대신 추론 시간은 크게 줄었다. 라마가 쓰는 그룹 쿼리 어텐션은 쿼리를 몇 개의 그룹으로 나눠 그룹마다 키·값 헤드를 하나씩 두는 절충안이다.

마지막 주제는 피드포워드의 활성 함수다. 라마는 ReLU 대신 SwiGLU를 쓰는데, 이 형태는 파라미터 행렬이 셋이라 원조 구조와 총 파라미터를 맞추려면 행렬 크기를 줄여야 한다. 논문의 실험에서 이 함수는 혼란도와 여러 벤치마크에서 좋은 성적을 냈지만, 저자들은 왜 잘 되는지에 대한 설명을 제시하지 않고 그 성공을 신의 은총으로 돌린다고 적었다. 발표자는 이것이 농담처럼 들리지만 딥러닝 연구의 실상에 가깝다며, 수백억 파라미터 모델에서는 소규모 모델로 단순화하거나 절제 실험으로 비교하는 방식이 현실적이라고 덧붙인다.

주요 인사이트

• 라마의 변경 사항은 대부분 성능 지표보다 비용 구조에 관한 선택이다. RMS 정규화는 통계량을 하나 줄이고, KV 캐시는 중복 연산을 없애며, 그룹 쿼리 어텐션은 메모리 전송량을 줄인다.
• KV 캐시가 성능 개선의 끝이 아니라는 점이 이 영상의 백미다. 계산 병목을 해결하자 메모리 접근이 병목으로 바뀌었고, 그 새 병목이 다중 쿼리·그룹 쿼리 어텐션을 낳았다.
• 회전 위치 임베딩을 값(V)에 적용하지 않는 이유는 위치가 개입해야 할 지점이 두 토큰의 관계 점수, 즉 쿼리와 키의 내적이기 때문이다. 적용 시점도 가중치 행렬을 곱한 뒤로 원조 트랜스포머와 다르다.
• 장기 감쇠는 부수 효과가 아니라 의도된 성질이다. 멀리 떨어진 두 단어의 관계 강도가 자연히 약해지도록 내적의 상한이 거리와 함께 줄어든다.
• SwiGLU 논문의 결론이 “왜 잘 되는지 설명하지 않는다”는 문장이라는 대목은, 현재 딥러닝에서 상당수 설계 선택이 이론적 해명보다 절제 실험과 탐색으로 굳어진다는 사실을 보여준다.

자주 묻는 질문

관련 AI 소식