어텐션을 보는 라우터, MoE의 선택이 달라지는 지점

오늘 다룰 논문은 arXiv에 공개된 Attention-Aware Routing: Coupling Routing and Attention in MoEs이며 식별자는 arXiv:2609.20974입니다. 이 논문은 MoE 라우터가 토큰의 hidden state만 보고 전문가를 고르는 관행을 깨고 어텐션 가중치의 시간적 흐름과 주파수 패턴을 라우팅 신호로 함께 쓰자고 제안합니다. 원문 논문 링크는 제목 아래에 걸려 있습니다.

라우터는 왜 토큰 하나만 보고 결정을 내릴까

안녕하세요, 패트릭입니다. 제가 이 논문을 집어 든 이유는 꽤 단순합니다. MoE를 오래 보다 보면 라우터가 하는 일이 이상할 정도로 가볍게 느껴질 때가 있기 때문입니다. 수십 개의 전문가가 대기하고 있는데, 정작 누가 일할지를 정하는 결정은 현재 토큰의 hidden state 한 줄로 끝나는 경우가 많습니다. 문맥이 어디로 흐르고 있는지, 앞 문장에서 어떤 단어가 계속 시선을 끌고 있는지 같은 정보는 어텐션 어딘가에 이미 계산되어 있는데도 말이죠.

생각해보면 사람도 비슷하게 일하지 않습니다. 회의에서 발언자를 고를 때 지금 손 든 사람만 보지 않고, 그동안 누가 어떤 주제에 오래 시선을 뒀는지를 함께 봅니다. 그런데 MoE 라우터는 그 두 번째 정보를 버리고 있었습니다. 논문이 진단한 출발점도 바로 여기입니다. 라우터는 토큰의 hidden state에 의존해 전문가를 고르고 가중치를 매기는데, 이 신호에는 문맥 정보가 제한적으로만 담긴다는 것입니다.

여기서 재미있는 질문이 생깁니다. 어텐션 가중치는 이미 문맥을 요약한 값인데, 왜 라우팅에서는 쓰지 않았을까요. 이유는 partly 관성에 가깝습니다. 라우터는 가볍고 빨라야 한다는 믿음이 있었고, 어텐션을 들여다보는 일은 괜히 무겁고 얽히는 일로 여겨졌기 때문입니다. 그래서 라우터와 어텐션은 같은 층에 살면서도 서로 말을 섞지 않는 이웃처럼 지냈습니다. 저는 이 단절이 논문이 파고든 지점이라고 봅니다. 아키텍처를 키우는 이야기가 아니라, 이미 있는 신호를 라우터가 읽게 만드는 이야기이기 때문입니다.

그렇다면 어텐션 가중치는 무엇을 따로 기억하고 있을까

이 논문이 어텐션을 바라보는 방식이 제가 가장 흥미롭게 본 부분입니다. 저자들은 어텐션 가중치를 토큰 hidden state와는 분리된 문맥 상태의 요약으로 취급합니다. hidden state가 지금 이 토큰이 무엇인지 말해준다면, 어텐션 가중치는 이 토큰이 지금까지 어디를 보고 왔는지를 말해준다는 것이죠. 같은 단어가 반복돼도 앞 문맥에 따라 시선의 궤적이 달라지는데, 그 궤적 자체가 라우팅에 쓸모 있는 정보라는 주장입니다.

일상적인 비유를 들어보겠습니다. 같은 질문을 받은 두 사람이 있다고 해보죠. 말은 똑같이 "모르겠습니다"라고 했지만, 한 사람은 계속 창밖을 보고 있었고 다른 한 사람은 회의 자료의 특정 페이지를 오래 응시했습니다. 대답이라는 결과는 같아도, 그다음에 누구에게 일을 맡겨야 할지는 시선의 기록을 보면 달라집니다. 어텐션 가중치가 바로 그 시선의 기록입니다. 토큰 표현 하나로 뭉개면 사라지는 차이인데, 궤적을 펼치면 다시 보이는 차이죠.

문제는 그다음입니다. 시선의 기록은 길고 시끄럽습니다. 매 스텝의 어텐션 행렬을 통째로 라우터에 넣으면 차원이 폭발하고 노이즈까지 딸려옵니다. 그래서 논문은 슬라이딩 윈도우라는 장치를 씁니다. 최근 구간의 어텐션 가중치만 잘라서 보고, 그 안에서 시간적인 흐름과 주파수 성분을 뽑아냅니다. 저는 이 선택이 영리하다고 봅니다. 날것 그대로를 던지는 대신, 라우터가 소화할 수 있는 형태로 요약했기 때문입니다. 문맥 전체를 외우겠다는 욕심을 버리고, 지금 결정에 필요한 리듬만 건네는 셈입니다.

