Nestor: 어텐션과 전문가 선택을 겹쳐 학습 계산을 줄이는 모델

배경과 문제

언어모델을 키우는 목적은 더 많은 계산을 쓰는 것 자체가 아니라, 같은 계산으로 예측을 더 잘하는 데 있다. 조건부 계산은 입력마다 가중치 일부만 사용하는 방법이다. MoE는 각 토큰을 소수의 전문가 MLP에 보내 전체 저장 용량에 비해 토큰당 계산량을 작게 유지한다. 여기서 MLP는 토큰별 특징을 변환하는 부분이고, 어텐션은 다른 토큰의 정보를 모으는 부분이다.

일반적인 MoE는 MLP를 나누지만 어텐션 파라미터는 모든 토큰이 공유한다. 긴 문서에서는 토큰 사이 관계를 계산하는 비용이 커진다. 이 논문은 어텐션에 대해서도 입력마다 다른 연산자와 참고 문맥을 선택하고, 그 선택을 전문가 선택과 연결하면 어떤 이득이 생기는지 묻는다. 단순히 기존 어텐션을 빠르게 실행하는 구현 최적화가 아니라 모델 구조를 바꾸는 접근이다.

기존 접근의 한계

희소 어텐션은 참고할 토큰을 줄이는 방법이다. 고정된 가까운 구간만 보기도 하고, 입력 내용에 따라 묶음이나 관련 토큰을 고르기도 한다. 이런 방법은 보통 어텐션 맵을 줄이면서 파라미터는 공유한다. 반대로 어텐션 헤드의 혼합은 사용할 파라미터를 고르면서 전체 문맥을 볼 수 있다. 따라서 해싱 기반 방법까지 모두 입력 내용과 무관하다거나, 헤드 혼합에는 계산 절약이 전혀 없다고 말할 수 없다.

Nestor는 참고할 토큰과 사용할 어텐션 파라미터를 함께 고르고, 선택한 모듈 안에 전문가 MLP까지 넣는다. 기존 방식과 구별되는 지점은 선택을 블록의 여러 층에서 유지하며 두 희소성을 중첩한다는 것이다. 새 분야를 학습하면 기존 분야의 성능이 떨어지는 망각 문제도 평가하지만, 공유 어텐션만이 실제 망각의 원인이라고 입증한 것은 아니다.

연구방법을 단계별로

Nestor의 핵심 아이디어는 어텐션 연산자를 MoE가 MLP를 일반화한 것과 같은 방식으로 일반화하는 것이다. 먼저, 각 층의 어텐션 연산자를 C개의 '컬럼(column)'이라 불리는 별도의 어텐션 연산자로 대체한다. 각 토큰은 컬럼 라우터(column router)에 의해 Kc개의 컬럼에 배정된다. 컬럼 c의 어텐션은 그 컬럼에 배정된 과거 토큰들(Uc)을 대상으로 계산되므로, 어텐션 맵이 희소해지고 동시에 해당 토큰은 Kc개 컬럼의 어텐션 파라미터만 사용한다. 이는 하나의 라우팅 결정이 어텐션 맵과 파라미터를 동시에 희소화하는 '컬럼 라우티드 어텐션(column-routed attention)'이다.

다음으로, 각 컬럼 안에 자체 전문가 MLP들을 둔다. 컬럼 c에는 E개의 전문가가 있고, 컬럼 내 전문가 라우터가 각 토큰에 대해 Ke개의 전문가를 선택한다. 따라서 전문가는 (컬럼, 전문가) 쌍으로 인덱싱되며, 토큰은 자신이 배정된 컬럼의 전문가들만 사용한다. 이로써 희소성이 계층적으로 겹쳐진다: 먼저 컬럼을 선택하고, 선택된 컬럼 안에서 전문가를 선택한다. 또한 모든 컬럼이 공유하는 소수의 '공유 전문가(shared experts)'도 두어, 컬럼 간 공통 지식을 학습할 수 있게 한다.

계산 효율을 위해 층들을 M개씩 블록으로 묶고, 컬럼 라우팅은 블록의 입력에서 한 번만 계산하여 블록 내 모든 층에서 동일한 컬럼 선택을 유지한다. 반면 전문가 선택은 층마다 새로 계산된다. 실험에서는 C=8, 첫 블록 Kc=3에서 마지막 블록 Kc=8(전체 컬럼 활성)로 깊이에 따라 활성 컬럼 수를 늘리는 스케줄을 사용한다. 마지막 블록에서는 모든 컬럼이 활성화되어 전체 시퀀스에 대한 완전한 어텐션을 수행한다. 이는 초기 층에서는 일부 토큰을 선택하는 희소한 처리를, 후반 층에서는 전역적 처리를 하는 기존 LLM의 인터리브 어텐션과 유사하지만, 컬럼의 문맥은 입력에 따라 달라지며 시퀀스 어디에서든 토큰을 포함할 수 있다는 차이가 있다.

