EntropyMoE: Entropy-Aware Sparse Expert Routing for Tokenizer-Free LLMs

기술 요약: EntropyMoE

문제 정의
최근 토크나이저가 없는 대규모 언어 모델(LLM), 특히 Byte Latent Transformer(BLT)의 발전은 국소적 예측 불확실성(엔트로피)을 기반으로 원시 바이트를 동적 크기의 패치로 그룹화함으로써 효율성을 개선했습니다. 이 아키텍처에서 예측 가능한 바이트 시퀀스는 더 긴 패치를 형성하는 반면, 불확실한 시퀀스는 더 짧은 패치를 형성합니다. 그러나 기존의 바이트 패치 아키텍처는 패치의 의미론적 난이도나 세밀도와 관계없이 모든 패치에 균일하고 밀집된 피드포워드 연산을 적용합니다. 이러한 "일률적인(one-size-fits-all)" 접근 방식은 모델 용량을 패치의 의미론적 변화에 적응시키지 못합니다. 또한, 표준 Mixture-of-Experts(MoE) 모델은 일반적으로 고차원 은닉 상태(hidden states)를 기반으로 라우팅을 수행하는 토큰 네이티브 방식이며, 이는 패칭 과정에서 이미 계산된 명시적인 불확실성 신호를 무시하고 불필요한 라우팅 오버헤드를 발생시킵니다.

방법론: EntropyMoE
본 논문은 글로벌 패치 트랜스포머의 밀집 피드포워드 레이어를 Top-K MoE 레이어로 대체하기 위해 설계된 패치 네이티브 희소 전문가(patch-native sparse expert) 아키텍처인 EntropyMoE를 제안합니다. 핵심 혁신은 고차원 은닉 상태 대신 스칼라 패치 엔트로피를 라우팅 메커니즘의 유일한 입력으로 사용하는 데 있습니다.

• 아키텍처: BLT 스타일의 파이프라인을 기반으로 구축된 EntropyMoE는 로컬 바이트 인코더, 글로벌 셀프 어텐션, 디코더를 유지하면서 밀집 피드포워드 서브레이어를 전문가 풀(expert pools)로 대체합니다. 각 동적 패치는 단일 라우팅 단위 역할을 합니다.
• 엔트로피 전용 라우팅: 은닉 상태를 전문가 로짓(logits)으로 투영하는 기존의 라우터와 달리, EntropyMoE 라우터는 학습된 아핀 변환($s = aH + b$)을 통해 패치의 정규화된 스칼라 엔트로피($H_i$)만을 전문가 로짓으로 매핑합니다. 이는 라우팅 파라미터를$O(d \times M)$에서$O(M)$으로 줄여줍니다 (여기서$d$는 은닉 차원,$M$은 전문가 수). 보고된 실험(25개 레이어, 8개 전문가)에서 이는 은닉 상태 라우터의 400,000개 이상의 파라미터와 비교하여 단 400개의 학습 가능한 라우팅 파라미터만을 사용합니다.
• Top-K 선택: 라우터는 엔트로피로부터 유도된 로짓을 기반으로 패치당 정확히 $K=2$개의 전문가를 선택합니다. 패치의 은닉 상태는 선택된 전문가에 의해 처리되지만, 선택 과정에서는 명시적으로 제외됩니다.
• 바이트 가중 작업량 산정(Byte-Weighted Workload Accounting): 패치마다 길이가 다르기 때문에 패치 수를 세는 것만으로는 계산 부하를 정확히 반영할 수 없습니다. EntropyMoE는 바이트 가중 로드 밸런싱 목적 함수를 도입합니다. 전문가에 할당된 부하는 해당 전문가에 할당된 패치 길이의 합으로 계산되며, 이를 전체 바이트 수로 정규화합니다. 이를 통해 긴(저엔트로피) 패치를 다루는 전문가가 많은 수의 짧은(고엔트로피) 패치를 다루는 전문가에 비해 과소 활용되지 않도록 보장합니다.
• 초기화: 희소 모델은 밀집 BLT 체크포인트로부터 밀집 피드포워드 가중치를 모든 전문가에게 복사하여 초기화하며, 이를 통해 모델이 연속 학습(continuation training) 중에 분기되기 전까지 동일한 동작을 시작하도록 보장합니다.

주요 기여