시간과 주파수를 라우터에 얹으면 무엇이 달라질까

그렇다면 시간적 흐름과 주파수 패턴이라는 말은 실제로 무엇을 뜻할까요. 시간적 요소는 비교적 직관적입니다. 최근 몇 스텝 동안 특정 위치에 대한 어텐션이 꾸준히 높았는지, 아니면 잠깐 스치고 지나갔는지를 보는 것입니다. 꾸준히 머문 시선과 스치는 시선은 다음 행동을 다르게 이끄니까요. 수학 문제를 풀 때 중간 계산 결과가 몇 스텝 동안 계속 참조된다면, 그 문제는 같은 유형의 전문가에게 맡기는 편이 낫다는 신호가 될 수 있습니다.

주파수 성분은 조금 더 낯설게 들리지만, 아이디어 자체는 단순합니다. 어텐션의 변화가 매끈하게 이어지는지, 아니면 들쭉날쭉하게 뛰는지를 보는 것입니다. 어떤 토큰은 매 스텝 비슷한 위치를 참조하면서 안정적으로 흐르고, 어떤 토큰은 스텝마다 참조 위치가 크게 요동칩니다. 이런 요동의 결이 라우팅에 단서가 된다는 것이 논문의 시각입니다. 진동이 크다는 것은 문맥이 아직 정착되지 않았다는 뜻일 수 있고, 안정적이라는 것은 이미 방향이 잡혔다는 뜻일 수 있기 때문입니다.

저는 이 대목을 읽으면서 오디오 엔지니어를 떠올렸습니다. 파형을 시간 축으로만 보면 놓치는 것이 있는데, 주파수로 펼치면 저음이 새는지 고음이 깨지는지가 바로 보이니까요. 어텐션도 마찬가지입니다. 순서대로만 보면 그저 숫자의 나열이지만, 리듬과 진동으로 펼치면 문맥의 상태를 읽을 수 있습니다. 논문은 이 두 종류의 요약을 라우터의 입력에 덧붙입니다. 라우터 자체를 거대하게 만드는 것이 아니라, 라우터가 보는 풍경을 넓혀주는 방식이죠. 그래서 이름도 정직합니다. 어텐션을 의식하는 라우팅이라는 것입니다.

어텐션을 고치는 방법이 아니라는 점을 짚고 넘어가면

여기서 헷갈리기 쉬운 지점이 하나 있어서 먼저 짚고 가겠습니다. 이 방법은 어텐션 모듈을 고치는 방법이 아닙니다. 어텐션 가중치를 읽어서 라우팅에 쓸 뿐, 어텐션 계산 자체를 다시 학습하지 않습니다. 이 구분은 생각보다 중요합니다. 어텐션을 직접 건드리면 모델 전체의 표현이 흔들리지만, 읽기만 하면 라우팅 결정이라는 좁은 통로에만 영향을 주기 때문입니다.

그렇다면 읽기만으로 어텐션이 바뀔 수 있을까요. 직접은 아니지만 간접적으로는 바뀝니다. 라우터가 다른 전문가를 고르면 그 층의 출력이 달라지고, 그 출력이 잔차 스트림을 타고 다음 층으로 흘러가면 다음 층이 보는 입력이 달라지기 때문입니다. 입력이 달라지면 어텐션도 달라집니다. 논문이 포착한 것도 바로 이 간접 경로입니다. 한 층의 라우팅 변화가 다음 층의 어텐션 싱크를 키우는 현상이 관찰됐다는 것입니다.

이 대목은 오해하기 쉬워서 조금 조심해서 읽을 필요가 있습니다. 어텐션 싱크가 커졌다는 말이 곧 모델이 똑똑해졌다는 뜻은 아니기 때문입니다. 싱크는 특정 토큰에 어텐션이 몰리는 현상인데, 그 자체는 좋을 수도 있고 나쁠 수도 있습니다. 논문이 말하는 것은 방향입니다. 라우팅을 바꾸는 것만으로 어텐션의 모양까지 움직였다는 것, 다시 말해 라우터와 어텐션이 생각보다 단단히 연결되어 있다는 것입니다. 저는 이 연결을 확인한 대목이 숫자로 된 성능보다 오래 남는다고 봅니다. 성능은 다음 논문이 깨뜨릴 수 있지만, 연결의 존재는 설계자의 머릿속에 남기 때문입니다.

