토큰별 코드 하나를 모든 층이 공유해 PLE를 100배 압축한다

무엇인가

이 논문은 언어 모델의 임베딩 용량을 Transformer 연산과 함께 키우는 흐름에서, Gemma 4의 Per-Layer Embeddings(PLE)를 극단적으로 압축하는 문제를 다룬다. PLE는 토큰마다 디코더 층별로 별도 벡터를 두는 방식이다. E4B에서 PLE는 전체 파라미터의 35.2%, E2B에서는 45.8%를 차지하고, E4B는 폭 256의 토큰-층 벡터를 약 1,100만 개 저장한다. 약 100배 압축이면 벡터당 5바이트 안팎만 남는데, 기존 양자화기는 토큰-층 벡터마다 별도 코드를 부여하므로 이 예산에서 성능을 크게 잃는다. 나머지 모델 파라미터는 고정한 채 PLE 테이블만 압축하는 것이 문제 설정이다.

어떻게 동작하나

OPTIC은 토큰당 하나의 양자화 코드 튜플을 저장하고, 각 층의 전용 affine decoder가 이를 층별 벡터로 푼다. 토큰 t의 latent z_t는 M개 코드북 그룹에서 선택한 codeword를 이어 붙여 만들고, 층 l의 복원 벡터는 A_l z_t + b_l로 계산한다. 코드 할당 페이로드는 V×L이 아니라 V×M에 비례해 층 수만큼의 중복을 제거하고, decoder는 층마다 있지만 vocabulary 크기에는 비례하지 않는다. E4B의 약 100배 설정은 Z=644, M=161, K=256으로 전체 56,404,992바이트를 차지해, 비압축 BF16 PLE 5,637,144,576바이트 대비 99.94배 압축이다. 추론 시에는 요청된 토큰의 M개 codeword만 모아 해당 층 decoder를 적용하며, 전체 V×L×d 텐서를 만들지 않는다.

무엇과 다른가

핵심은 저장 벡터 좌표가 아니라 PLE가 모델에 실제로 기여하는 출력 경로에서 오차를 줄이는 것이다. PLE는 각 층의 gated residual branch에서 hidden state로 계산한 gate와 원소별 곱을 한 뒤 출력 투영을 통과한다. OPTIC은 정규화와 residual 합 이전의 투영 출력 차이 Δr = W_out(τ g ⊙ (e_hat - e))를 최소화하는 손실을 쓴다. 코드북과 decoder는 연속 업데이트하고 토큰 코드는 이산 재할당하는 교대 최적화를 수행하며, 재할당에는 fitting record의 gate를 평균한 층별 Gbar 행렬을 사용해 벡터별 오차를 가중한다. 초기화는 토큰별 L개 PLE 행을 이어 붙인 뒤 층별 Gbar로 스케일링하고 PCA, OPQ식 eigenvalue 배분, 그룹별 K-means로 잡는다. 이 밖에 행 단위 출력경로 가중 PCA-PQ와, 빈도 높은 토큰에 코드 그룹을 더 주는 frequency-tiered OPTIC 변형도 제시한다.

어떻게 쓰나

실험은 PLE에 맞춰 적용한 6개 기존 양자화 베이스라인(PQ, RVQ, OPQ-P, AQ, AQLM Euclidean beam encoder, INT8 randomized low-rank SVD)과 output-path-weighted PCA-PQ, OPTIC을 비교한다. 보정 데이터는 FineWeb, FineWeb-Edu, OpenWebMath, FineWeb2에서 뽑은 512토큰 시퀀스 1,536개(786,432 위치)에, vocabulary 전체 토큰당 하나의 합성 레코드를 더한다. 자연어 레코드는 비압축 모델에서 캐시한 문맥 의존 gate를, 합성 레코드는 all-ones gate를 쓴다. 평가는 PIQA, GSM8K, GPQA Diamond, AIME 2026, MMLU-Redux, IFBench, InfographicVQA의 7개 과제에서 참조 모델 대비 성능 유지율을 1로 캡한 뒤 동일 가중 기하평균으로 집계한다. W4A16 실험은 공식 양자화 인식 학습 파이프라인의 비양자화 체크포인트에서 PLE 압축 artifact를 fitting해 공식 W4A16 릴리스에 적용하며, PLE lookup은 BF16으로 남는다.

전제와 한계

