컴백부터 K-뷰티까지 — K-컬쳐의 모든 것을 메일로 받아보세요메일로 받아보기

METAL MEDIA

각 층이 자기 깊이의 기억만 보는 트랜스포머 구조를 뇌의 백질처럼 층끼리 직통 연결해 성능을 올리다

arXiv:2608.184862026-08-20

WhiteMatter: All-to-All Cross-Layer Connections via KV Mixing

각 층이 자기 깊이의 기억만 보는 트랜스포머 구조를 뇌의 백질처럼 층끼리 직통 연결해 성능을 올리다

트랜스포머는 원래 각 층이 자기 층에서 만든 정보만 저장해서 다음 토큰 처리에 쓰는데, WhiteMatter는 라우터라는 작은 신경망으로 모든 층의 정보를 섞어 소수의 공유 채널에 담고 각 층이 그중 하나를 골라 읽게 한다. 이렇게 하면 얕은 층도 더 깊은 층이 만든 정보를 참고할 수 있고, 채널 수를 층 수보다 줄이면 메모리로 저장해야 하는 캐시 크기도 줄어든다. 같은 조건에서 학습했을 때 층수가 50% 많은 일반 트랜스포머보다 성능이 좋았고, 캐시를 절반으로 줄여도 그 이득의 상당 부분을 유지했다.

METAL MEDIA 해설 도표

각 층이 자기 깊이의 기억만 보는 트랜스포머 구조를 뇌의 백질처럼 층끼리 직통 연결해 성능을 올리다

  1. 01문제의식: 트랜스포머는 이전 토큰을 다시 볼 때 그 토큰이 만든 모든 층의 정보 중 자기 층과 같은 깊이에서 나온 정보만 쓸 수 있어, 이미 계산해 둔 깊은 층의 정보를 활용하지 못한다.
  2. 02방법: 각 토큰마다 라우터가 모든 층의 은닉 상태를 읽어 k개의 공유 채널로 섞고, 각 층은 정해진 규칙에 따라 그중 하나의 채널을 읽어 어텐션에 사용한다. 채널 수 k를 층 수 L보다 작게 하면 저장해야 할 캐시가 줄어든다.
  3. 03훈련 시 어려움: 한 토큰의 채널을 만들려면 그 토큰의 모든 층 계산이 끝나야 하는데, 그 계산은 다시 이전 토큰의 채널에 의존해 순환 구조가 생긴다. 이를 병렬로 풀기 위해 순서를 섞은 반복 계산 방식을 여러 번 통과시키는 절차를 도입했다.
  4. 04결과: 80억 토큰으로 학습했을 때, 같은 16층 구조에서 일반 트랜스포머 대비 검증 퍼플렉시티(문장 예측이 얼마나 헷갈리는지 나타내는 지표, 낮을수록 좋음)를 8.2% 낮췄고 24층 일반 모델보다도 낮았다. 캐시를 절반으로 줄여도 6.3% 낮은 퍼플렉시티를 유지했으며, 같은 캐시 크기의 비교 모델(LCKV)보다도 5.0% 낮았다.
  5. 05속도: 반복 계산 순서를 조절한 방식이 순수 순차 계산보다 최대 13.9배, 다른 병렬 반복 방식보다 11.2배 빠르게 수렴했다.
METAL MEDIA이 원문을 바탕으로 재구성한 해설 도표이며, 논문 저자의 원문 figure가 아닙니다.