베이스를 얼리고 라우팅만 가르치면 무엇이 보일까

실험 설계로 넘어가 보죠. 논문은 기반 트랜스포머를 통째로 얼리고 라우팅 파라미터만 학습합니다. 저희가 새로운 AI 아키텍처를 볼 때도 benchmark 숫자만 보지는 않습니다. 실제로 무엇이 바뀌었고, 그 변화가 시스템에서 어떤 비용 구조를 만드는지를 함께 봅니다. 그런 눈으로 보면 이 얼리기 결정은 꽤 단호합니다. 좋아진 것이 있다면 그 공로는 라우팅에만 돌리겠다는 선언이니까요.

이런 통제는 왜 필요할까요. MoE 실험에서 가장 흔한 함정은 전체를 함께 미세조정해놓고 라우팅 덕분이라고 말하는 것입니다. 어텐션도 바뀌고 피드포워드도 바뀌었는데 라우팅이 좋아졌다고 주장하면, 무엇이 좋아진 것인지 알 수 없습니다. 기반을 얼리면 그 변명이 사라집니다. 라우터가 어텐션 신호를 읽는 법만 배웠는데 성능이 바뀌었다면, 적어도 그 변화는 라우팅의 몫이라고 말할 수 있죠.

물론 이런 설계에도 대가가 있습니다. 전체를 풀어 학습하면 숫자는 더 잘 나올 수도 있습니다. 기반까지 함께 움직이면 모델이 새 신호에 몸 전체를 맞추기 때문입니다. 그럼에도 저자들이 얼리기를 택한 것은 숫자를 키우는 대신 해석을 지키겠다는 선택입니다. 저는 이런 선택을 보면 논문의 자신감이 읽힙니다. 크게 이기는 그림보다 작게라도 정확히 이기는 그림을 택했기 때문입니다. 실제 시스템 관점에서 보면 다음 질문이 바로 생깁니다. 라우팅만으로 얻은 이득이 전체 미세조정의 이득과 겹치는지, 아니면 서로 더해지는지 말이죠. 논문은 그 질문에 답하지 않지만, 적어도 질문을 던질 수 있는 발판은 만들었습니다.

틀린 답만 짧아지고 맞은 답은 그대로인 이유

이제 행동의 변화로 들어가 보겠습니다. 논문이 보고한 장면 가운데 저는 이 대목이 가장 사람 같습니다. 이 방법을 적용하면 틀린 답은 짧아지는데 맞은 답의 길이는 그대로였다는 것입니다. 길게 헤매는 생성이 줄었다는 표현이 보고서에 등장하는데, 말 그대로 모델이 틀릴 때 빙빙 도는 일이 줄었다는 뜻입니다.

왜 이런 일이 생길까요. 생성 모델이 틀릴 때는 대개 초반에 방향을 잘못 잡고, 그 잘못된 방향을 스스로 정당화하면서 문장을 늘리는 패턴을 보입니다. 이미 어긋난 궤적을 길게 끌고 가는 것이죠. 반면 맞을 때는 필요한 계산을 끝내고 멈춥니다. 어텐션의 리듬을 읽는 라우터는 이 둘을 구분하는 데 유리합니다. 궤적이 요동치고 참조가 흩어지는 상태에서는 일찍 방향을 끊는 쪽의 전문가를 고를 수 있기 때문입니다.

여기서 한 가지 짚을 점이 있습니다. 짧아졌다는 것이 곧 정확해졌다는 뜻은 아닙니다. 틀린 답이 짧아진 것은 비용 관점에서는 분명히 좋지만, 정답률 자체와는 다른 축의 개선입니다. 논문도 이 둘을 섞지 않습니다. 정답률 개선은 GSM8K 같은 benchmark에서 따로 보고하고, 길이 변화는 생성 행동의 변화로 따로 보고합니다. 저는 이런 분리가 마음에 듭니다. 하나의 숫자로 모든 것을 설명하려는 유혹을 참았기 때문입니다. 운영 측면에서 보면 이 차이는 상당히 큽니다. 틀린 답이 짧아지면 낭비되는 토큰이 줄고, 서빙 큐에서 기다리는 다음 요청이 빨리 들어갑니다. 맞은 답의 길이가 그대로라는 것은 필요한 추론을 잘라낸 것이 아니라는 방증이기도 하고요.