라우터는 보조 손실 없이 균형을 맞추기 위해 역전파 없이 업데이트되는 편향(bias)을 사용한다. 또한 잔차 연결을 '어텐션 잔차(attention residual)'로 대체하여, 각 층의 입력이 토큰 임베딩, 이전 블록 출력, 현재 블록의 컬럼 스트림 평균의 학습된 소프트맥스 혼합이 되도록 한다. 비교 모델에도 어텐션 잔차를 적용했다. 다만 전문가 구성과 라우팅 등 여러 설계 선택이 함께 있으므로, 추가 절제 실험과 함께 해석해야 한다.

결과와 숫자의 의미

사전학습 실험은 100M, 200M, 500M, 1B 활성 파라미터 규모의 dense, MoE, Nestor 모델을 DCLM-baseline 데이터셋에서 문맥 길이 4096, 활성 파라미터당 50토큰으로 학습시켰다. 이 주 비교에서 Nestor는 모든 규모에서 가장 낮은 held-out 손실을 기록했다. 1B는 활성 규모의 이름이며 최대 모델의 실제 저장량은 Nestor 5.53B, MoE 5.49B, dense 1.10B로 다르다. 예를 들어 1B 규모에서 Nestor의 손실은 2.571 nat, MoE는 2.592 nat, dense는 2.743 nat였다. 각 아키텍처의 손실-계산 곡선을 L(F)=AF^{-α}+L∞로 피팅하여, Nestor와 같은 손실에 도달하기 위해 기준 모델이 필요한 계산량을 추정했다. 그 결과 MoE는 Nestor의 1.13~1.29배, dense는 3.1~4.1배의 계산이 필요했다. 이 비율은 규모가 커져도 줄어들지 않았으며, 최대 1B에서 MoE는 1.22배, dense는 4.1배였다. 이는 동일 손실까지의 필요 계산량을 스케일링 법칙으로 추정한 값이며, 실제 벽시계 시간이나 처리 속도를 측정한 것은 아니다.

지속적 사전학습의 도메인 적응 실험에서는 500M Nestor, 500M MoE, 그리고 사전학습 손실이 비슷한 1B dense 모델을 먼저 32k 문맥으로 확장한 후, 책(PG-19), 수학(OpenWebMath), 코드(AlgebraicStack), 독일어(FineWeb-2) 각각 0.5B 토큰으로 별도 적응시켰다. 희소 모델의 라우터는 동결하고 나머지 파라미터만 학습했다. 새 도메인 손실 감소(적응 이득)는 세 모델 모두 MoE 대비 1% 이내로 비슷했다. 그러나 기존 DCLM 손실 증가(망각)는 Nestor가 MoE보다 평균 3% 적었고, dense 모델은 21% 더 컸다. 25%의 기존 데이터를 재생(replay)하면 세 모델 모두 망각이 약 70% 감소했지만 적응 이득은 약 6% 줄었고, Nestor는 MoE보다 망각이 5% 적었다. 이는 Nestor의 컬럼 구조가 새 도메인 학습 시 기존 도메인의 표현을 덜 손상시킨다는 것을 시사한다.

긴 문맥 확장 실험에서는 4k로 사전학습된 모델을 8k, 16k, 32k, 64k로 각각 독립적으로 확장했다. 동일 토큰 수(0.5B) 조건에서는 Nestor의 MoE 대비 이점이 사전학습 때와 비슷하게 유지되었다. 그러나 동일 FLOPs 조건에서는 Nestor의 토큰당 FLOPs가 문맥이 길수록 MoE보다 작아지므로(64k에서 약 0.80~0.84배), 절약된 계산으로 더 많은 토큰을 학습할 수 있다. 이 추가 토큰을 반복 데이터에 쓰면 손실이 더 낮아졌고, 새 데이터에 쓰면 가장 큰 향상을 보였다. 예를 들어 1B 규모 64k에서 동일 FLOPs 새 데이터 조건의 Nestor 손실은 MoE보다 0.058 nat 낮았다. 이는 문맥이 길어질수록 Nestor의 계산 효율성 이점이 커진다는 것을 보여준다.

구체적인 예시

설명용으로 코드와 수학이 섞인 문서를 생각해보자. 코드 선언을 참고할 토큰은 한 컬럼으로, 수학의 가정을 참고할 토큰은 다른 컬럼으로 보낼 수 있다. 각 컬럼은 자신에게 배정된 과거 토큰을 읽고 그 안의 전문가로 출력을 계산한다. 이는 직관을 위한 예시다. 실제 학습된 컬럼에 사람이 정한 코드·수학 역할이 그대로 생겼다고 확인한 실험 결과는 아니다.