무엇을 했나

  1. 문제의식: 트랜스포머는 이전 토큰을 다시 볼 때 그 토큰이 만든 모든 층의 정보 중 자기 층과 같은 깊이에서 나온 정보만 쓸 수 있어, 이미 계산해 둔 깊은 층의 정보를 활용하지 못한다.
  2. 방법: 각 토큰마다 라우터가 모든 층의 은닉 상태를 읽어 k개의 공유 채널로 섞고, 각 층은 정해진 규칙에 따라 그중 하나의 채널을 읽어 어텐션에 사용한다. 채널 수 k를 층 수 L보다 작게 하면 저장해야 할 캐시가 줄어든다.
  3. 훈련 시 어려움: 한 토큰의 채널을 만들려면 그 토큰의 모든 층 계산이 끝나야 하는데, 그 계산은 다시 이전 토큰의 채널에 의존해 순환 구조가 생긴다. 이를 병렬로 풀기 위해 순서를 섞은 반복 계산 방식을 여러 번 통과시키는 절차를 도입했다.
  4. 결과: 80억 토큰으로 학습했을 때, 같은 16층 구조에서 일반 트랜스포머 대비 검증 퍼플렉시티(문장 예측이 얼마나 헷갈리는지 나타내는 지표, 낮을수록 좋음)를 8.2% 낮췄고 24층 일반 모델보다도 낮았다. 캐시를 절반으로 줄여도 6.3% 낮은 퍼플렉시티를 유지했으며, 같은 캐시 크기의 비교 모델(LCKV)보다도 5.0% 낮았다.
  5. 속도: 반복 계산 순서를 조절한 방식이 순수 순차 계산보다 최대 13.9배, 다른 병렬 반복 방식보다 11.2배 빠르게 수렴했다.
Figure 4: Held-out language-modeling quality versus non-embedding parameter count at an 8B-token budget. The connected vanilla points form the depth-scaling reference; point labels report per-token KV-cache size relative to the L=16 vanilla model. WhiteMatter is shown in half- and full-cache configurations; the LCKV configurations have four and seven warmup layers.
Figure 4: Held-out language-modeling quality versus non-embedding parameter count at an 8B-token budget. The connected vanilla points form the depth-scaling reference; point labels report per-token KV-cache size relative to the L=16 vanilla model. WhiteMatter is shown in half- and full-cache configurations; the LCKV configurations have four and seven warmup layers.
Table 1: Downstream evaluation. LAMBADA and WikiText report perplexity; the remaining columns report normalized accuracy in percent. Bold denotes the best result among the 16-layer models.
ModelLAMBADA ↓WikiText ↓PIQA ↑HellaSwag ↑ARC-E ↑OBQA ↑
Vanilla 16L127.4749.3460.8831.6747.3929.00
LCKV w=4107.5248.8162.5732.5245.6631.20
LCKV w=7102.9749.0262.2432.4046.2130.00
WhiteMatter k=871.5844.4062.3533.6145.7129.60
WhiteMatter k=1660.7343.2863.5533.8046.2129.40
Vanilla 24L97.4044.7162.7333.2147.9431.80
Vanilla 32L79.3941.4463.8234.3547.9032.20
Figure 5: Prefill convergence wall time versus group count g. Jacobi (g=1) and autoregressive evaluation (g=T) form the two endpoints. The 4-layer model was trained with exact autoregressive execution at length 1024; evaluation used T=2048 and the same channel-read policy for pass selection and timing.
Figure 5: Prefill convergence wall time versus group count g. Jacobi (g=1) and autoregressive evaluation (g=T) form the two endpoints. The 4-layer model was trained with exact autoregressive execution at length 1024; evaluation used T=2048 and the same channel-read policy for pass selection and timing.
Table 2: Measured per-token FLOPs for training, prefill, and decoding. Each pair of columns reports GFLOPs per token and the ratio to the 16-layer vanilla model.
ModelTrainingPrefillDecode
GFLOP/tok×GFLOP/tok×GFLOP/tok×
Vanilla 16L0.4441.000.1421.000.1791.00
Vanilla 24L0.6651.500.2121.500.2691.50
Vanilla 32L0.8872.000.2832.000.3592.00
LCKV w=41.4213.200.9356.610.1730.97
LCKV w=71.1742.650.7385.210.1750.97
WhiteMatter k=81.0282.320.4323.050.1770.99
WhiteMatter k=161.1112.500.4673.300.1841.03
Figure 6: Training-schedule and pool-rank ablations. (a) Model performance across training iteration schedules. Each cell reports one combination of iteration parameters, averaged over two seeds. The vertical axis nests the number of no-gradient passes nno-grad (inner labels) within the number of gradient-carrying passes ng (outer labels); the horizontal axis shows the training iteration schedule. (b) Test perplexity versus pool rank for WhiteMatter and two ablations.
Figure 6: Training-schedule and pool-rank ablations. (a) Model performance across training iteration schedules. Each cell reports one combination of iteration parameters, averaged over two seeds. The vertical axis nests the number of no-gradient passes nno-grad (inner labels) within the number of gradient-carrying passes ng (outer labels); the horizontal axis shows the training iteration schedule. (b) Test perplexity versus pool rank for WhiteMatter and two ablations.
Table 3: Complete zero-shot downstream results. LAMBADA and WikiText report perplexity; all other columns report accuracy in percent.
ModelLAMBADAWikiTextPIQAWino.BoolQHella.ARC-EARC-COBQA
Vanilla 16L127.4749.3460.8850.0448.1731.6747.3925.0029.00
LCKV w=4107.5248.8162.5750.9148.9332.5245.6623.7231.20
LCKV w=7102.9749.0262.2452.4160.6432.4046.2125.9430.00
WhiteMatter k=871.5844.4062.3551.3860.4033.6145.7125.2629.60
WhiteMatter k=1660.7343.2863.5551.3861.0433.8046.2125.4329.40
Vanilla 24L97.4044.7162.7350.9956.6733.2147.9426.3731.80
Vanilla 32L79.3941.4463.8250.5156.1834.3547.9026.9632.20
Figure 7: Whole-brain white-matter tractography. A population-averaged human structural connectome reconstructed from diffusion MRI, rendered as fiber tracts in sagittal (left) and coronal (right) views. The tracts span the brain and arc between distant regions in every direction; color encodes local fiber orientation (red: left–right, green: anterior–posterior, blue: superior–inferior). Rendered with DSI Studio (33) from its population-averaged human template (32), built from Human Connectome Project data (5).
Figure 7: Whole-brain white-matter tractography. A population-averaged human structural connectome reconstructed from diffusion MRI, rendered as fiber tracts in sagittal (left) and coronal (right) views. The tracts span the brain and arc between distant regions in every direction; color encodes local fiber orientation (red: left–right, green: anterior–posterior, blue: superior–inferior). Rendered with DSI Studio (33) from its population-averaged human template (32), built from Human Connectome Project data (5).