약 100배 PLE 압축에서 OPTIC은 BF16 E4B의 7개 과제 집계 성능을 97.8% 유지하고, W4A16 E4B에서는 96.9%를 유지한다. 같은 조건에서 가장 강한 기존 베이스라인인 PQ는 BF16 E4B에서 66.3%에 그쳤고, W4A16에서는 가장 강한 기존 베이스라인이 82.8%를 기록했다. 논문은 E4B에서 각각 31.5%p와 14.0%p 앞선다고 보고한다. 과제별로 BF16 E4B의 GSM8K는 OPTIC 92.8%, 비압축 92.9%, PQ 73.5%였고, IFBench는 41.7%, 44.0%, 16.0%, InfographicVQA는 67.5%, 69.7%, 38.4%였다. 95% paired-bootstrap 구간은 BF16 E4B에서 OPTIC 95.5–99.0%, PQ 60.6–70.5%였고, W4A16에서는 OPTIC 93.9–97.8%, AQLM Euclidean beam encoder 76.8–87.1%로 분리된다. E2B에서는 frequency-tiered 변형이 BF16 90.0%, W4A16 93.0%로 가장 높고, 기본 OPTIC은 89.1%와 91.5%를 유지한다.

압축률을 더 높이면 OPTIC은 BF16 E4B에서 594배에 90.5%, 1,002배에 84.1%를 유지하고, W4A16 E4B에서는 87.7%와 81.8%를 유지한다. 같은 BF16 E4B에서 행 단위 output-path-weighted PCA-PQ는 약 100배에서 93.3%를 내지만 약 500배에서 25.1%, 약 1,000배에서 0%로 무너진다. 즉 벡터당 바이트가 줄수록 토큰 코드를 층 간 공유하는 OPTIC의 이점이 커진다. 코드 공유 ablation에서 토큰-층마다 독립 latent code를 쓰면 BF16 E4B 79.3%, W4A16 86.8%로 OPTIC의 97.8%, 96.9%보다 크게 낮다. decoder rank를 20으로 제한하면 70.55%까지 떨어져(제한 없음 95.09%), 층별 decoder의 표현력도 중요하다. fitting 목적함수는 Euclidean fitting보다 1.98%p, diagonal weighting보다 0.31%p 높은 paired mean retention을 보였고, 코드 재할당은 세 paired fit 모두에서 0.65–1.14%p 개선했다.

PLE 구조 분석도 압축이 단순 평균 제거가 아님을 보여준다. 층별 vocabulary 평균은 테이블 제곱 Frobenius norm의 8.1%를 설명하지만, 모든 벡터를 층 평균으로 바꾸면 GSM8K가 92.9%에서 43.5%, GPQA Diamond가 56.9%에서 23.0%로 떨어지고 AIME 2026은 0이 되어 7과제 유지율이 0이 된다. 저장 좌표에서 층별 토큰 고유 변동의 90%를 잡으려면 256개 주성분 중 212–227개가 필요하지만, 출력 경로 메트릭에서는 1–205개(중앙값 158)로 층마다 큰 차이를 보인다. 그래서 출력 경로 기준 fitting과 층별 decoder가 필요하다는 것이 저자들의 해석이다.

개발자에게 이 방법은 Gemma 4 E2B/E4B처럼 PLE가 파라미터의 큰 비중을 차지하는 모델을 vLLM으로 서빙할 때 유용하다. RTX 4090에서 고정 1GiB KV 캐시로 BF16 E4B를 약 100배 compact PLE로 서빙하면 GPU 메모리가 29.7% 줄고, 8개 prefill/generation 워크로드에서 wall time 증가는 BF16 0.40%, W4A16 0.69% 이하였다. 다만 이 압축은 gated PLE 출력 경로에 특화되어 있어 보정용 gate와 출력 투영, 바이트 예산, 대상 체크포인트가 맞아야 한다. 저자들은 한계로 E2B/E4B만 평가했고 주요 비교와 Table 3가 단일 fitting seed를 쓰며, paired-bootstrap 구간은 벤치마크 샘플링 불확실성을 나타낸다고 밝힌다. 코드 공유 ablation은 동일 바이트 예산에서 코드와 decoder의 바이트 배분 변화까지 포함한 결합 효과이고, 출력 경로 손실은 gate를 고정하고 각 층 투영 출력에서만 오차를 측정한다. 베이스라인은 원래 방법을 PLE에 맞게 변형한 것이며, 표본 수가 적은 과제 하나가 유지율을 크게 흔들 수 있다.