Gated Recurrent Transformers: Expressive Depth through Recurrent Modulation
목차
개요
Gated Recurrent Transformer(GRT)는 트랜스포머의 깊이(depth)와 파라미터 수를 분리하기 위해 제안된 재귀 깊이(recurrent depth) 아키텍처다. 고정 깊이의 prelude 블록과 coda 블록이 하나의 공유 core를 감싸고, 그 core를 R번 반복 적용한다. 핵심 기여는 GRU에서 착안한 경량 projection과 elementwise update gate를 도입해 매 재귀 단계마다 서로 다른 입력을 구성하는 것이다.
기존 weight sharing은 계속 변화하는 표현에 동일한 변환을 강제하기 때문에, 첫 번째 재귀 이후의 은닉 상태와 여덟 번째 재귀 이후의 은닉 상태를 같은 방식으로 처리하게 된다. 논문은 이것이 깊이를 가치 있게 만드는 기능적 다양성(functional diversity)을 붕괴시킨다고 지적한다. GRT의 게이트는 현재 은닉 상태, 고정된 prelude 출력, 그리고 매 단계마다 재샘플링되는 노이즈를 조건으로 삼아 하나의 가중치 텐서가 여러 개의 특화된 레이어처럼 동작하도록 만든다.
주요 결과는 다음과 같다. isoFLOPS 제약에서 3레이어 GRT가 12레이어 GPT-2 Small 베이스라인의 품질에 도달했고, 9개의 scale-by-budget 셀 전부에서 MoR와 heavy-tail depth sampling을 앞섰다. isoPARAMS 제약에서는 동일 파라미터·동일 데이터 예산으로 검증 손실 2.76을 기록해 비재귀 대응 모델의 2.84를 개선했다. Large 스케일에서는 파라미터 62% 감소, 피크 디코딩 메모리 59% 감소를 컴파일 기준 생성 지연 10% 증가로 확보했다.
배경과 문제 정의
깊이와 파라미터의 결합
표준 트랜스포머는 레이어를 하나 추가할 때마다 새로운 가중치 집합이 함께 추가된다. 전체 고유 파라미터 수는 대략 L 곱하기 d 제곱 규모로 증가하며, 이는 고정된 하드웨어 예산에서 유효 깊이를 제약하는 메모리 병목이 된다.
논문은 메모리 효율성 외에 더 근본적인 계산적 동기를 제시한다. 유한한 상태와 심볼을 반복 적용하는 것만으로 계산 가능한 모든 것을 계산할 수 있다는 Turing의 통찰은, 추론 시스템의 힘이 파라미터의 폭이 아니라 반복의 깊이에 있음을 시사한다는 것이다. Chain-of-thought는 출력 공간에서 시퀀스 길이를 소모하며 test-time compute를 실현하지만, 재귀 깊이는 추가 토큰 생성이나 추가 파라미터 저장 없이 은닉 상태 내부에서만 더 깊이 계산하는 상보적 방식이다.
관련 연구
가중치 공유의 기원은 1986년 Rumelhart, Hinton, Williams의 역전파 논문 말미에 제안된 synchronous iterative net으로 거슬러 올라간다. RNN, LSTM, GRU는 시퀀스 차원을 따라 공유 가중치를 적용하는 반면, GRT는 네트워크 깊이 차원을 따라 공유 변환을 반복한다.
트랜스포머 시대에는 ALBERT가 모든 레이어의 가중치를 공유했고, Universal Transformers가 토큰별 적응적 halting을 추가했다. 다만 두 방법 모두 전체 레이어 스택을 공유해 입력의 표현 단계와 무관하게 동일한 변환을 강제했다. Geiping et al.은 prelude, shared core, coda에 서로 다른 역할을 부여해 이 제약을 완화했고, Koishekenov et al.은 사전학습 LLM의 초기·중간·후기 레이어가 각각 인코딩·추론·디코딩 단계를 담당한다고 보고 중간 레이어만 재귀시켰다. GRT는 이 중간 레이어 선택적 재귀 원리를 따른다.
계산량을 동적으로 조절하는 계열로는 위치별 종료 신호를 학습하는 ACT, 토큰을 레이어 부분집합으로 라우팅하는 Mixture-of-Depths, 이를 재귀 설정으로 확장한 Mixture-of-Recursions가 있다. 세 방법 모두 전용 종료 신호나 라우팅 메커니즘을 필요로 하는 반면, GRT는 라우팅이나 halting 로직 없이 공유 블록의 업데이트를 얼마나 흡수할지만 조절하는 단순한 게이팅을 사용한다.
가중치 결속 반복을 고정점으로 수렴시키는 deep equilibrium model과도 구분된다. GRT는 가중치 결속 업데이트는 공유하지만 평형 목적함수는 사용하지 않고, 고정된 수의 이산 단계를 펼쳐 역전파하며, 깊이 샘플링으로 각 단계를 단일 극한에 대한 접근이 아니라 독립적인 exit로 학습시킨다.
한편 Kaplan et al.은 파라미터 공유 트랜스포머가 동일 파라미터 기준(isoPARAMS)에서는 더 좋지만 동일 연산 기준(isoFLOPS)에서는 더 나쁘다고 관찰했다. GRT는 두 체제를 모두 다루며 isoFLOPS를 주 기여로 삼는다.
방법론
Prelude-Shared-Coda 구조
GRT는 트랜스포머 블록을 세 집합으로 분할하고 npre+nrec×R+ncoda 표기를 사용한다. 논문은 고정 컨텍스트 인코더와 공유 재귀 core를 분리하는 것이 모든 레이어를 균일하게 공유하는 것보다 학습이 안정적이라고 보고한다.
| 구성 요소 | 블록 수 | 역할 |
|---|---|---|
| Prelude | npre | 한 번만 적용되어 안정적이고 컨텍스트를 담은 조건 신호 h(pre)를 생성 |
| Shared core | nrec | R번 재귀 적용, 이전 은닉 상태와 고정된 prelude 앵커를 함께 처리 |
| Coda | ncoda | 최종 정제 상태 h(R)를 받아 logits 생성용 최종 은닉 상태로 투영 |
전체 고유 파라미터 수는 (npre + nrec + ncoda) 곱하기 d 제곱 규모이며 R과 무관하다. 예를 들어 2+5×4+2 구성은 forward pass당 2 + 5 × 4 + 2 = 24회의 블록 실행을 수행하지만 가중치는 2 + 5 + 2 = 9개 블록만 저장한다. 이는 isoFLOP으로 매칭된 24레이어 GPT-2 Medium 대비 고유 블록 수가 2.6배 적은 것이다.
시퀀스 길이 S에서 토큰당 forward pass 비용은 다음과 같다.
1
FLOPs = (npre + nrec·R + ncoda) · (24d^2 + 4Sd) + R · 10d^2
여기서 24d^2 + 4Sd는 표준 블록 비용(어텐션·MLP 투영과 두 번의 어텐션 행렬곱)이고, 10d^2는 재귀 projection Wproj와 게이트 MLP fg의 단계당 오버헤드다.
Recurrent Projection과 Elementwise Gate
각 재귀의 출력을 다음 재귀의 입력으로 그대로 넣으면 원래 입력 컨텍스트를 망각할 위험이 있고, 깊은 재귀에서 그래디언트 소실·폭발이 발생하며, 단일 공유 가중치 텐서의 표현 용량에 과도한 부담을 준다. GRT는 GRU의 게이팅에서 착안해 재귀 단계마다 은닉 표현을 상태로 취급한다. 공유 core가 상태를 읽고, 변환하고, 갱신된 상태를 다시 쓰는 구조는 시간이 아니라 깊이 위에서 동작하는 RNN 셀과 유사하다.
전체 업데이트는 다음 세 식으로 정의된다.
1
2
3
h~(r) = Wproj [ h(r-1) + eps_x , h(pre) ] (2)
o(r) = Bshared( h~(r) ) (3)
h(r) = g(r) * h(r-1) + (1 - g(r)) * o(r) (4)
여기서 대괄호는 특징 축 연결(concatenation), g(r)은 0에서 1 사이 값을 갖는 T×d 학습 가능 elementwise 게이트, Wproj는 d×2d 학습 투영, eps_x는 평균 0 분산 sigma 제곱의 주입 노이즈, 별표는 요소별 곱이다.
게이트 자체는 정규화된 현재 상태와 정규화된 prelude 출력을 조건으로 하는 소형 피드포워드 네트워크가 생성한다.
1
g(r) = sigmoid( fg([ LN(h(r-1)), LN(h(pre)) ]) / tau + eps_g ) (5)
tau는 온도 하이퍼파라미터로 모든 실험에서 1.0으로 고정했고, eps_g는 학습 중 주입되는 스칼라 단위 게이트 노이즈이며, fg는 SiLU 활성화와 은닉 차원 dgate = d를 갖는 2층 MLP다.
fg의 두 번째 선형 레이어 bias는 +4로 초기화된다. 이렇게 하면 학습 시작 시점에 g(r)이 약 0.98이 되어 식 (4)의 복사 분기가 지배하고, residual stream이 재귀를 거의 변하지 않고 통과한다. 학습이 진행되면서 모델은 어떤 요소를 덮어쓸지 스스로 학습한다. 망각 게이트의 양수 초기화는 순환 신경망에서 오래된 관행이며, GRT는 같은 원리를 시간이 아닌 깊이 축에 적용한 것이다.
토큰 t, 재귀 r에서 게이트의 i번째 요소를 보면 값이 1에 가까울 때 residual stream이 그대로 유지되고, 0에 가까울 때 블록 출력이 상태를 완전히 대체한다. 이 요소 단위 특화는 Gemma 3n과 Gemma 4가 사용하는 Per-Layer Embeddings(PLE)보다 메모리 효율적인 대안이다. PLE는 레이어별 컨텍스트를 제공하기 위해 L개의 고유 임베딩 레이어를 필요로 하지만, GRT는 게이트를 h(pre)에 조건화해 매 단계에 원본 입력에 대한 일정한 시야를 제공하면서 이전 상태와 확률성으로 이를 변조한다.
상태 노이즈는 재귀 projection 입력을 교란해 단계 간 취약한 정확 일치 패턴 학습을 억제하고, 게이트 노이즈는 게이트가 학습 중 거의 상수 값으로 붕괴하는 것을 방지한다.
학습 중 재귀 깊이 샘플링
각 학습 스텝마다 1부터 R 사이의 균일 분포에서 r을 샘플링한다. 이는 두 가지 목적을 수행한다. 첫째, 모든 exit 지점을 암묵적으로 학습시켜 보조 손실 항이나 multi-exit 모델에서 흔한 그래디언트 간섭 없이 추론 시 early exit을 가능하게 한다. 둘째, stochastic depth 정규화로 작용해 고정 깊이 스케줄 대비 최종 검증 손실을 개선한다.
Forward Pass 의사코드
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def grt_forward(x, model, R, r_min=1):
h = token_embed(x) + pos_embed(x) # [S, d]
for block in model.prelude_blocks:
h = h + block(h)
h_pre = h # 고정 앵커
r = R if not training else uniform_sample(r_min, R) # 깊이 샘플링
for step in range(r):
eps = sample_gaussian(0, sigma**2)
h_tilde = W_proj(cat([h + eps, h_pre])) # concat 후 d로 투영
o = B_shared(h_tilde) # 매 단계 동일 가중치
eps_g = sample_gaussian(0, sigma_g**2)
g = sigmoid(f_gate(LN(h), LN(h_pre)) / tau + eps_g)
h = g * h + (1 - g) * o # 게이트된 residual 업데이트
for block in model.coda_blocks:
h = h + block(h)
return lm_head(h) # [S, |V|]
실험 셋업
트랜스포머 백본은 nanoGPT 코드베이스 위에 구현했다. 모든 모델은 시퀀스 길이 T = 1024, GPT-2 BPE 토크나이저(어휘 50,257개)로 학습했다.
옵티마이저는 AdamW이며 beta1 = 0.9, beta2 = 0.95, weight decay 0.1, 2,000 스텝 선형 learning rate warmup, 그리고 최대 6e-4에서 6e-5로 향하는 코사인 스케줄을 사용했다. 그래디언트 노름은 1.0에서 클리핑했고 모든 실행은 bfloat16 혼합 정밀도로 진행했다. 모든 GRT 실행에서 상태 노이즈 0.1, 게이트 노이즈 0.1, 게이트 온도 1.0을 사용했다.
학습은 20,000 스텝, 스텝당 유효 배치 491,520 토큰(배치 크기 8, 그래디언트 누적 60, 시퀀스 길이 1024)으로 진행해 실행당 약 9.8B 토큰을 소비했다. Dense 베이스라인인 GPT-2 Small, Medium, Large는 동일한 설정으로 처음부터 학습했으며 사전학습 가중치는 사용하지 않았다.
재귀 계열 경쟁자 네 가지도 동일 레시피로 처음부터 학습했다. Mixture-of-Recursions(MoR), Geiping et al.의 prelude-coda 레이아웃에 heavy-tail Poisson 깊이 샘플링을 적용한 변형, 재귀별 LoRA 어댑터로 가중치 결속을 완화한 Relaxed Recursive Transformers(RRT), 그리고 모든 루프 반복을 감독하는 Ouro다.
학습 구성
| Run | p | b | R | c | d | heads | Peak LR | Steps | Hardware |
|---|---|---|---|---|---|---|---|---|---|
| GPT-2 Small (baseline) | - | - | - | - | 768 | 12 | 6e-4 | 20k | 2x H200 |
| GRT Small (isoFLOP) | 1 | 1 | 10 | 1 | 768 | 12 | 6e-4 | 20k | 2x H200 |
| GPT-2 Medium (baseline) | - | - | - | - | 1024 | 16 | 6e-4 | 20k | 2x H200 |
| GRT Medium (isoFLOP) | 2 | 5 | 4 | 2 | 1024 | 16 | 6e-4 | 20k | 2x H200 |
| GRT Medium (isoParam) | 2 | 20 | 4 | 2 | 1024 | 16 | 6e-4 | 20k | 2x H200 |
| GPT-2 Large (baseline) | - | - | - | - | 1280 | 20 | 6e-4 | 20k | 2x H200 |
| GRT Large (isoFLOP) | 1 | 5 | 6 | 5 | 1280 | 20 | 6e-4 | 20k | 2x H200 |
| GRT Large (isoParam) | 3 | 30 | 6 | 3 | 1280 | 20 | 6e-4 | 20k | 2x H200 |
모든 재귀 실행은 게이트 bias 초기화 +4, 상태 노이즈 0.1, 게이트 노이즈 0.1, 게이트 온도 1.0을 사용했다.
주요 결과
isoFLOPS와 isoPARAMS 검증 손실
isoFLOPS 체제에서 GRT는 베이스라인 파라미터의 36~37%만 저장하면서 경쟁력 있는 품질을 달성했다.
isoFLOPS 체제 결과
| Model | Config | Layers | FLOPs/fwd | Params | Val Loss |
|---|---|---|---|---|---|
| GPT-2 Small Baseline | 12L | 12 | 1.84 G | 124M | 3.15 |
| MoR | 1+1x10+1 | 3 | 1.84 G | 45M | 3.30 |
| Heavy-tail Poisson | 1+1x10+1 | 3 | 1.84 G | 34M | 3.23 |
| Ouro | 1+1x10+1 | 3 | 1.84 G | 35M | 3.19 |
| RRT | 1+1x10+1 | 3 | 1.84 G | 33M | 3.14 |
| GRT (ours) | 1+1x10+1 | 3 | 1.84 G | 35M | 3.14 |
| GPT-2 Medium Baseline | 24L | 24 | 7.35 G | 354M | 2.84 |
| MoR | 2+5x4+2 | 9 | 7.35 G | 137M | 3.02 |
| Heavy-tail Poisson | 2+5x4+2 | 9 | 7.35 G | 124M | 2.97 |
| Ouro | 2+5x4+2 | 9 | 7.35 G | 126M | 2.93 |
| RRT | 2+5x4+2 | 9 | 7.35 G | 125M | 2.95 |
| GRT (ours) | 2+5x4+2 | 9 | 7.35 G | 127M | 2.89 |
| GPT-2 Large Baseline | 36L | 36 | 21.1 G | 774M | 2.71 |
| MoR | 1+5x6+5 | 11 | 21.1 G | 309M | 2.91 |
| Heavy-tail Poisson | 1+5x6+5 | 11 | 21.1 G | 291M | 2.93 |
| Ouro | 1+5x6+5 | 11 | 21.1 G | 290M | 2.82 |
| RRT | 1+5x6+5 | 11 | 21.1 G | 289M | 2.85 |
| GRT (ours) | 1+5x6+5 | 11 | 21.1 G | 293M | 2.77 |
Small 스케일에서 GRT는 세 개 시드 모두에서 dense 베이스라인보다 낮은 손실을 기록했다. GRT는 3.145 플러스마이너스 0.004, 베이스라인은 3.188 플러스마이너스 0.056이며, GRT의 최악 시드 3.148이 베이스라인의 최고 시드 3.154보다도 낮다. Medium과 Large에서는 표준 예산에서 dense 베이스라인이 각각 0.05, 0.06 nats 앞서지만, 토큰 예산이 늘어나면 두 격차 모두 좁혀진다.
재귀 베이스라인과의 격차는 스케일이 커질수록 벌어진다. RRT가 네 방법 중 가장 강해 Small에서 3.143 대 GRT 3.141로 시드 편차 범위 안의 동률을 기록하지만, Medium에서 0.06, Large에서 0.08 nats 뒤진다. MoR와 heavy-tail Poisson은 모든 스케일에서 0.08~0.16 nats 뒤처진다. Ouro는 Small에서 3.19로 경쟁력이 있으나 Medium은 2.93, Large는 2.82로 수렴했다.
논문은 이 격차의 원인을 각 방법이 무엇을 고정하는지에서 찾는다. RRT의 LoRA 델타는 학습 시점에 결정되어 모든 입력에 동일하게 적용되므로 확보되는 다양성이 사전에 고정되고 재귀 횟수에 따라 늘어나지 않는다. 반면 GRT의 게이트는 곧 갱신할 상태 자체를 조건으로 삼아, 균일 공유가 상태를 붕괴시키는 지점에서 연속된 상태를 분리 가능하게 유지한다.
isoPARAMS 체제 결과
| Model | Config | Layers | FLOPs/fwd | Params | Val Loss |
|---|---|---|---|---|---|
| GPT-2 Small Baseline | 12L | 12 | 1.84 G | 124M | 3.15 |
| MoR | 1+10x10+1 | 12 | 15.64 G | 135M | 3.12 |
| Heavy-tail Poisson | 1+10x10+1 | 12 | 15.64 G | 124M | 3.06 |
| Ouro | 1+10x10+1 | 12 | 15.64 G | 127M | 3.10 |
| RRT | 1+10x10+1 | 12 | 15.64 G | 125M | 3.08 |
| GRT (ours) | 1+10x10+1 | 12 | 15.64 G | 127M | 3.04 |
| GPT-2 Medium Baseline | 24L | 24 | 7.35 G | 354M | 2.84 |
| MoR | 2+20x4+2 | 24 | 26.4 G | 367M | 2.78 |
| Heavy-tail Poisson | 2+20x4+2 | 24 | 26.4 G | 356M | 2.74 |
| GRT (ours) | 2+20x4+2 | 24 | 26.4 G | 357M | 2.76 |
| GPT-2 Large Baseline | 36L | 36 | 21.1 G | 774M | 2.71 |
| MoR | 3+30x6+3 | 36 | 109.0 G | 795M | 2.69 |
| Heavy-tail Poisson | 3+30x6+3 | 36 | 109.0 G | 776M | 2.70 |
| GRT (ours) | 3+30x6+3 | 36 | 109.0 G | 779M | 2.65 |
고정 파라미터 예산에서 재귀 R을 늘리면 Medium에서 0.08 nats, Large에서 0.06 nats 개선된다. 다만 Medium isoPARAMS에서는 heavy-tail Poisson이 2.74로 GRT의 2.76보다 근소하게 앞선다. Large에서는 GRT가 2.65로 모든 비교군을 앞선다.
다운스트림 벤치마크
lm-eval-harness를 사용해 Large 스케일에서 아홉 개 벤치마크의 zero-shot 정확도를 측정했다.
Large 스케일 다운스트림 평가
| Task | GRT isoFLOPS (288M, 11L) | GPT-2 Large (774M, 36L) | GRT isoPARAMS (774M, 36L) |
|---|---|---|---|
| ARC-Challenge | 25.43 | 23.63 | 25.26 |
| ARC-Easy | 42.89 | 43.31 | 45.54 |
| BoolQ | 60.18 | 57.58 | 52.29 |
| HellaSwag | 35.62 | 37.21 | 41.52 |
| LAMBADA (OpenAI) | 39.05 | 40.05 | 45.27 |
| LAMBADA (Standard) | 30.60 | 29.73 | 38.02 |
| OpenBookQA | 28.40 | 29.80 | 30.20 |
| PIQA | 64.31 | 65.51 | 66.05 |
| Winogrande | 52.25 | 51.62 | 53.28 |
| Average | 42.08 | 42.05 | 44.15 |
isoFLOPS 변형은 파라미터의 37%만으로 dense 모델의 평균 성능과 동률을 이뤘다(42.08 대 42.05). isoPARAMS 변형은 아홉 개 과제 중 여덟 개에서 dense 베이스라인을 앞섰고 평균 +2.10점을 기록했다. 가장 큰 개선은 LAMBADA Standard의 +8.29점이며, 크게 하락한 항목은 BoolQ의 -5.29점이다.
Early Exit 분석
GRT는 보조 손실 없이도 추론 시 early exit을 창발적 속성으로 지원한다. 동일한 추론 FLOPs 조건에서 더 적은 재귀로 exit한 GRT가 더 이른 레이어에서 exit한 dense 모델보다 낮은 손실을 보인다. 균일 깊이 샘플링이 모든 중간 재귀 상태를 최종 손실 예측용으로 학습시키기 때문에, 단일 체크포인트가 재학습이나 보조 손실 없이 빠르고 얕은 추론에서 고품질 깊은 추론까지 연속적인 연산-품질 다이얼을 제공한다. 논문은 재귀 단계 절반만 실행해도 정확도의 92%를 유지한다고 보고한다.
재귀 단계에 따른 다음 토큰 예측의 정성적 변화도 분석했다. 지식 집약적 토큰은 재귀 궤적 초반에 안정화되고, 추론 기반 과제는 깊은 반복에서 점진적 오류 수정과 확률 질량 첨예화를 보이며, 개방형 프롬프트는 연속적인 언어적 정제를 보인다. 예를 들어 “The author of the book 1984 was George” 프롬프트에서 Orwell의 확률은 prelude 단독 출력 8%에서 1단계 46%, 2단계 83%, 3단계 이후 90~91%로 수렴한다.
지연 시간과 디코딩 메모리
GRT는 isoFLOPS로 매칭된 dense 모델과 동일한 수의 블록 적용을 수행하되 더 적은 고유 가중치 텐서에 분산시킨다.
생성 지연 시간과 피크 메모리 (배치 4, 프롬프트 길이 1024, 128 토큰 생성)
| Model | Config | Params | Eager (ms/tok) | Compiled (ms/tok) | Peak mem. |
|---|---|---|---|---|---|
| GPT-2 Medium | 24L | 354M | 2.95 | 2.54 | 747 MB |
| GRT Medium | 2+5x4+2 | 127M | 3.62 | 2.81 | 402 MB |
| GPT-2 Large | 36L | 774M | 4.33 | 3.33 | 1570 MB |
| GRT Large | 1+5x6+5 | 293M | 5.33 | 3.67 | 639 MB |
torch.compile 환경에서 GRT-Large는 생성 지연 10% 증가를 대가로 파라미터 62%와 피크 메모리 59%를 줄인다. Medium 스케일의 오버헤드는 11%다. Eager 모드에서는 두 경우 모두 오버헤드가 23%로, 격차의 절반가량은 산술 연산이 아니라 추가된 요소별 연산의 커널 실행 오버헤드에서 온다.
디코딩 시 GRT는 순진하게 구현하면 재귀 단계마다 하나씩 총 R개의 KV 캐시를 요구해 메모리가 R배가 된다. 논문은 마지막 단계 캐시만 재사용, 첫 단계 캐시만 재사용, 단계 간 K/V 평균이라는 세 가지 압축 전략을 평가했고 재귀 레이어 기준 3.27배 축소를 달성했다. 평균 전략이 가장 좋은 정확도를, 마지막 단계 전략이 가장 나쁜 정확도를 보였다.
종단간 디코딩 메모리 (Medium isoFLOPS 체크포인트, bf16, T = 1024)
| Configuration | HellaSwag | Mem. @ B=1 | Mem. @ B=32 |
|---|---|---|---|
| GPT-2 Medium (24L, 354M) | 34.45 | 0.75 GiB (1.00x) | 3.66 GiB (1.00x) |
| GRT isoFLOPS, 전체 R배 KV | 33.65 | 0.41 GiB (0.55x) | 3.32 GiB (0.91x) |
| GRT isoFLOPS, 평균 KV | 33.90 | 0.35 GiB (0.47x) | 1.44 GiB (0.39x) |
배치 크기 1에서는 디코딩이 가중치 지배적이라 순진한 R배 캐시로도 dense 대비 0.55배에 머문다. 배치 크기 32에서는 캐시가 지배적이 되어 순진한 전략은 0.91배까지밖에 회수하지 못한다. K/V 평균 전략은 이를 0.39배로 낮추면서 HellaSwag 정확도도 전체 캐시 대비 개선한다(33.90 대 33.65). 논문은 이 스케일에서 평균화가 약한 정규화로 작용하는 것으로 보인다고 해석한다.
Ablation
Small 모델에서 20,000 스텝 전체 예산으로 순차적 component ablation을 수행했다.
Component Ablation (Small, 20,000 스텝)
| 구성 | Val. Loss | 변화 |
|---|---|---|
| Dense 12레이어 베이스라인 | 3.157 | - |
| Recurrence만 | 3.264 | +0.107 |
| + Prelude / coda | 3.229 | -0.035 |
| + State noise | 3.211 | -0.018 |
| + Prelude 재주입 | 3.189 | -0.022 |
| + Elementwise gate (full GRT) | 3.141 | -0.048 |
재귀만 적용하면 검증 손실이 0.107 nats 악화된다. 단계를 구별할 메커니즘이 없으면 공유 블록이 초기 처리와 후기 처리를 하나의 가중치 텐서로 뭉개기 때문이다. Elementwise 게이트가 -0.048 nats로 단일 최대 기여 요소이며, 전체 모델만이 dense 베이스라인 아래로 내려간다.
Medium 구성에서 2,000 스텝만 학습한 짧은 지평 대조군에서는 마지막 두 행의 순서가 뒤집힌다.
Component Ablation (Medium, 2,000 스텝)
| 구성 | Val. Loss | 변화 |
|---|---|---|
| Dense 24레이어 베이스라인 | 3.621 | - |
| Recurrence | 4.163 | +0.542 |
| + Prelude / coda | 4.118 | -0.045 |
| + State noise | 4.099 | -0.019 |
| + Prelude 재주입 | 3.901 | -0.198 |
| + Elementwise gate (full GRT) | 3.786 | -0.115 |
짧은 지평에서는 prelude 재주입이 -0.198로 게이트의 -0.115보다 크다. 논문은 구조적 요소가 일찍 효과를 내는 반면 게이트는 학습에 따라 기여가 누적되는 학습 메커니즘이기 때문이라고 해석한다.
하이퍼파라미터 민감도도 검증했다. 게이트 온도는 Medium 5,000 스텝에서 0.5, 1.0, 2.0에 대해 각각 3.62, 3.60, 3.63을 기록해 1.0이 최적이었다. 상태 노이즈는 Medium 5,000 스텝에서 0.0일 때 0.1 대비 0.031 nats 악화, 0.2일 때 0.018 nats 악화를 보여 0.05~0.1 구간이 최적이었다.
전체 20,000 스텝 지평에서 Small 구성으로 반복한 결과도 응답 곡면이 수렴 시점에 날카로워지지 않음을 보여준다.
게이트 bias 및 상태 노이즈 민감도 (Small, 20,000 스텝)
| 설정 | Val. Loss |
|---|---|
| 게이트 bias +4 (기본) | 3.141 |
| 게이트 bias +2 | 3.152 |
| 게이트 bias 0 | 3.160 |
| 게이트 bias -2 | 3.151 |
| 상태 노이즈 0 | 3.229 |
| 상태 노이즈 0.1 (기본) | 3.211 |
| 상태 노이즈 0.2 | 3.230 |
게이트 bias를 기본 +4에서 0으로 내리면 0.019 nats 손해를 보고 -2에서는 3.151로 회복해, 곡면이 뾰족한 정점이 아니라 얕고 비단조적이다. 상태 노이즈도 제거 시 0.018, 두 배로 늘릴 시 0.019로 대칭적 최적점을 보인다. 어떤 설정에서도 손실 급증이나 발산은 발생하지 않았다.
데이터 예산에 따른 거동도 분석했다. 작은 스케일에서 GRT는 관측된 모든 데이터 볼륨에서 dense 베이스라인보다 우월한 Pareto frontier를 유지한다. 큰 스케일에서는 저데이터 영역에서 dense 모델이 앞서다가 토큰 예산이 늘어나면서 격차가 크게 좁혀지는 교차 현상이 나타난다. Large 구성에서는 높은 데이터 예산에서 두 궤적이 수렴하며, 논문은 재귀 공유가 충분한 데이터가 주어질 때 파라미터 용량을 포화시키는 데 특히 효과적임을 시사한다고 해석한다.
설계 속성과 학습 시간 비교도 제시된다.
재귀 깊이 방법론 설계 비교 (Small, 20,000 스텝 기준)
| Property | Dense | MoR | Poisson | RRT | Ouro | GRT |
|---|---|---|---|---|---|---|
| 완전 가중치 공유 | n/a | Yes | Yes | No (LoRA) | Yes | Yes |
| 가변 추론 깊이 | No | Yes | Yes | No | Yes | Yes |
| 입력 의존 게이팅 | No | Yes | No | No | Yes | Yes |
| 단계별 노이즈 주입 | No | No | 초기화만 | No | No | Yes |
| 반복별 손실 | No | Yes | No | No | Yes | No |
| 창발적 early exit | No | Yes | Yes | No | Yes | Yes |
| 학습 시간 | 3h18m | 12h19m | 4h16m | 6h42m | 8h20m | 6h18m |
| Val. loss | 3.15 | 3.30 | 3.23 | 3.14 | 3.19 | 3.14 |
GRT는 완전 가중치 공유, 입력 의존 업데이트, 가변 추론 깊이를 모두 결합한 유일한 방법이다. 다만 학습 시간은 dense 열과 직접 비교할 수 없다. Dense 베이스라인만 torch.compile을 켜고 학습했고, 자체 방법을 포함한 모든 재귀 방법은 해당 하드웨어에서 컴파일이 동작하지 않아 비활성화한 상태로 실행했기 때문이다.
재귀 동역학 분석
Large 1+5×6+5 체크포인트로 공유 블록이 각 재귀 단계에서 무엇을 계산하는지 분석했다. 진단은 추론 노이즈를 끈 상태로 검증 분할에서 수행했다.
손실 수렴은 강하게 앞단에 몰려 있다. Prelude 출력만으로는 손실이 5.29이고, 공유 core를 한 번 적용하면 3.77로 1.52 nats 감소한다. 2단계에서 3.14에 도달하며, 나머지 네 단계는 0.46 nats만 추가로 기여해 최종 2.68로 수렴한다. 즉 첫 두 재귀 단계가 prelude 표현에서 최종 출력까지의 전체 손실 감소분의 약 77%를 차지한다. 최종 단계 출력 분포에 대한 KL 발산은 prelude에서 2.61, 2단계에서 0.46, 5단계에서 0.014로 떨어져 출력 분포가 4단계 이후 사실상 확정된다.
게이트의 유용성에 대한 하한도 측정했다. 게이트를 0으로 강제하면 손실이 12.03으로 파국적으로 나빠지고, 게이트를 1로 강제해 모든 블록 제안을 폐기하면 5.26으로 prelude 단독 출력과 동일해진다. 두 극단이 학습된 모델의 2.68을 감싸고 있으며, 학습된 게이트는 재귀를 건너뛰지도 상태를 무분별하게 덮어쓰지도 않는 세밀한 균형을 획득한 것이다.
게이트는 write-heavy에서 copy-heavy로 전환된다. 중간 단계(2~3단계)에서 평균 약 0.82, 표준편차 약 0.18로 가장 열려 있고 가장 변동이 크며, 6단계로 갈수록 평균 0.87, 표준편차 0.12로 조여진다. 게이트 값이 0.95를 초과하는 copy-포화 차원은 1단계 19.7%에서 6단계 28.8%로 단조 증가하는 반면, 0.05 미만의 write-포화 차원은 전 구간에서 1e-4 미만으로 희귀하다. 유효 게이트 개방도는 1단계 0.182에서 6단계 0.066으로 단조 감소한다.
업데이트 동역학은 크기가 줄고 방향은 정렬되는 양상이다. 블록의 원시 제안 노름은 여섯 단계 내내 93~99로 거의 일정한데, 실제 residual stream에 들어가는 게이트 적용 업데이트는 1단계 17.6에서 6단계 6.5로 줄어든다. 이 발산은 전적으로 게이트가 닫히면서 발생한다. 제안 코사인 유사도는 1단계 0.19에서 6단계 0.51로 증가하고, 적용 업데이트의 코사인 유사도는 1단계에 이미 0.94이며 6단계에 0.99에 접근한다. 단계별 이득 추정치는 2단계에서 10.4로 정점을 찍고 6단계에서 4.7로 감소하는 오목한 프로파일을 보인다.
CKA 분석은 재귀 단계 간 표현이 동일 깊이 표준 트랜스포머의 레이어 간 표현보다 서로 더 유사함을 확인한다. 이는 공유 블록이 깊이별로 특화된 변환이 아니라 폭넓게 적용 가능한 변환을 학습했음을 시사한다.
명시적 난이도 신호 없이도 암묵적 연산 라우팅이 일어난다. Prelude 이후 손실 기준으로 토큰을 십분위로 정렬하면, 초기 난이도와 여섯 단계에 걸친 총 개선량 사이에 R 제곱 0.998의 거의 선형인 관계가 나타난다. 가장 쉬운 십분위는 0.49 nats만 개선되는 반면 가장 어려운 십분위는 5.02 nats 개선되어 열 배 차이가 난다. 게이트 값과 난이도의 상관은 모든 단계에서 -0.07에서 -0.04 사이로 약한 음의 값이며, 불확실한 토큰에 대해 게이트가 약간 더 많이 쓰도록 허용하는 것과 일치한다.
가중치 수준의 기전도 보고된다. 게이트 가중치의 비대칭성으로 Wx의 주요 특이값이 Wh보다 44% 크며, 이는 게이트가 앵커보다 은닉 상태 변화에 더 강하게 반응하도록 유도한다. Contrastive projection은 평균 행 코사인 -0.19를 보여 공유 블록이 prelude 이후 무엇이 바뀌었는지를 강조하는 차분 신호를 받도록 한다. 헤드 특화는 가중치로 고정되어 동일한 14개의 previous-token 헤드, 8개의 broadcast 헤드, 11개의 self-attention 헤드가 매 단계에 나타난다. 즉 블록은 매번 같은 기능적 프로그램을 적용하고, 각 단계의 출력을 얼마나 반영할지는 헤드가 아니라 게이트가 결정한다.
추론 시 학습 때 본 R = 6을 넘겨 확장하는 것도 평가했다. 8, 10, 12단계를 적용하면 각각 2.69, 2.72, 2.74로 성능이 소폭 악화된다. 모델이 6단계에서 표현적으로 수렴했고 추가 단계는 정제가 아니라 노이즈를 도입한다는 뜻이다.
한계와 주의사항
논문이 명시한 한계는 세 가지다. 첫째, 재귀 깊이 R이 추론 시 고정되어 있으며 토큰별 halting 메커니즘이 없다. 둘째, 게이트 bias와 노이즈 크기가 GPT-2 계열을 벗어나면 재조정이 필요할 수 있다. 셋째, 스케일에 따라 달라지는 최적 공유 비율에 대한 체계적 연구가 남아 있다.
추가로 유의할 점들이 있다. KV 캐시가 R배로 확장되는 비용은 실재하며, 파라미터 절감 효과를 잠식하는지는 배치 크기에 달려 있다. 배치 크기 32에서 순진한 R배 캐시는 dense 대비 0.91배에 그쳐 메모리 이점이 거의 사라진다.
Medium과 Large의 표준 토큰 예산에서는 dense 베이스라인이 여전히 앞서며, 격차가 좁혀지는 것은 토큰 예산을 늘렸을 때다. isoPARAMS Medium 설정에서는 heavy-tail Poisson이 GRT보다 근소하게 좋은 검증 손실을 기록한다. 다운스트림 평가에서도 isoPARAMS 모델이 BoolQ에서 5.29점 하락하는 등 모든 과제에서 일관되게 개선되지는 않는다.
학습 시간 비교는 컴파일 조건이 달라 dense 열이 유리하게 측정되었으므로 재귀 방법들 사이에서만 비교 가능하다. 또한 추론 시 재귀를 학습 때보다 깊게 확장하면 성능이 오히려 나빠지므로, 학습 시 설정한 R이 사실상 상한으로 작동한다.
모든 실험이 GPT-2 계열 아키텍처, 20,000 스텝, 약 9.8B 토큰 규모에서 수행되었다는 점도 결과 해석 시 고려해야 한다.
결론
GRT는 현재 은닉 상태, 고정 prelude 출력, 주입된 확률적 노이즈를 조건으로 하는 학습된 요소별 게이트가 모델 깊이와 파라미터 수를 분리하는 원리적 전략임을 보였다. isoFLOPS 제약에서 dense 베이스라인 고유 파라미터의 36~37%만 저장하면서 모든 스케일에서 MoR와 heavy-tail Poisson 샘플링을 앞섰다. isoPARAMS 제약에서는 일관된 검증 손실 개선과 최대 +8.29점의 다운스트림 개선을 얻었다.
설계상의 핵심은 공유 변환이 가중치에서는 고정되어 있지만 동작에서는 동적이라는 점이다. 게이트가 매 단계 서로 다른 컨텍스트 기반 입력을 구성하기 때문에, 동일한 파라미터를 재사용하면서도 각 패스에서 다르게 계산할 수 있다. 여기에 학습 중 균일 깊이 샘플링을 결합하면 보조 손실 없이 추론 시 연속적인 연산-품질 트레이드오프가 창발한다.
향후 방향으로는 동적 토큰별 halting과 dense 언어 모델에서 재귀 student로의 지식 증류가 제시되었다.