한 층의 선택이 다음 층의 시선을 바꾸는 과정

그렇다면 층과 층 사이에서는 정확히 무슨 일이 일어날까요. 논문이 짚은 경로는 이렇습니다. l층에서 라우팅이 바뀌면 그 층의 출력이 달라지고, 이 출력이 잔차 스트림에 더해져 다음 층으로 전달됩니다. 다음 층의 어텐션은 바뀐 입력을 보고 계산되므로, 직접 건드리지 않았는데도 모양이 달라집니다. 보고서에 따르면 이 변화는 l+1층의 어텐션 싱크를 키우는 방향으로 나타났습니다.

이 장면을 조금 더 천천히 따라가 보겠습니다. 잔차 스트림은 각 층이 남기는 메모가 계속 누적되는 복도 같은 것입니다. l층에서 다른 전문가가 선택되면 복도에 붙는 메모의 내용이 달라집니다. 다음 층은 그 복도를 걸으면서 어텐션으로 어디를 볼지 정하는데, 메모가 달라졌으니 시선도 달라집니다. 특정 위치에 시선이 더 오래 머문다면, 그것이 바로 싱크가 증폭된 모습입니다. 어텐션을 직접 가르치지 않았는데 시선의 모양이 바뀌었다는 것이 이 방법의 묘미입니다.

저는 여기서 조금 의심했습니다. 싱크가 커지는 것이 항상 좋은 일일까요. 꼭 그렇지는 않습니다. 싱크는 일종의 닻인데, 닻이 잘 박히면 배가 안정되지만 엉뚱한 곳에 박히면 배가 엉뚱한 곳에 멈춥니다. 논문이 보여준 것은 라우팅이 닻의 위치를 움직일 수 있다는 사실이지, 닻이 항상 옳은 곳에 박힌다는 보장은 아닙니다. 그래서 다음 장면이 중요합니다. 이 간접 효과가 어떤 과제에서는 도움이 되고 어떤 과제에서는 해가 되는지, 층을 가려가며 확인하는 대목 말이죠.

3.37%포인트를 숫자 이상으로 읽는 법

이제 숫자를 볼 차례입니다. 논문은 OLMoE에서 라우팅만 미세조정하는 기준선과 비교해 GSM8K이 3.37%포인트 올랐다고 보고합니다. 일반 파일에만 남겼을 때와 SOUL.md에 옮겨 적었을 때처럼, 무엇과 비교했는지가 숫자의 의미를 정합니다. 여기서 기준선이 라우팅만 학습한 SFT라는 점이 중요합니다. 전체를 푼 상태와 겨룬 것이 아니라, 같은 조건에서 어텐션 신호를 읽었는지 여부만 갈라서 겨룬 것입니다.

이 숫자는 그대로 믿기 전에 조건을 봐야 합니다. 3.37%포인트는 수학 추론이라는 특정 과제에서의 개선입니다. 모든 과제가 고르게 좋아졌다는 뜻이 아닙니다. 뒤에 다루겠지만, 사실 회상 같은 과제에서는 오히려 흔들리는 장면도 나옵니다. 그래서 저는 이 숫자보다 그 뒤의 구조가 더 중요하다고 봅니다. 라우팅 결정에 문맥의 리듬을 넣었더니, 단계별로 이어가는 추론에서 이득이 났다는 구조 말이죠.

수학 추론이 왜 이런 신호와 잘 맞을까요. 수학 문제는 앞선 계산을 계속 참조하면서 풀어야 합니다. 중간 결과가 어디에 있는지, 그 결과를 얼마나 꾸준히 바라보고 있는지가 다음 단계의 전문가 선택과 직결됩니다. 어텐션의 시간적 흐름은 바로 그 꾸준함을 재는 자입니다. 기준선이 토큰 하나로만 판단했다면, 이 방법은 지금까지의 응시 기록까지 함께 본 셈입니다. 3.37이라는 숫자는 그 기록을 보는 일이 헛일이 아니었다는 증거로 읽힙니다. 물론 단일 benchmark의 단일 수치이므로 과장해서는 안 됩니다. 하지만 통제된 조건에서 나온 차이이기에, 가볍게 넘길 수치도 아닙니다.

