기술 요약: 크로네커 임베딩

1. 문제 제기

현대적인 대규모 언어 모델 (LLM) 은 토큰 식별자를 연속 벡터로 매핑하기 위해 $|V| \times d_{model}$ 형태의 학습된 입력 임베딩 테이블에 의존합니다. 최첨단 규모에서 이 테이블은 컨텍스트 계산이 시작되기 전에 수억에서 수십억 개의 학습 가능 파라미터를 소비합니다. 예를 들어, 어휘 크기가 13 만 개이고 $d_{model}=4096$인 모델은 입력 임베딩에만 약 5 억 3,700 만 개의 파라미터를 할당합니다. 이는 상당한 엔지니어링 마찰을 초래합니다. 테이블은 데이터 병렬 랭크 간에 분할되어야 하며, 텐서 병렬 설정에서 통신 오버헤드를 발생시키고, 옵티마이저 상태를 위해 막대한 메모리 (혼합 정밀도 Adam 의 경우 파라미터당 약 16 바이트) 를 필요로 합니다. 더 나아가, 이 논문은 토큰 ID 에서 벡터로의 고정된 매핑을 학습해야 할 필요성 자체가 열린 질문이라고 주장합니다. 특히 엔지니어링 비용과 이러한 임베딩이 실제로 발달시키는 구조의 성격을 고려할 때 그렇습니다.

2. 방법론: 크로네커 임베딩

저자들은 학습된 입력 임베딩 테이블을 위한 결정론적이고 구조화된 대체안인 **크로네커 임베딩 (Kronecker Embeddings)**을 제안합니다. 이 방법은 표준 BPE 또는 SentencePiece 토큰화를 유지하지만, 룩업 테이블을 고정된 인코더와 단일 학습된 투영 (projection) 으로 대체합니다.

핵심 메커니즘

UTF-8 바이트 시퀀스 $b = (b_1, \dots, b_L)$을 가진 토큰에 대해 임베딩은 다음과 같이 계산됩니다:

1. 바이트 - 위치 인코딩: 토큰을 구성 바이트와 그 위치로 분해합니다.
2. 크로네커 곱: 결정론적 코덱 벡터 $\kappa(b) \in \mathbb{R}^D$를 구성합니다. 여기서 $D = d_c \times d_p$입니다.
   • $d_c = 256$ (전체 바이트 알파벳).
   • $d_p$는 최대 바이트 위치 (예: 32) 입니다.
   • 코덱은 원 - 핫 바이트 벡터와 원 - 핫 위치 벡터의 크로네커 곱들의 합입니다: $\kappa(b) = \frac{1}{\sqrt{L}} \sum_{p=1}^L c_{b_p} \otimes p_p$.
3. 정규화: 코덱 출력은 초기 학습을 안정화하기 위해 토큰별 $z$-정규화 (평균 0, 표준편차 1) 를 거칩니다.
4. 학습된 투영: 단일 학습 가능한 선형 레이어 $W_{proj} \in \mathbb{R}^{D \times d_{model}}$가 코덱 벡터를 모델의 은닉 차원으로 투영합니다.

운영 변형

• gpu_table: 시작 시 전체 $|V| \times D$ 코덱 테이블을 미리 계산합니다. 계산량은 무시할 수 수준 (가athering) 이지만, 메모리 사용량은 높습니다 (수 기가바이트).
• gpu_dynamic: 순전파 (forward pass) 중 컴팩트한 바이트 버퍼 (131k 어휘 기준 약 4.5 MB) 를 사용하여 코덱을 실시간으로 재계산합니다. 이는 작은 계산 오버헤드 (스텝 시간의 0.010.24%) 를 대규모 메모리 절감으로 교환합니다.

귀납적 편향

이 방법은 바이트 수준 국소성을 도입합니다. 두 문자열이 동일한 위치에서 바이트를 공유할 때만 유사한 임베딩을 받습니다. 이는 표준 BPE 임베딩이 조형적 변형을 클러스터링하는 경향과 달리, 대소문자 구분 (예: swift 대 SWIFT) 과 오타에 대한 견고성 (예: seperate 대 separate) 을 유지합니다.

3. 주요 기여

A. 교차 모델 프로브 결과

저자들은 1 억 3,500 만 개에서 6,710 억 개 파라미터에 이르는 6 개의 현대적 언어 모델 (Llama-3.2, Qwen3, DeepSeek-V3 등 포함) 의 입력 임베딩을 분석했습니다:

• 조형적 클러스터링: 학습된 BPE 임베딩은 형태적 관련어 (예: running, runner) 보다 훨씬 더 조형적 변형 (예: run, Run, .run) 을 클러스터링합니다. 이 패턴은 토크나이저 계열과 규모를 막론하고 유지됩니다.
• 크로네커 탈출: 크로네커 임베딩은 이러한 조형적 클러스터링을 탈출하여 바이트가 유사한 문자열을 검색합니다.
• 형태론적 구조: 임베딩 층에서 두 방법 모두 엄격한 형태론적 클러스터링을 발달시키지 않습니다. 1 억 2,400 만 개 체크포인트에 대한 계층적 프로브에서 BPE 방식의 첫 번째 트랜스포머 레이어는 형태론적 기하학을 공발생/맥락적 기하학과 교환하는 반면, 크로네커는 초기 레이어를 통해 바이트 수준 기하학을 유지합니다.

B. 통제된 훈련 비교

