🤖 machine learning

TriRoute: Unified Learned Routing for Joint Adaptive Attention, Experts, and KV-Cache Allocation

TriRoute는 단일 경량 컨트롤러를 통해 모든 토큰에 대해 어텐션 해상도, 전문가 선택 및 KV-캐시 비트 너비를 공동으로 최적화하는 통합 학습형 라우팅 프레임을 도입하여, 개별적인 조건부 연산 방식들과 비교하여 우수한 파레토 효율성과 강건성을 달성합니다.

원저자: Andrii Balashov, Olena Ponomarova

게시일 2026-07-09

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Andrii Balashov, Olena Ponomarova

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대하고 고급스러운 레스토랑 주방의 매니저라고 상상해 보세요. 고객이 요리(문장의 단어를 의미함)를 주문할 때마다, 당신의 주방은 그 요리를 준비하는 데 얼마나 많은 노력을 기울일지 결정해야 합니다.

전통적인 대규모 언어 모델(LLM)에서 주방은 경직되어 있습니다. "the"와 같은 흔한 단어든, "Nakamura"와 같은 희귀하고 복잡한 이름이든, 모든 단어에 대해 똑같은 처우를 합니다. 셰프는 동일한 강도로 칼질하고, 젓고, 접시에 담습니다. 이는 엄청난 낭비입니다.

TriRoute라는 논문은 이 새로운 스마트한 매니저를 소개합니다. 이 매니저는 식사의 질을 망치지 않으면서도 비용을 절감하기 위해, 모든 단어에 대해 세 가지 유형의 결정을 동시에 내릴 수 있습니다.

주방의 세 가지 레버 (The Three Levers of the Kitchen)

저자들은 이전 방식들이 단 하나의 레버만 한 번에, 즉 독립적으로 조절하여 비용을 아끼려 했다고 주장합니다. 반면 TriRoute는 이 세 가지 레버를 동시에 당깁니다.

"어텐션(Attention)" 레버 (얼마나 뒤를 돌아볼 것인가):
- 기존 방식: 모든 단어가 문장의 이전 모든 단어를 살펴봅니다.
- TriRoute의 방식: "and"와 같이 지루한 단어에 대해서는 매니저가 "뒤를 돌아보는 과정을 건너뛰라"고 말합니다. 이름과 같이 핵심적인 단어에 대해서는 "맥락을 이해하기 위해 문장 전체를 보라"고 말합니다.
- 비유: 메뉴 전체를 읽을지, 아니면 가격만 슥 훑어볼지를 결정하는 것과 같습니다.
"전문가(Expert)" 레버 (얼마나 많은 두뇌를 사용할 것인가):
- 기존 방식: 모든 단어가 모델의 완전하고 무거운 두뇌에 의해 처리됩니다.
- TriRoute의 방식: 단순한 단어의 경우 "널(null) 전문가"에게 보냅니다(기본적으로 아무것도 하지 않고 그냥 통과시킴). 복잡한 단어의 경우 최고의 전문가들에게 보냅니다.
- 비유: 간단한 주문은 그냥 적기만 하는 웨이터에게 보내고, 복잡하고 맞춤화된 요리는 헤드 셰프에게 보내는 것과 같습니다.
"메모리(Memory)" 레버 (얼마나 명확하게 기억할 것인가):
- 기존 방식: 모든 단어를 높은 정밀도로(금색 잉크로 쓰는 것처럼) 주방의 메모리 로그에 기록하여 많은 공간을 차지합니다.
- TriRoute의 방식: 단순한 단어는 연필로 쓰고(낮은 정밀도, 공간 절약), 희귀하고 중요한 단어는 금색 잉크로 써서 나중에 주문이 들어왔을 때 쉽게 찾을 수 있도록 합니다.
- 비유: 나중에 필요할 가능성에 따라 영수증을 먼지 쌓인 상자에 보관할지(저품질), 아니면 금고에 보관할지(고품질) 결정하는 것과 같습니다.

문제점: 잘못된 결정의 "도미노 효과"

저자들은 컴퓨터가 이 세 가지 결정을 따로 배우거나, 도움 없이 한꺼번에 배우게 하면 문제가 발생한다는 것을 발견했습니다.

그들은 이를 **"붕괴 카스케이드(Collapse Cascade)"**라고 부릅니다.
매니저가 게을러져서 모든 뒤돌아보기(Attention)를 건너뛰기로 결정했다고 가정해 봅시다. 갑자기 두뇌 전문가들(Experts)은 지루하고 동일한 데이터의 흐름을 받게 되어 단어들을 차별화하려는 노력을 멈춥니다. 그러면 메모리 관리자(Bits)는 아무것도 명확하게 저장할 이유가 없다고 판단하여 모두를 연필로 쓰게 만듭니다. 전체 시스템이 저렴하고 저품질인 버전으로 붕괴됩니다.

