전문가 혼합(MoE) 모델 직접 만들기: 라우터와 top-k 게이팅, 희소 연산까지 코드로 이해하기

핵심 메시지

• GPT-4와 딥시크 같은 모델은 일반적인 트랜스포머 대신 전문가 혼합(MoE) 구조를 쓴다.
• MoE는 어텐션 뒤의 밀집 신경망을 여러 개의 작은 전문가 신경망으로 바꾸고, 라우터가 토큰을 일부 전문가에게만 보낸다.
• 전문가를 모두 메모리에 올려야 해서 GPU 용량은 더 필요하지만, 실제 계산은 일부만 하므로 처리량이 빨라진다.
• 라우터가 몇몇 전문가에만 쏠리지 않도록 학습 중 가우시안 잡음을 섞는 noisy top-k 게이팅을 쓴다.
• 학습 데이터로는 심슨 대본 약 700만 자를 쓰고, 문자 단위 137개 어휘로 토큰화해 다음 문자를 예측하게 한다.

쉽게 이해하기

강의는 파이토치 노트북 하나로 전문가 혼합 모델을 처음부터 만들어 보는 과정을 따라간다. 원본 코드는 안드레이 카르파티의 간단한 트랜스포머 구현에서 출발해 다른 사람들이 MoE를 얹은 것을 다시 손본 것으로, 주석과 설명을 크게 늘렸다. GPU 없이 구글 코랩의 CPU만으로도, 심지어 맥북에서도 돌아간다.

먼저 데이터와 토큰화다. 심슨 대본 약 700만 자를 불러와 서로 다른 문자 137개를 확인하고, 각 문자에 1부터 137까지 번호를 붙이는 가장 단순한 방식으로 숫자로 바꾼다. 그다음 블록 크기 8, 배치 크기 4로 잘라 넣고 다음 토큰을 맞히게 한다. 강연자는 이 방식이 자기지도 학습이라 별도의 라벨 작업이 필요 없고, 그래서 모델들이 텍스트만 들이부어도 그토록 많이 배울 수 있었다고 짚는다.

다음은 어텐션이다. 2017년 무렵 등장한 이 개념이 현대 AI 모델의 토대라고 설명하면서, 어텐션은 문장 조각 속 단어들 사이의 관계를 모두 담은 더 나은 벡터를 만들어 내는 일이라고 직관적으로 정리한다. 코드에서는 쿼리와 키를 곱해 전치하고 하삼각 행렬을 만드는데, 이는 미래 정보가 뒤로 새어 들어오지 않게 막는 인과 마스킹 때문이다. 소프트맥스는 어텐션 가중치를 정규화하기 위해 쓴다. 헤드를 여러 개 두는 멀티헤드는 각 헤드가 저마다의 관점으로 보게 하려는 것이며, 이론이 아니라 실험으로 확인된 개선이라고 강조한다.

여기서부터 갈라진다. 보통의 트랜스포머라면 어텐션 뒤에 밀집 신경망이 오지만, MoE는 작은 신경망 여러 개를 두고 계산할 때마다 그중 일부만 골라 쓴다. 일부만 쓰기 때문에 희소(sparse) 방식이라 부른다. 전문가를 전부 올려 두어야 하니 GPU 메모리는 더 들지만 실제 연산량은 줄어들어, 대형 사업자들이 이 구조를 선호하는 이유가 된다.

어느 전문가에게 보낼지는 라우터가 정한다. 어텐션에서 나온 입력을 보고 어떤 전문가들이 맡을지 고르는 간단한 선형 계층이며, top-k 로짓과 인덱스를 뽑아 상위 몇 개만 선택한다. 소프트맥스를 얹어 가중치 합이 1이 되게 하면 낮은 전문가는 자연스럽게 떨어져 나간다. 다만 학습 중에 모델이 특정 전문가에만 지나치게 의존하는 일이 생길 수 있어, 가우시안 잡음을 로짓에 더해 더 많은 전문가를 쓰도록 밀어 주는 noisy top-k 게이팅을 기본값으로 삼았다.

마지막으로 모든 조각을 합쳐 학습한다. 심슨 데이터에서는 최대 1만 회 반복이 좋은 결과를 냈다. 실험 기록을 보면 처음에는 손실이 매우 높고 출력도 무작위 글자에 가깝지만, 4,000스텝쯤 되면 영어 문장에 훨씬 가까운 형태가 나오기 시작한다. 강연자는 이 정도 데이터와 모델 크기로 영어의 뉘앙스까지 기대할 수는 없지만 개선 과정만큼은 분명히 보인다고 말한다.

주요 인사이트

• MoE의 장점은 흔히 '똑똑해진다'로 오해되지만, 강의가 짚는 실제 이점은 자원 배분이다. 파라미터는 많이 두되 한 번의 계산에는 일부만 참여시켜, 메모리를 더 쓰는 대신 처리 속도를 얻는 맞바꿈이다.
• 라우터 쏠림 문제와 그 해법인 잡음 주입은 MoE가 이론적으로 깔끔한 구조가 아니라 경험적으로 다듬어진 구조라는 점을 보여 준다. 강연자가 반복해서 '실험해 보라'고 말하는 이유가 여기에 있다.
• 문자 단위 토큰화와 심슨 대본이라는 작은 데이터로도 구조 자체는 그대로 재현된다. 대형 모델을 이해하는 데 대형 자원이 필요한 것은 아니라는 점이 이 노트북의 교육적 가치다.
• 강연자는 MoE를 몰라도 어텐션만큼은 한 줄씩 이해하라고 권한다. AI 면접에서 반드시 물어보는 주제이자 나머지 구조가 모두 그 위에 쌓이기 때문이다.

자주 묻는 질문

관련 AI 소식