FineWeb-Edu 의 25 억 토큰으로 훈련된 nanoGPT GPT-2 124M 아키텍처에 대한 3 시드 비교:

• 성능: 크로네커는 BPE-연동 (tied) 베이스라인보다 2.5 ± 0.2% 낮은 검증 손실을 달성했습니다 (0.083 nats 차이, 약 9% 낮은 퍼플렉시티).
• 수렴: 손실 차이는 훈련 중 벌어졌으며 수렴 시 좁아지는 대신 안정화되었습니다.
• 효율성: 크로네커는 BPE 베이스라인의 수렴 손실에 도달하는 데 1.43 배 적은 옵티마이저 스텝이 필요했습니다.
• 파라미터 수: 1 억 2,400 만 개 규모에서 입력 측 학습 가능 파라미터는 91% 감소했습니다 (3,860 만 개에서 310 만 개로). 다만, 크로네커의 $D \neq d_{model}$로 인해 구조적으로 출력 헤드의 가중치 연동이 불가능하여 분리된 출력 헤드가 필요해 전체 학습 가능 카운트는 약간 (310 만 개) 증가했습니다. 최첨단 규모 (예: 1,200 억 MoE) 에서는 입력 측 감소 폭이 **94%**에 달합니다 (5 억 300 만 개 파라미터 절감).

C. 견고성 및 충실도 프로브

• 철자 견고성: 110 개의 깨끗한/오타 프롬프트 쌍에서 크로네커는 **55.5%**의 경우 상위 1 예측을 유지한 반면, BPE 는 47.3% 였습니다 (+8.2%p). 또한 깨끗한 데이터와 오타 데이터 분포 간의 KL 발산은 낮았으며 (-7.6%), 최종 은닉 상태의 코사인 유사도는 높았습니다 (+2.6%).
• 생성 충실도: 정성적 프로브에 따르면 크로네커 모델은 자동회귀 생성을 통해 바이트로 새로운 문자열과 오타 (예: kronekticus, netwrok) 를 "반향"하는 반면, BPE 모델은 이를 분할하고 잊어버렸습니다.

D. 메커니즘 관찰

• 안정된 목표: BPE 임베딩 노름은 훈련 중 std=0.020 에서 0.026 으로 이동하는 반면, 크로네커 투영 노름은 약 1.0 에서 안정적으로 유지됩니다. 이는 크로네커가 트랜스포머 본체에 대한 안정적인 표현적 목표를 제공함을 시사합니다.
• 메모리 효율성: 실시간 변형은 생산 어휘 크기에서 입력 경로 메모리 발자국을 ~2.15 GB(미리 계산된 테이블) 에서 ~4.5 MB(바이트 버퍼) 로 줄입니다.

4. 중요성 및 주장

이 논문은 크로네커 임베딩이 다음과 같은 표준 입력 임베딩을 위한 파라미터 효율적인 플러그 - 인 교체안을 제공한다고 주장합니다:

1. 임베딩 병목 현상 제거: 최첨단 규모에서 입력 측 학습 가능 파라미터를 91~94% 감소시켜 트랜스포머 본체를 위한 용량을 확보하거나 옵티마이저 상태 메모리를 줄입니다.
2. 훈련 역학 개선: 통제된 환경에서 더 낮은 검증 손실과 더 빠른 수렴을 달성하며, 이는 안정적인 표현적 목표와 파라미터 효율적인 최적화 때문일 가능성이 있습니다.
3. 견고성 강화: 조형적 오류에 대한 내재적 견고성을 제공하고 생성 시 바이트 수준의 충실도를 유지합니다. 이는 표준 BPE 임베딩에 결여된 특성입니다.
4. 임베딩 역할 재정의: 학습된 임베딩이 형태론을 포착해야 한다는 가정을 재검토하며, 대신 주로 조형적 변형을 포착함을 보여줍니다. 크로네커는 독특한 바이트 수준의 귀납적 편향을 제공합니다.

저자들은 이 방법이 토크나이저 대체품도, 문자 단위 모델도 아니며, 토큰 시퀀스를 기반으로 작동하지만 임베딩 룩업을 대체한다고 명시합니다. 가중치 연동 불가, compute/commute와 같이 의미적으로 멀지만 바이트가 유사한 단어들의 클러스터링, 그리고 10 억 개 이상의 파라미터 규모와 다국어 데이터셋에서의 추가 검증 필요성 등 한계를 인정합니다.

5. 한계 및 향후 작업

이 논문은 몇 가지 열린 질문을 식별합니다:

• 확장성: 10 억~70 억 개 이상 규모에서 검증 손실 우위와 계층적 결과가 유지되는지 여부.
• 분리된 베이스라인: 통제된 비교는 BPE-연동 베이스라인을 사용했습니다. 향후 작업은 BPE-분리 베이스라인과 비교해야 합니다.
• 메커니즘: 성능 향상의 정확한 메커니즘 (안정된 목표 대 파라미터 효율성 대 귀납적 편향) 은 여전히 열린 질문입니다.
• 출력 측: 현재 설계는 출력 헤드를 표준 학습 행렬로 남겨둡니다. 저자들은 향후 작업이 어휘 없는 해독을 위해 크로네커를 출력 측으로 확장할 수 있다고 가정하지만, 이는 아직 테스트되지 않았습니다.

이 연구는 코드, 재현 가능한 스크립트, 그리고 규모에서의 안정성을 입증하기 위한 4 개의 사전 훈련된 LightningLM 모델 (20 억~1,200 억 MoE) 과 함께 공개되었습니다.