해결책: 저자들은 세 명의 결정권자가 모두 활발하고 다양하게 움직이도록 강제하는 특별한 "균형 잡기"(수학적 레시피)를 만들었습니다. 이는 마치 코치가 팀에게 "모두가 힘든 일을 피하기로 결정해서는 안 됩니다. 팀이 강하게 유지되도록 노력을 분산시켜야 합니다"라고 말하는 것과 같습니다.

결과: 더 똑똑하고 저렴한 주방

연구진은 소형부터 중형 크기의 모델까지 테스트를 진행했습니다. 결과는 다음과 같습니다.

더 나은 가치: 엄격한 예산(예: "평소 컴퓨터 전력의 55%만 사용하고 메모리의 40%만 사용")을 설정했을 때, TriRoute는 세 가지 별개의 시스템의 다이얼을 단순히 낮추는 기존 방식보다 훨씬 더 나은 결과를 만들어냈습니다.
"희귀한 것들" 보호: 이것이 가장 중요한 발견입니다. 기존 방식은 비용을 아끼기 위해 희귀하고 어려운 단어(이름, 코드, 수학 문제 등)를 소홀히 다루어 모델이 실수를 유발하곤 했습니다. TriRoute는 어려운 일에 에너지를 아껴 쓰는 법을 배웠습니다. 쉬운 단어("the", "is")는 건너뛰지만, 희귀한 단어("Nakamura", "quantum")를 위해서는 모든 역량을 쏟아붓습니다.
해석 가능성: 매니저가 실제로 무엇을 하고 있는지 살펴보았을 때, 그것은 매우 타당했습니다. 매니저는 문장의 시작 부분, 희귀한 이름, 숫자에 "금색 잉크"와 "전체 주의력"을 쏟아부은 반면, "the"와 같은 흔한 단어는 매우 저렴하게 처리했습니다.

핵심 요약

TriRoute는 단일하고 가벼운 "매니저"로서, 모든 단어에 대해 정확히 얼마만큼의 주의력, 두뇌, 그리고 메모리가 필요한지 결정하는 법을 배웁니다. 이 세 가지 결정을 따로가 아닌 함께 내림으로써, 모델의 어려운 작업 수행 능력을 유지하면서도 실행 비용을 약 절반으로 줄입니다.

이는 모든 제품을 똑같이 취급하는 공장과, 어떤 제품에는 고급 마감이 필요하고 어떤 제품은 빠르고 저렴하게 만들 수 있는지 아는 스마트한 공장의 차이와 같습니다.

기술 요약: TriRoute

문제 정의

대규모 언어 모델(LLM)을 배포하는 데는 막대한 추론 비용이 발생하지만, 현재의 자원 할당 방식은 통계적으로 낭비가 심합니다. 조밀한(Dense) 트랜스포머는 의미적 중요성과 관계없이 모든 토큰에 균일한 부동 소수점 연산(FLOPs)과 Key/Value(KV) 캐시 메모리를 적용합니다. 문법적으로 사소한 기능어(function word)는 문단의 의미를 지탱하는 희귀한 고유 명사와 동일한 자원을 소비합니다.

기존의 조건부 연산 기술들은 이러한 균일성을 깨뜨리려 시도하지만, 개별적인 축(axis) 상에서 독립적으로 작동합니다:

MoE (Mixture-of-Experts): 토큰을 일부 전문가(expert) 집합으로 라우팅하여 피드 포워드 네트워크(FFN)를 희소화합니다.
MoD (Mixture-of-Depths): 예측 가능한 것으로 간주되는 토큰에 대해 전체 트랜스포머 블록(어텐션 + FFN)을 건너뜁니다.
KV-캐시 양자화 (KV-Cache Quantization): 키(Key)와 값(Value)을 더 낮은 비트 너비로 저장하여 어텐션 메모리를 압축합니다.

본 논문은 이러한 결정들이 **강하게 결합(strongly coupled)**되어 있다고 주장합니다. 예를 들어, 희귀한 엔티티 토큰은 미래의 쿼리에 의해 충실히 검색될 수 있도록 완전한 어텐션 해상도와 높은 정밀도의 캐시 저장이 필요할 수 있지만, 그 FFN 변환은 예측 가능하여 건너뛸 수 있을 수도 있습니다. 반대로, 기능어는 세 가지 차원 모두에서 건너뛸 수 있는 대상일 수 있습니다. 이들을 별도의 하이퍼파라미터를 가진 독립적인 최적화 문제로 취급하면 이러한 상관관계를 포착하지 못해, 최적의 자원 할당에 실패하고 "테일 케이스(tail-case)" 강건성(희귀 엔티티, 코드, 산술 연산에서의 성능)을 저하시킬 수 있습니다.

