150M 파라미터의 비-트랜스포머 언어 모델 'WarpState'를 3억 토큰으로 직접 학습시켜 보았습니다

WarpState란 무엇인가?

WarpState는 표준 트랜스포머 스택이 아니야.

기본 아이디어는 다음 세 가지를 합치는 거지:

1. 로컬 타일드 어텐션 (Local tiled attention)

전체 시퀀스에 걸친 글로벌 셀프 어텐션 대신, 토큰을 고정된 128토큰 단위의 청크로 나눴어.

각 청크 안에서는 일반적인 인과적 스케일드 닷 프로덕트 어텐션(causal scaled-dot-product attention)을 사용해.

모든 청크는 학습 중에 토큰 하나하나 어텐션을 돌리는 게 아니라, 큰 배치 단위의 GPU 작업으로 처리할 수 있어.

그래서 로컬 경로는 대략 이런 식이야:

tokens
↓
128-token chunks
↓
causal local attention
↓
local representation

2. 빠르고 느린 텐서 메모리 (Fast + slow tensor memory)

처리가 끝난 청크들은 지속적인 텐서 메모리로 압축돼.

WarpState는 모든 어텐션 헤드마다 두 개의 행렬을 유지해:

Fast State (빠른 상태)
Slow State (느린 상태)

빠른 상태는 비교적 짧은 메모리 타임스케일로 초기화되고, 느린 상태는 정보를 훨씬 더 오래 유지하도록 초기화돼.

개념적으로 보면:

current chunk
↓
K and U
↓
bounded tensor write
↓
┌───────────────┐
│ Fast memory │
│ Slow memory │
└───────────────┘
↓
future chunks

메모리 쓰기는 제한된 외적(outer-product) 형태의 업데이트를 기반으로 해:

write = tanh(K)^T × tanh(U) / chunk_size

그리고 상태들은 대략 이런 식으로 업데이트돼:

Fast = decay_fast × Fast + (1 - decay_fast) × write

Slow = decay_slow × Slow + (1 - decay_slow) × write

감쇠율(decay rates)은 헤드마다 독립적으로 학습돼.

시작값은 대략 이래:

Fast decay ≈ 0.90
Slow decay ≈ 0.99

모델은 빠른 메모리와 느린 메모리를 각각 얼마나 읽어올지도 학습해.

3. 로컬 어텐션과 메모리 간의 학습된 라우팅 (Learned routing between local attention and memory)

모든 토큰에 대해, 모델은 다음 중 어디에서 정보를 얼마나 가져올지 결정하는 게이트를 생성해:

local chunk attention
vs
long-range tensor memory

대략 이런 식임:

output =
gate × local_attention
+
(1 - gate) × memory_read

그러니까 모델이 이전 청크에서 가져온 정보는 압축된 상태에 의존하면서도, 토큰 간의 정밀한 로컬 관계는 그대로 활용할 수 있는 거지.

Shared recurrent depth

WarpState에서 또 특이한 점은 16개의 거대한 레이어를 전부 따로따로 두지 않았다는 거야.

현재 모델에는 4개의 물리적 WarpState 코어만 들어있는데, 이걸 16번의 논리적 깊이(logical depth) 패스에 걸쳐서 재사용해:

Core 0
Core 1
Core 2
Core 3
Core 0
Core 1
Core 2
Core 3
...

각 논리적 깊이마다 학습된 작은 스케일과 바이어스 값이 있어서, 같은 물리적 코어라도 몇 번째 깊이에서 쓰이느냐에 따라 다르게 작동할 수 있어.

간단히 말하면:

x = x × (1 + depth_scale) + depth_bias

x → shared WarpState core

모든 거대한 가중치 행렬을 복제하지 않고도 더 깊은 반복 연산을 수행하려는 의도지.

자동 회귀 생성 중에, 두 깊이가 같은 물리적 코어 가중치를 공유하더라도 각 논리적 깊이는 각자 독립적인 메모리 캐시를 따로 가져.

Other details

현재 버전 스펙은 이래:

d_model: 1280
heads: 20
head_dim: 64
physical cores: 4
logical depth: 16
FFN hidden: 4480
chunk size: 128
RMSNorm
SwiGLU
RoPE inside each local chunk
tied input/output embeddings

입력 프로젝션은 하나로 합쳐져서 다음을 출력해:

Q
K
V
local/memory gate
memory U

이걸 한 번의 프로젝션으로 다 뽑아내는 거지.

Training results