왜 중요한가

채팅 모델이 답을 생성할 때 저장해 두는 기억(캐시)의 크기와 계산량은 서비스 비용과 응답 속도에 직결되는 문제다. 이 연구는 층 수를 늘리지 않고도 성능을 끌어올리면서 캐시 크기를 줄일 수 있는 방법을 보여줘, 더 적은 자원으로 더 똑똑한 모델을 돌리는 데 참고가 될 수 있다.

Figure 8: Convergence timing for the larger cyclic-trained model. This 8-layer, D=1024, k=8 model is evaluated at T=4096.
Figure 8: Convergence timing for the larger cyclic-trained model. This 8-layer, D=1024, k=8 model is evaluated at T=4096.

이 논문의 용어

  • 퍼플렉시티(perplexity) · 언어모델이 다음 단어를 얼마나 잘 예측하는지 나타내는 지표로, 낮을수록 예측을 잘한다는 뜻
  • KV 캐시 · 트랜스포머가 이전 토큰들의 정보를 다시 계산하지 않고 재사용하기 위해 저장해 두는 키(Key)와 값(Value) 데이터
  • 라우터(router) · 여러 층의 정보를 입력받아 어떻게 섞을지 가중치를 정하는 작은 신경망
  • 자기회귀 디코딩 · 이전에 생성한 토큰들을 보고 다음 토큰을 하나씩 순서대로 만들어내는 방식
  • 야코비 반복 / Gauss-Seidel 반복 · 순환적으로 얽힌 계산을 여러 번 근사적으로 풀어나가는 수치 계산 방법으로, 순서를 얼마나 섞느냐에 따라 속도와 병렬성이 달라진다

저자 · Wenbo Zhang, Xiang Ren

arXiv에서 원문 보기

최신 논문

논문 전체 보기 →

METAL MEDIA 최신 기사

그림 출처: Wenbo Zhang et al., arXiv:2608.18486, CC BY 4.0