1. 패치 네이티브 희소 연산: 동적 바이트 패치가 근본적인 라우팅 단위가 되는 토크나이저-프리 LLM을 위한 희소 조건부 연산의 정식화로, 라우팅 결정을 고차원 의미론적 표현으로부터 분리했습니다.
2. 엔트로피 전용 라우터: 패칭 파이프라인에 내재된 인과적 불확실성 신호를 활용하는 엔트로피 전용 Top-2 라우터를 도입했습니다. 이는 패치 세밀도(패처에 의해 결정), 전문가 선택(엔트로피에 의해 결정), 작업량 산정(바이트 길이에 의해 결정)을 분리합니다.
3. 통제된 경험적 검증: 동일한 데이터, 업데이트 횟수, 활성 파라미터 수 및 초기화 조건 하에서 Dense BLT, Hidden-only MoE, Hidden+Entropy MoE와 엄격하게 비교했습니다.

실험 결과
저자들은 FineWeb-Edu 데이터셋을 사용하여 BLT-1B 백본에 대해 EntropyMoE를 평가했습니다.

• 모델링 품질: EntropyMoE는 테스트된 모든 구성 중 가장 낮은 홀드아웃 Bits Per Byte (BPB) (0.8351)를 달 기록하여, Dense BLT 베이스라인(0.8442) 및 두 가지 은닉 상태 기반 MoE 변형 모델을 모두 능가했습니다. 이 개선은 통계적으로 유의미했습니다 (95% 신뢰 구간이 0 미만에 위치).
• 다운스트림 성능: 다운스트림 정확도는 전반적으로 베이스라인과 유사한 수준을 유지했습니다. 특히, EntropyMoE는 HellaSwag 태스크에서 Dense BLT 대비 +1.16%의 통계적으로 유의미한 향상을 보였으며, PIQA 및 MMLU와 같은 다른 태스크에서도 경쟁력 있는 성능을 유지했습니다.
• 라우팅 효율성: EntropyMoE는 은닉 상태 라우터와 비교했을 때 라우팅 파라미터를 3 자릿수(약 1,000배) 감소시켰습니다(400 vs ~400k). 테스트된 희소 모델 중 가장 좋은 반복 시간(iteration time)과 유효 바이트 처리량(valid-byte throughput)을 달성했으나, MoE 디스패치 및 통신 오버헤드로 인해 Dense BLT보다는 약 2배 느렸습니다.
• 라우팅 패턴: 분석 결과, EntropyMoE는 매우 선택적인 라우팅 구조를 유도하는 것으로 나타났습니다. 할당을 광범위하게 분산시키는 은닉 상태 라우터와 달리, EntropyMoE는 소수의 전문가 집합에 할당을 집중시켰습니다 (상위 2개 전문가가 질량의 89.1%를 차지). 또한 저엔트로피 패치와 고엔트로피 패치 사이의 라우팅 동작에서 명확한 차이를 보여주었으며 (Hidden-only의 0.0040 대비 Entropy의 JSD는 0.1263), 이는 엔트로피가 전문가 활용을 효과적으로 차별화함을 나타냅니다.
• 작업량 산정: 바이트 가중 로드 밸런싱은 패치 수 기반 산정보다 BPB를 일관되게 개선하였으며, 이는 가변 길이 패치 시나리오에서 바이트 커버리지를 고려하는 것의 필요성을 입증했습니다.

의의 및 주장
본 논문은 패치 엔트로피가 토크나이저-프리 LLM의 희소 조건부 연산을 위한 효과적이고 저차원적인 라우팅 좌표 역할을 한다고 주장합니다. 패칭을 위해 이미 계산된 불확실성 신호를 활용함으로써, EntropyMoE는 전체 의미론적 표현을 라우팅 결정에 소비하지 않고도 희소 연산을 조직합니다.

저자들은 결과가 패치 엔트로피를 조건부 연산의 유망한 기초로 확립하고 개선된 바이트 수준 모델링을 입증했다는 점을 명시하면서도, 필수적인 전문가 정체성이나 밀집 백본 대비 연산량 매칭 우위를 확립했다고는 주장하지 않습니다 (현재의 희소 구현은 여전히 더 느리기 때문입니다). 관찰된 개선은 단순한 희소 용량 때문이 아니라 엔트로피 조건부 라우팅 기하학과 바이트 인식 작업량 균형 덕분인 것으로 간주됩니다. 은닉 상태 MoE 변형 모델들이 유사한 용량을 가졌음에도 동일한 이득을 얻지 못했기 때문입니다. 이 연구는 MoE 모델링을 토크나이저 기반 표현 너머로 확장하여, 더욱 효율적이고 적응 가능한 바이트 수준 언어 모델링을 위한 경로를 제시합니다.