방법론: TriRoute

저자들은 세 가지 고립된 메커니즘을 하나의 경량화된 학습형 컨트롤러로 대체하는 통합 아키텍처인 Tri-Route를 제안합니다. 모든 레이어 $\ell$ 의 모든 토큰 $t$ 에 대해, 컨트롤러는 세 가지 이종 축(heterogeneous axes)에 걸쳐 조정된 정책 $\pi^{(\ell)}_t = (a_t, e_t, b_t)$ 를 생성합니다:

어텐션 해상도 ( $a_t$ ): 건너뛰기(0 비용), 로컬(윈도우 어텐션), 또는 풀(전체 인과적 어텐션) 중 하나를 선택합니다.
전문가 선택 ( $e_t$ ): 토큰을 희소한 FFN 전문가 집합 또는 **널 전문가(null expert, $f_0 \equiv 0$ )**로 라우팅합니다. 널 전문가는 모델이 어텐션은 유지하면서 FFN만 건너뛸 수 있게 해주며, 이는 MoD가 표현할 수 없는 영역입니다.
KV-캐시 비트 너비 ( $b_t$ ): 미래의 쿼리를 위해 해당 토큰의 Key/Value 쌍을 기록할 때의 정밀도(2, 4, 8, 또는 16비트)를 결정합니다.

기술적 구성 요소

공유 컨트롤러: 2개 층의 MLP 트렁크(trunk)가 정규화된 잔차 상태(residual state)와 저렴한 인과적 사이드 특징(위치, 거리, 예측 엔트로피)을 처리하여 세 개의 별도 헤드(Attention, Experts, Bits)를 위한 로짓(logits)을 생성합니다. 이 설계는 공유된 "토큰 중요도" 개념이 모든 축에 정보를 제공하면서도 각기 다른 출력 서브스페이스를 유지할 수 있게 합니다.
이종 완화 (Heterogeneous Relaxation): 학습에는 이산적 결정(discrete decisions)을 미분할 수 있는 기법이 필요합니다. 저자들은 다음을 사용합니다:
- 범주형 어텐션 및 비트 너비 결정을 위한 Gumbel-Softmax 및 Straight-Through (ST) 추정기.
- 전문가 선택을 위한 Top-k Softmax 게이팅.
- 교차 축 그래디언트 균형 (Cross-Axis Gradient Balancing): 어텐션을 건너뛰는 것이 캐시 비트를 줄이는 것보다 손실(loss)에 훨씬 큰 영향을 미치기 때문에, 저자들은 (그래디언트 노름의 EMA를 사용하는) 축별 그래디언트 리스케일링 메커니즘을 도입하여 어텐션 축이 공유 트렁크를 지배하는 것을 방지합니다.
결합 인식 균형 (Coupling-Aware Balancing): 한 축에서의 붕괴가 다른 축의 붕괴를 유발하는 "붕괴 카스케이드(collapse cascade)"(저분산 입력으로 인한 현상)를 방지하기 위해 다음을 채택합니다:
- 축별 화이트닝 (Per-axis Whitening): 각 헤드에 전달되기 전 트렁크 특징을 분리된 방식으로 정규화합니다.
- 엔트로피 하한선 (Entropy Floor): 옵션 사용 엔트로피가 특정 임계값 아래로 떨어질 때만 활성화되는 페널티를 적용하여, 예산 제약과 싸우지 않으면서도 다양성을 보장합니다.
미분 가능한 비용 모델 및 라그랑주 제어: 단일 라그랑주 제약 조건이 FLOPs와 메모리를 모두 포함하는 목표 평균 예산 $C^\star$ 를 강제합니다. 듀얼 변수(Lagrange multiplier)가 온라인으로 업데이트됨으로써, 실무자는 세 개의 별도 희소성 하이퍼파러미터를 튜닝하는 대신 단일 스칼라 노브(knob)만으로 전체 파레토 프런티어를 탐색할 수 있습니다.

주요 기여