내가 제일 궁금했던 건 이 아키텍처가 실제 사전 학습(pretraining)을 버틸 수 있느냐 하는 거였어.

결과는 성공.

NaN이나 그래디언트 붕괴, 혹은 눈에 띄는 최적화 실패 없이 약 3억 개의 토큰으로 전체 학습을 끝냈거든.

학습 막바지에도 그래디언트 노름(gradient norm)은 대략 이 정도 수준이었어:

0.65 – 0.75

그때 학습률(learning rate)은 이미 이 정도로 줄어든 상태였고:

3e-5

최대 VRAM 점유율은 4.52 GB 정도를 유지했어.

모델이 학습 과정에서 언어 구조를 확실히 익혔다는 것도 확인했어.

초기 체크포인트는 그냥 영어처럼 보이는 노이즈만 뱉어냈거든.

근데 나중 체크포인트로 갈수록 의미론적으로 그럴듯한 클러스터가 형성되고, 문단 구조도 꽤 갖춰지더라고.

예를 들어, 페이스북에 대해 물어보면 최종 모델은 이런 것들과 연관 지어 대답해:

online platform
social media
sharing content
sharing information
interaction with other people
community

물론 아직 제대로 된 챗봇이라고 하기엔 무리가 있지.

분명한 한계점들도 보여:

반복 루프
의미론적 끌개(semantic attractors) 현상
약한 사실 관계 기억력
가끔 발생하는 역할 혼동
긴 문장 생성 시 붕괴 현상

게다가 이 모델은 그냥 베이스 사전 학습만 거친 상태야.

인스트럭션 튜닝, SFT, RLHF 같은 건 하나도 안 했으니까, 내가 첨부한 채팅 스크린샷은 챗봇 벤치마크라기보다는 그냥 성능 확인용 테스트 정도로 봐줘.

또 다른 중요한 한계는 학습 예산이야.

1억 5천만 파라미터 모델을 겨우 3억 토큰으로 학습시킨 거라, 파라미터당 학습 토큰 수가 대략:

~2 tokens per parameter

이 정도밖에 안 돼. 그래서 이건 제대로 학습된 150M 언어 모델이라기보다는, 이 아키텍처가 학습 가능하다는 걸 증명한 결과물로 봐주면 좋겠어.

What surprised me most

나한테 흥미로웠던 건, 이런 조건에서도 모델이 의미 있는 언어 표현을 학습할 수 있었다는 점이야:

• 물리적 코어가 4개뿐인데도,

• 그 코어를 계속 재사용하고,

• 어텐션을 128 토큰의 로컬 윈도우로 제한하고,

• 청크 간의 정보 이동을 고정 크기 텐서 상태로 처리했는데도 말이지.

그래서 장기 메모리 크기가 기존의 전체 KV 캐시처럼 컨텍스트에 비례해서 선형적으로 늘어나지 않아.

확실한 결론을 내리기 전에 테스트해보고 싶은 게 아직 많아.

다음 단계는 아마 이런 것들이 될 거야:

• 전체 검증 데이터셋에 대한 결정론적 평가;

• 똑같은 데이터로 학습한 파라미터 동일 수준의 트랜스포머 베이스라인 비교;

• 빠른/느린 메모리 상태 분석;

• 긴 컨텍스트에서의 동작 측정;

• 반복/끌개 문제 조사;

• 나중에는 더 큰 학습 예산으로 테스트해보기.

일단은 첫 번째 전체 실행 결과를 공유하고 싶었어. 왜냐하면 이 시점부터 아키텍처가 단순한 아이디어를 넘어 실제로 학습된 언어 모델이 되었거든.

아키텍처에 대한 피드백은 언제든 환영이야. 특히 메모리 업데이트나 공유 코어(shared-core) 설계에 대한 비판은 더더욱 환영함.

고마워, 정말 유용한 피드백이야.

청크 재귀에 관해서 말하자면, 현재 병렬 스캔은 없어. 128토큰 청크 전체에 대한 로컬 어텐션은 하나의 SDPA 작업으로 배치 처리되지만, fast/slow 메모리 업데이트는 청크별로 순차적으로 진행돼. 1024토큰 시퀀스라면 논리적 패스당 8번의 상태 업데이트가 일어나는 셈이지. 현재 학습 순전파(forward)에서는 각 논리적 깊이 패스마다 상태가 초기화되고, 그 패스 내의 청크들을 가로질러 전달돼. 생성(generation)은 조금 다른데, 왜냐하면 각...