최근 공개된 연구 논문 '트라이루트(TriRoute)'는 대규모 언어 모델(LLM)의 추론(inference) 비용을 절감하면서도 성능을 유지하는 새로운 접근 방식을 제시합니다. 기존에는 어텐션(attention) 해상도, 전문가(expert) 선택, KV 캐시(KV-cache) 비트 폭 할당 등 세 가지 주요 최적화 기법이 독립적으로 적용되었지만, 트라이루트는 이들을 하나의 경량 컨트롤러로 통합하여 각 토큰(token)에 필요한 연산 자원을 유기적으로 조절합니다.
트라이루트 컨트롤러는 모든 레이어(layer)의 각 토큰에 대해 어텐션 모드(건너뛰기/로컬/전체), 희소한 FFN 전문가 세트(MoE), 그리고 KV 캐시 비트 폭을 조율된 정책으로 결정합니다. 이 컨트롤러는 Gumbel-Softmax와 스트레이트-스루 추정(straight-through estimation)을 활용한 이질적인 완화(heterogeneous relaxation) 기법과 부하 분산(load-balanced) 톱-k 게이팅(gating)을 통해 종단 간(end-to-end) 학습됩니다. 특히, 예산 제약(Lagrangian budget constraint)을 두어 평균 연산 및 메모리 비용을 조절 가능한 노브(knob)처럼 사용할 수 있게 하여, 비용 효율성을 극대화합니다. 연구진은 개별적으로 훈련할 때 발생하는 '라우팅 붕괴 연쇄(routing-collapse cascade)' 문제를 해결하기 위해 축별 정규화(per-axis normalization)와 결합 인식 균형 손실(coupling-aware balancing loss)을 도입했습니다.
이 기술은 1.3B(13억) 매개변수(parameter) 규모의 디코더 전용(decoder-only) 모델에서 기존의 독립적인 MoD(Mixture-of-Depths), MoE(Mixture-of-Experts), KV-양자화(KV-quantization) 조합보다 우수한 파레토 효율성(Pareto-dominance)을 보였습니다. 동일한 추론 FLOPs와 메모리 사용량에서 더 나은 성능을 달성했으며, 특히 순수 퍼플렉시티(perplexity) 최적화에서 저하될 수 있는 희귀 엔티티, 코드, 산술 연산 등 '꼬리 사례(tail-case)'에 대한 견고성(robustness)을 더 잘 보존했습니다. 분석 결과, 컨트롤러는 문장 시작 위치, 희귀한 서브워드(subword), 고유 명사(named entities)에는 전체 어텐션과 고정밀 캐시를 할당하고, 기능어(function words)에는 저렴한 라우팅을 적용하는 등 해석 가능한 구조를 보여주었습니다. 이는 LLM의 추론 비용을 효과적으로 관리하면서도, 모델의 중요한 능력인 희귀 정보 처리 능력을 유지하거나 개선할 수 있음을 의미합니다. 결과적으로, LLM 서비스 제공자들은 더 적은 비용으로 고품질의 서비스를 제공할 수 있게 될 것입니다.