미스트랄 7B·8x7B 구조 해설: 슬라이딩 윈도우 어텐션과 MoE 작동 원리 정리

핵심 메시지

• 미스트랄은 라마와 거의 같은 디코더 전용 구조지만, 셀프 어텐션에 슬라이딩 윈도우를 쓰고 KV 캐시를 롤링 버퍼로 운영한다는 점이 다르다.
• 슬라이딩 윈도우 어텐션은 각 토큰이 창 안의 이웃만 보게 해 내적 연산을 줄이는데, 층을 여러 번 쌓으면 합성곱 신경망의 수용 영역처럼 먼 토큰 정보까지 간접적으로 전달된다.
• 창 밖 토큰은 어차피 마스킹되므로 KV 캐시도 창 크기만큼만 유지할 수 있고, 덮어쓰기로 순서가 어긋난 캐시는 포인터를 기준으로 되돌려 원래 순서를 복원한다.
• 긴 프롬프트는 창 크기와 같은 청크로 나눠 캐시를 미리 채우되, 질의는 현재 청크로 두고 키·값에는 이전 청크가 담긴 캐시를 함께 넣어야 청크 사이 연결이 끊기지 않는다.
• 8x7B의 희소 전문가 혼합은 토큰마다 게이트가 여덟 개의 피드포워드 중 점수가 높은 두 개만 골라 실행해, 파라미터 수는 늘리면서 추론 비용은 억제한다.

쉽게 이해하기

발표자는 미스트랄 7B와 8x7B를 벤 뒤 바닐라 트랜스포머와 무엇이 다른지부터 정리한다. 두 모델 모두 인코더 블록 형태의 층을 32번 반복하고 마지막에 정규화·선형 계층·소프트맥스를 붙인 디코더 전용 구조이며, 사실상 라마와 매우 비슷하다. 차이는 셀프 어텐션에 슬라이딩 윈도우가 들어가고, 추론용 KV 캐시가 롤링 버퍼 방식이며, 피드포워드의 활성화 함수가 바뀌었다는 점이다. 8x7B에서는 피드포워드 하나 대신 여덟 개가 병렬로 놓여 전문가 역할을 한다.

슬라이딩 윈도우 어텐션의 출발점은 인과 마스크다. 자기회귀 모델이므로 각 토큰은 뒤에 오는 토큰을 볼 수 없는데, 여기에 더해 창 크기보다 멀리 떨어진 앞쪽 토큰까지 가린다. 그만큼 계산해야 할 내적이 줄어 학습과 추론이 빨라지지만, 멀리 떨어진 단어끼리 직접 연결되지 않아 성능이 떨어질 수 있다는 우려가 남는다. 발표자는 책을 읽을 때 5장의 문장을 1장의 문장과 굳이 연결하지 않아도 되는 것처럼, 국소적 맥락만으로 충분한 경우가 많다는 직관도 함께 제시한다.

그럼에도 먼 토큰이 완전히 단절되지는 않는다. 한 층을 통과하면 각 위치의 벡터는 창 안의 토큰 정보를 흡수하고, 그 결과가 다음 층의 입력이 되면서 정보가 겹겹이 합쳐진다. 발표자는 소프트맥스 결과와 값 행렬의 곱을 손으로 따라가며, 층이 하나씩 쌓일 때마다 각 위치가 참조하는 토큰 집합이 넓어진다는 것을 보여준다. 이는 합성곱 신경망에서 커널을 반복 적용할수록 한 특징이 원본 이미지의 더 넓은 영역을 반영하게 되는 수용 영역 개념과 같은 원리다.

추론 최적화 부분에서는 KV 캐시를 다시 짚는다. 토큰을 하나씩 생성할 때 이전 단계에서 이미 계산한 내적을 반복하지 않도록, 새로 나온 토큰을 키와 값에만 덧붙이고 질의는 마지막 토큰으로 교체한다. 여기에 슬라이딩 윈도우를 결합하면 창 밖 값은 어차피 0이 곱해지므로 캐시를 창 크기로 제한할 수 있는데, 이것이 롤링 버퍼 캐시다. 가장 오래된 자리를 덮어쓰다 보면 저장 순서가 문장 순서와 어긋나므로, 마지막으로 쓴 위치를 가리키는 포인터를 기준으로 뒤쪽과 앞쪽을 이어 붙여 원래 순서를 복원한다.

마지막 축은 희소 전문가 혼합과 배치 처리다. 8x7B에서는 게이트가 토큰마다 여덟 개 전문가의 점수를 내고 상위 두 개만 골라 실행한 뒤, 그 두 점수에만 소프트맥스를 적용해 가중합을 만든다. 모델 용량은 키우면서도 토큰당 행렬 곱은 두 번으로 묶어두는 방식이다. 또 길이가 제각각인 여러 사용자 프롬프트를 패딩으로 채워 낭비하는 대신, 하나의 긴 시퀀스로 이어 붙이고 xFormers의 블록 대각 마스크로 프롬프트끼리 섞이지 않게 막는 구현도 소개한다.

주요 인사이트

• 성능 최적화의 상당 부분이 새로운 수학이 아니라 '계산하지 않아도 되는 것을 계산하지 않기'에서 나온다. 마스크로 어차피 지워질 내적, 캐시에 남겨둘 필요 없는 창 밖 토큰, 결과를 쓰지 않을 패딩 토큰이 모두 제거 대상이다.
• 국소 어텐션이 곧 문맥 상실은 아니다. 층을 쌓는 것만으로 참조 범위가 넓어지므로, 창 크기는 '모델이 볼 수 있는 최대 거리'가 아니라 '한 층에서 직접 볼 거리'로 이해해야 한다.
• 긴 프롬프트를 청크로 나눠 캐시를 채울 때 질의와 키·값의 범위를 다르게 잡는 설계가 핵심이다. 현재 청크만으로 마스크를 만들면 청크 경계에서 문맥이 끊겨 버린다.
• 희소 전문가 혼합에서 상위 두 개를 고른 뒤에 소프트맥스를 적용하는 순서에도 이유가 있다. 전체에 소프트맥스를 먼저 걸면 선택된 가중치의 합이 1이 되지 않아, 전문가 수를 바꿔가며 모델을 비교할 때 출력 범위가 달라진다.

자주 묻는 질문

관련 AI 소식