논문의 이론 분석은 이러한 계층 구조를 단순한 선형 모형으로 설명한다. 거친 라벨은 참고할 문맥과 통계량을, 세부 라벨은 그 통계량을 사용하는 방식을 결정한다. 중첩 예측기는 같은 컬럼의 여러 전문가에 관련된 표본을 함께 활용해 표현을 추정한다. 해당 가정에서 평면적인 예측기보다 추정해야 할 유효 차원이 작아지는 결과를 보인다. 라벨이 알려져 있고 선형 어텐션을 쓰는 모형이므로 실제 Nestor가 모든 데이터에서 우월하다는 증명은 아니다.

한계와 오해하기 쉬운 점

이 연구의 가장 큰 한계는 규모와 평가 지표다. 최대 모델이 1B 활성 파라미터에 불과하며, 평가는 held-out 손실만 사용했다. 실제 다운스트림 벤치마크나 대규모 프론티어 모델에서의 성능, 실서비스 속도는 검증되지 않았다. 또한 1.13~1.29배(MoE)와 3.1~4.1배(dense)의 계산 효율성은 스케일링 법칙을 피팅하여 추정한 값이다. 이는 실제 처리 속도나 벽시계 시간 개선을 측정한 것이 아니며, FLOPs 계산에는 어텐션이 포함되어 있다. 만약 선형 층의 FLOPs만 센다면 Nestor는 기준 모델보다 4~8% 더 많은 FLOPs를 사용한다. 따라서 'Nestor가 4배 빠르다'와 같은 해석은 오해다.

도메인 적응 실험에서 '망각이 3% 적다'는 것은 기존 DCLM 손실 증가량이 MoE 대비 3% 작다는 뜻이지, 모든 지식의 3%만 망각한다는 뜻이 아니다. 실제 망각량은 손실 증가로 측정되며, 예를 들어 코드 도메인 적응 후 Nestor의 DCLM 손실은 2.779에서 3.073으로 0.294 nat 증가했다. 또한 25% 재생은 총 예산 0.5B 토큰 내에서 기존 데이터 비율을 25%로 한 것이며, 망각을 약 70% 줄이는 대신 적응 이득이 약 6% 감소하는 트레이드오프가 있다. 긴 문맥 확장에서도 동일 토큰과 동일 FLOPs 조건을 구분해야 한다. 동일 토큰에서는 Nestor의 이점이 유지될 뿐 커지지 않으며, 동일 FLOPs에서 이점이 커지는 것은 절약된 계산으로 더 많은 토큰을 학습하기 때문이다.

이론 분석은 계층적 잠재 구조, 선형 어텐션, 알려진 라벨 등 단순화된 가정에 기반한다. 따라서 실제 대규모 모델의 보편적 성능을 보장하지 않는다. 또한 뇌피질의 컬럼 구조와의 연결은 비유와 동기일 뿐, 뇌 메커니즘을 증명한 것이 아니다. 실험은 대부분 시드 1개로 수행되었고, 일부 200M 설정만 시드 2개를 사용했다. 공식 README의 11B 코퍼스 예시는 논문과 다른 결과를 낼 수 있으며, 정확한 재현을 위해서는 105B 코퍼스에서 같은 예산으로 샘플링해야 한다. 마지막으로, 비교 모델에도 어텐션 잔차를 적용했으므로, Nestor의 개선이 희소성 하나만의 효과로 귀속될 수는 없다.

실제로 어떻게 활용할까

활용 관점에서는 같은 연산 예산을 어떻게 쓰는지 비교하는 틀을 배울 수 있다. 긴 문맥에서 토큰당 연산을 줄인 뒤 절약분을 추가 학습에 쓰는 조건과, 같은 토큰 수로 학습하는 조건은 다른 질문에 답한다. 새 분야에 적응할 때에도 새 자료의 손실 감소와 기존 자료의 손실 증가를 함께 기록해야 한다. 이 논문의 도메인 실험은 각각 별도의 적응이며, 여러 분야를 연속해서 배운 실제 서비스의 성능 검증은 아니다.

공개 코드와 체크포인트는 구조 변형을 연구하는 출발점이다. 실제 적용 전에는 더 큰 모델, 실제 과제의 정확도, GPU 실행 시간과 통신 비용까지 확인해야 한다. 컬럼 내부에 희소·선형 어텐션을 더하는 방법은 향후 방향으로 제시된다. 기존 컬럼을 그대로 보존하면서 새 도메인 컬럼만 학습한다고 단정하거나, 문서 요약·코딩 서비스에서 이미 효과를 입증했다고 확장해서는 안 된다.