모든 층에 같은 장치를 달면 왜 사실 회상이 흔들릴까

앞에서 미뤄뒀던 이야기로 돌아올 차례입니다. 이 장치를 모든 층에 무턱대고 달면 사실 회상이 나빠질 수 있다는 것입니다. 반면 수학 추론의 이득은 장치를 깊은 층에 두었을 때 유지됩니다. 같은 방법인데 과제마다 반응이 엇갈린다는 점이 흥미롭습니다. 그렇다면 층의 깊이는 무엇을 가르는 것일까요.

제 해석은 이렇습니다. 얕은 층은 문장의 표면, 이름과 날짜와 같은 구체적인 조각을 다루는 데 바쁩니다. 여기서 어텐션의 리듬을 섣불리 섞으면 정확한 조각을 집는 일이 방해받을 수 있습니다. 사실 회상은 조각을 정확히 집는 일에 가깝습니다. 반면 깊은 층은 이미 모인 조각들을 엮어 추론의 방향을 정하는 곳입니다. 여기서는 리듬이 도움이 됩니다. 어디를 오래 봤는지, 어디에서 요동쳤는지가 다음 추론 단계의 선택과 직결되기 때문입니다.

이 비대칭은 실제 설계에 바로 쓸 수 있는 교훈을 줍니다. 좋은 장치를 모든 곳에 뿌리는 것이 능사가 아니라는 것입니다. 층마다 하는 일이 다른데 같은 처방을 내리면 어딘가에서는 탈이 납니다. 논문은 층을 가려가며 장치를 붙이는 실험을 통해 이 점을 보여줍니다. 저는 이 대목에서 논문의 태도가 읽힌다고 느꼈습니다. 만능 장치를 주장하는 대신, 어디에서 쓰고 어디에서 빼야 하는지를 함께 보고했기 때문입니다. 이런 보고가 있어야 다음 연구자가 같은 실수를 반복하지 않습니다. 사실 회상이 흔들렸다는 대목은 실패담이 아니라 사용 설명서에 가깝습니다.

라우팅을 층별 탐침으로 쓰고 나서 남는 질문

마지막으로 논문의 시선을 한 번 더 비틀어보겠습니다. 층을 가려가며 이 장치를 붙이는 실험은 성능 비교를 넘어 하나의 탐침으로 읽을 수 있습니다. 각 층의 어텐션이 라우팅에 쓸 만한 정보를 얼마나 품고 있는지를 재는 막대라는 것입니다. 깊은 층에서 이득이 남고 얕은 층에서 잡음이 생겼다는 것은, 층마다 어텐션이 품은 정보의 결이 다르다는 뜻이기도 합니다.

그렇다면 다음 질문은 자연스럽게 이어집니다. 이 탐침을 다른 모형이나 다른 과제에도 꽂아보면 무엇이 보일까요. 수학 추론이 아닌 코드 생성이나 긴 문서 요약에서는 어느 층의 리듬이 라우팅을 도울까요. 저는 이 질문이 논문이 남긴 가장 큰 선물이라고 봅니다. 새로운 부품을 하나 제안한 것을 넘어, 부품을 이용해 모형 내부를 들여다보는 방법을 덤으로 줬기 때문입니다.

물론 여기까지 일반화하기에는 데이터가 부족합니다. 보고서가 보여준 것은 OLMoE에서의 통제된 실험과 몇 가지 과제의 반응입니다. 다른 MoE 구조에서도 같은 층별 패턴이 나타나는지, 윈도우 길이나 주파수 요약 방식을 바꾸면 결론이 달라지는지는 아직 모릅니다. 그럼에도 방향은 분명합니다. 라우터와 어텐션을 분리된 부품으로 보던 시각에서, 서로 읽고 영향을 주고받는 한 쌍으로 보는 시각으로 넘어갔다는 것입니다. 제가 앞으로 보고 싶은 것은 이 한 쌍을 서빙 현실에서 어떻게 다룰지에 대한 답입니다. 슬라이딩 윈도우를 유지하는 메모리와 계산 비용은 얼마인지, 배치 추론에서 리듬 요약이 병목이 되지는 않는지 말이죠. 그 답이 나오면, 이 논문은 읽기 좋은 논문을 넘어 쓰기 좋은 설계로 남을 것입니다.