LaCT(Large Chunk Test-Time Training)는 테스트타임 학습(TTT, Test-Time Training)의 빠른 가중치(fast weights)를 작은 미니배치 대신 매우 큰 청크 단위로 갱신하는 장문 컨텍스트 모델링 방법이다. 논문 “Test-Time Training Done Right”는 2K~100만 토큰 청크, 최대 모델 파라미터 40% 규모의 비선형 상태, 14B 파라미터 자기회귀 비디오 확산 모델 실험을 통해 TTT를 현대 GPU에 맞게 다시 설계할 수 있음을 보인다.
배경: 기존 TTT의 병목
테스트타임 학습은 추론 중 모델 일부 가중치를 임시로 업데이트해 현재 시퀀스의 기억을 저장한다. 개념적으로는 RNN의 hidden state와 비슷하지만, state가 단순 벡터가 아니라 작은 학습 문제로 업데이트되는 빠른 가중치라는 점이 다르다.
문제는 기존 TTT 계열이 보통 16~64 토큰 같은 작은 온라인 미니배치로 가중치를 갱신했다는 데 있다. 이렇게 하면 시퀀스 의존성은 촘촘하게 반영할 수 있지만, GPU에서는 병렬성이 낮아 FLOPs 활용률이 5% 미만으로 떨어질 수 있다. 이미지 세트, 비디오, N차원 그리드처럼 1차원 순서만으로 설명하기 어려운 데이터에도 부적합하다.
LaCT의 핵심 아이디어
LaCT는 정반대 방향을 택한다. 빠른 가중치를 매 16~64 토큰마다 조금씩 갱신하는 대신, 2,048 토큰부터 100만 토큰까지의 큰 청크를 한 번에 처리한다.
| 설계 요소 | 설명 |
|---|---|
| 대형 청크 업데이트 | 청크 내부 토큰을 큰 배치처럼 다뤄 GPU 병렬성을 높인다 |
| 비선형 빠른 가중치 | 상태 크기를 모델 파라미터의 최대 40% 수준까지 키워 기억 용량을 확장한다 |
| Window attention | 큰 청크 내부에서 지역 의존성을 보강한다 |
| Muon 최적화 | 테스트타임 온라인 업데이트에 더 복잡한 옵티마이저를 붙일 수 있다 |
| 순수 PyTorch 구현 | 커스텀 커널 의존을 줄여 연구 반복 속도를 높인다 |
논문은 LaCT가 NVIDIA A100에서 최대 70% GPU 활용률을 달성할 수 있다고 보고한다. 핵심은 “긴 컨텍스트를 더 자주 업데이트하자”가 아니라, “더 큰 단위로 업데이트해 하드웨어가 잘 먹는 행렬 연산으로 바꾸자”에 가깝다.
검증한 태스크
Novel View Synthesis
이미지 세트를 컨텍스트로 받아 새로운 시점을 합성하는 태스크다. 논문은 가장 긴 실험에서 100만 토큰 이상의 컨텍스트 길이를 사용한다. 이 설정은 대형 청크가 1차원 텍스트뿐 아니라 이미지 집합 같은 구조화된 입력에도 맞는다는 점을 보여준다.
Language Modeling
언어 모델에서는 마지막 2K 토큰 검증 손실 등을 통해 상태 크기와 청크 크기 변화가 성능에 미치는 영향을 분석한다. LaCT는 큰 비선형 상태를 효율적으로 다룰 수 있어 장문 컨텍스트에서 단순 어텐션 확장과 다른 연구 경로를 제시한다.
Autoregressive Video Diffusion
논문은 14B 파라미터 양방향 비디오 확산 트랜스포머를 LaCT와 sliding window attention으로 자기회귀 모델처럼 바꿔 최대 56K 비주얼 토큰 길이의 일관된 비디오 생성을 실험한다.
왜 중요한가
긴 컨텍스트 모델링은 kv-caching 압축, sliding window, sparse attention, Mamba 계열 상태공간 모델처럼 여러 방향으로 발전하고 있다. LaCT는 이 흐름에서 “추론 중 적응되는 큰 메모리”라는 별도 축을 제안한다.
특히 실무적으로 중요한 점은 다음이다.
- 커스텀 CUDA 커널 없이 PyTorch 코드로 실험 가능한 구조를 지향한다.
- 텍스트뿐 아니라 이미지 세트와 비디오 같은 N차원 데이터까지 같은 프레임으로 다룬다.
- 장문 컨텍스트 비용을 단순히 KV 캐시 최적화 문제로만 보지 않고, 추론 중 학습되는 상태 용량 문제로 재정의한다.
한계
LaCT는 아직 연구 단계의 아키텍처다. 공개 제품이나 범용 LLM 런타임에서 바로 사용할 수 있는 최적화 기법이라기보다, TTT 계열을 현대 GPU와 멀티모달 장문 데이터에 맞게 다시 설계하는 방향성에 가깝다. 또한 큰 청크는 청크 내부 순서 정보를 약하게 만들 수 있어 window attention 같은 보완 설계가 필요하다.
관련 문서
- llm-architecture-tips-long-context-cost — 장문 컨텍스트 비용을 줄이는 최신 LLM 아키텍처 흐름
- kv-caching — LLM 추론 속도를 높이는 KV 캐시 원리
- turboquant — KV 캐시·임베딩 압축 양자화
- delta-mem — LLM 어텐션에 결합되는 초소형 온라인 메모리
- sana-wm — 장시간 비디오 생성을 다루는 NVIDIA 오픈 세계 모델
참고 자료
- Test-Time Training Done Right — arXiv (2025-05-29)
- Project page — Tianyuan Zhang