Engram 모델의 폭주! 2b 모델 업데이트...
Engram gone wild! 2b model update...
핵심 요약
Engram 테이블을 활용해 모델의 지식 저장 효율을 높이는 2.6b 모델의 훈련 진행 상황과 테스트 결과를 공유합니다.
- Engram 구조 — 4.3b 크기의 Engram 테이블을 도입하여 모델의 지식 저장 부담을 외부로 분산함
- 학습 효율화 — SVD로 압축된 OLMo 3 기반의 고정된 LM Head와 Embed를 사용하여 훈련 시간 단축
- 성능 테스트 — 1억 토큰 학습 후 테스트 결과, 반복 학습된 데이터가 Engram으로 효과적으로 오프로드됨을 확인
- 스폰서 구인 — 모델 훈련 완성을 위해 컴퓨팅 자원을 지원해 줄 스폰서를 찾는 중
업데이트: https://old.reddit.com/r/LocalLLaMA/comments/1wis23s/update_small_model_engram/
대충 일주일 정도 지났으니 다시 돌아왔다. 사람들이 모델에 대해 이것저것 물어보더라고.
코드나 모델 같은 걸 원하던데, 지금 당장 그런 거 줘봤자 제대로 학습도 안 된 모델이라 쓸모도 없을 거다. 그러니까 바로 본론으로 들어가자.
수많은 테스트 끝에 확정된 스펙은 다음과 같다:
전체 2.6b 모델. 임베딩, LM 헤드, AttnRes 등 포함.
2.2b는 학습 완료. 임베딩/LM 헤드는 고정됨(각각 약 205m).
4.3b ENGRAM 테이블. 그래, 덩치가 좀 크다.
이제 아키텍처 쪽 얘기 좀 해보자:
표준 어블레이션(ablation) 연구에서 "최적"이라고 밝혀진 3:1 비율의 SWA 레이어를 포함했다. 4k / 4k / 8k SWA 레이어 뒤에 글로벌 레이어가 붙는 방식이다.
첫 번째 "블록"(4개 레이어)이 끝나면 Engram 테이블이 등장한다. Engram 테이블 자체는 아주 적은 수의 어텐션 헤드로 돌아가기 때문에 데이터를 무작정 때려 박는 방식이 아니다. 이게 바로 컨텍스트를 인식하는 Engram이다. Qwen이나 DS(DeepSeek)가 쓰는 방식이랑 정확히 일치하는지는 모르겠지만, 내가 쓰는 거랑 꽤 비슷하다. 차이점은 당연히 Engram 사이즈겠지.
이건 10개 블록에 마지막 글로벌 레이어 하나가 더해진 구조다(총 41개 레이어). + 임베드 + LM 헤드.
자, 한 발 물러서서 보면 최적화 문제는 이거다:
백본의 연산 능력을 어떻게 극대화하고, 지루한 작업은 테이블로 떠넘길 것인가?
여기서 Attention Based Residuals(Moonshot)가 구원투수로 등장한다. 왜 AttnRes냐고? 첫 번째 블록 이후의 모든 블록이 Engram 테이블을 어떤 식으로든 활용할 수 있게 해주거든. 잔차 스트림(residual stream)을 통해 어텐션을 거치면서 Engram에서 데이터를 가져올지, 가져온다면 정확히 얼마나 가져올지 결정하는 거지.
왜 Engram 테이블을 더 크게 안 만드냐고? 솔직히 말하면, 더 크게 만들어도 상관은 없을 거다. 하지만 연산과 테이블의 이런 불균형을 시도한 모델이 있었는지 모르겠네. 이론적으로 DeepSeek 연구에 따르면 잘 작동한다고 한다.
더 깊게 쌓을 수도 있겠지. 근데 학습 비용이 너무 비싸다.
왜 LM/Embed를 고정했냐고? OLMo 3 모델에서 SVD(특이값 분해)를 통해 다운 프로젝션한 거거든. 그러니까 ~5k에서 ~2k로 줄이려는 수학적 압축 시도인 셈이지. 덕분에 시간을 엄청나게 아꼈다.
현재 데이터는 ~KD 포맷으로 들어가 있다. 위키피디아 코퍼스에서 교사 모델(teacher)로부터 추출한 32개 로짓을 저장해 둔 상태다. 원-핫(1-hot) 방식으로 학습시키는 대신, 32개의 소프트 타겟을 맞추도록 하는 거지. 하드 크로스 엔트로피는 모델을 너무 쥐어짜서 "수조 개의 토큰"이 필요하다는 소리가 나오는 거다.
LM 헤드, 임베드, 토크나이저, 교사 모델을 빌려오면... 학습량이 훨씬 줄어든다.
---
학습은 어디까지 왔냐고? 어제 태평양 표준시 기준 새벽 4시쯤에 1억 토큰을 넘겼다.
현재 HF 저장소에 체크포인트, 데이터, 그리고 1억 4백만 토큰 지점의 safetensor가 다 올라와 있다.
--
이 시점에서 모델을 좀 검사해보고 싶었다. 완전히 쓰레기인지 아닌지는 확인해야 할 거 아냐? 필요한 데이터의 극히 일부만 본 상태니까.
그래서 딱 1억 토큰만 본 상태에서 문장 완성 테스트랑 여러 어블레이션을 시도해 봤다. Engram이 정확히 뭘 저장하는지(대부분 추측이지만) 연구해 본 거지.