통합된 정식화: TriRoute는 어텐션 해상도, 전문가 선택, KV-캐시 정밀도를 단일 컨트롤러와 예산 하에 공동으로 학습하는 최초의 아키텍처입니다.
안정적인 공동 학습 레시피: 저자들은 축별 화이트닝, 엔트로피 하한선, 그래디언트 균형을 통해 "교차 축 라우팅 붕괴 카스케이드"를 식별하고 완화함으로써, 이종 이산 결정의 엔드 투 엔드 학습을 가능하게 했습니다.
설계 공간 탐색: 저자들은 공유 트렁크와 별도의 헤드를 갖는 구조가 최적의 입도(granularity)임을 입증하며, 이는 별도의 라우터를 사용하거나 단일 공유 헤드를 사용하는 것보다 우수합니다. 또한, per-head 어텐션 라우팅과 토큰 레벨 FFN/비트 라우팅의 조합이 "스윗 스팟(sweet spot)"임을 보여줍니다.
테일 케이스 보존: 평균 퍼플렉시티(perplexity) 최적화 방식의 적응형 모델들이 흔히 희귀 입력에서 성능이 저하되는 것과 달리, Tri-Route는 자원을 필요한 곳에 정확히 할당함으로써 희귀 엔티티, 코드, 산술 연산에 대한 강건성을 유지합니다.

결과

컴퓨팅 최적화된 토큰 수에서 학습된 160M에서 1.3B 파라미터 규모의 디코더 전용 모델들을 대상으로 평가했습니다:

파레토 우위 (Pareto Dominance): TriRoute는 독립적으로 튜닝된 최적의 조합(MoD + MoE + KV-양자화)을 엄격히 능가하는 비용-품질 프런티어를 그려냅니다. 약 55%의 FLOPs와 40%의 KV 메모리 매칭 예산에서, Tri-Route는 밀집(dense) 모델의 다운스트림 정확도의 96~99%를 회복합니다.
성능 향상: 독립적인 조합과 비교했을 때, Tri-Route는 모든 규모에서 퍼플렉시티를 0.3~~0.4 개선하고 다운스트림 정확도를 0.7~~1.0 포인트 향상시킵니다.
테일 강건성: 희귀 엔티티, 코드, 수학 버킷에서 독립적인 조합은 상당한 저하(예: 희귀 엔티티에서 +2.1 퍼플렉시티)를 보였습니다. 반면 Tri-Route는 밀집 모델의 0.4 퍼플렉시티 이내를 유지하며, 공동 조율이 고가치 입력에 대한 "조용한 성능 저하"를 방지함을 입증했습니다.
해석 가능성: 컨트롤러가 다음과 같은 해석 가능한 패턴을 학습함을 분석을 통해 밝혀냈습니다:
- 읽기/쓰기 충실도: 희귀 서브워드와 고유 명사는 완전한 어텐션과 높은 정밀도의 캐시를 할당받습니다.
- 디커플링 (Decoupling): 희귀 엔티티는 종종 FFN을 건너뛰면서(널 전문가) 어텐션과 고정밀 캐시는 유지하는데, 이는 블록 단위의 MoD로는 불가능한 동작입니다.
- 깊이 특화 (Depth Specialization): 초기 레이어는 전반적으로 높은 어텐션을 유지하는 반면, 후기 레이어는 비핵심 토큰에 대해 어텐션을 건너뛰는 등 더 선택적인 모습을 보입니다.

의의 및 주장

본 논문은 공유된 학습형 컨트롤러가 고정된 추론 예산을 중요한 곳에 사용하는 원칙적인 방법임을 주장합니다. 세 가지 축의 조건부 연산을 통합함으로써 Tri-Route는 수동으로 튜닝된 고립된 메커니즘의 비효율성을 피합니다.

저자들은 자신들의 결과가 "적절한" 계산량이란 토큰당 단일 스칼라가 아니라, 이종 자원에 걸친 결합된 선택의 벡터라는 점을 시사한다고 강조합니다. 통합된 접근 방식은 파레토 프런티어를 개선할 뿐만 아니라, 어떤 입력이 절감을 위해 비용을 지불할지를 근본적으로 변화시켜 희귀하고 중요한 토큰이 과소 지원되지 않도록 보장합니다.

마지막으로 저자들은 실제 배포에 관한 겸허한 노트를 남겼습니다. 이론적인 FLOP 절감 효과는 상당하지만, 실제 벽시계 시간(wall-clock time) 가속도는 현재 시스템 수준의 격차(래기드 어텐션, 혼합 정밀도 캐시 관리, 라우터 오버헤드)로 인해 제한적입니다. 저자들은 이 FLOP-to-latency 격차를 좁히는 것을 실질적인 채택을 위한 주요 다음 단계로 보고 있습니다.