TL;DR
루프형 언어 모델은 반복 반복으로 추론을 수행하기 때문에 KV-cache의 메모리 사용량이 깊이에 따라 선형 증가한다. MELT는 레이어별 단일 KV-cache를 공유하고, gated latent state를 통해 과거 루프의 정보를 누적·교체하며 메모리 증가를 억제한다. 이를 통해 깊이가 증가해도 메모리 footprint는 상수에 가깝게 유지되면서도 LoopLM 수준의 성능을 유지한다. chunk-wise training과 two-phase 학습(Interpolated transition, Attention-aligned distillation)을 통해 MELT를 Ouro에서 파인튜닝하더라도 안정적으로 수렴하고, 메모리-효율성과 추론 성능의 균형을 달성한다.
왜 중요한가
루프형 언어 모델은 반복 반복으로 추론을 수행하기 때문에 KV-cache의 메모리 사용량이 깊이에 따라 선형 증가한다. MELT는 레이어별 단일 KV-cache를 공유하고, gated latent state를 통해 과거 루프의 정보를 누적·교체하며 메모리 증가를 억제한다. 이를 통해 깊이가 증가해도 메모리 footprint는 상수에 가깝게 유지되면서도 LoopLM 수준의 성능을 유지한다. chunk-wise training과 two-phase 학습(Interpolated transition, Attention-aligned distillation)을 통해 MELT를 Ouro에서 파인튜닝하더라도 안정적으로 수렴하고, 메모리-효율성과 추론 성능의 균형을 달성한다.
핵심 기여
Memory-efficient looped architecture
각 레이어에 대해 고정 크기의 KV-cache를 공유하도록 설계하고, 루프 간에 업데이트되는 latent state를 통해 메모리 복잡도를 O(L)로 낮춘다. 이는 기존 LoopLM의 KV-cache 확장(O(L×T))과 달리 추론 깊이가 늘어나도 메모리 증가를 억제한다.
Learnable gating for memory management
element-wise gating(z_t)으로 latent state h_t^l를 업데이트하고, 이를 통해 과거 루프의 정보를 선택적으로 보존한다. z_t가 1에 수렴하면 그래디언트 흐름이 안정되어 Gradient Superhighway가 형성된다.
Two-phase training with distillation
Phase 1: chunk-wise training과 interpolated transition으로 LoopLM에서 MELT로 부드러운 전이를 수행하고, 다중 루프에 걸친 지식을 교사 없이 가깝게 학습한다. Phase 2: frozen LoopLM을 교사로 하는 attention-aligned distillation을 도입해 MELT 표현을 LoopLM과 정렬한다.
Empirical efficiency and performance
MELT-1.6B는 Ouro-1.4B-Thinking 대비 메모리 footprint를 크게 줄이면서도 유사한 성능을 보여주고, 일반 Transformer 대비 메모리 효율이 뛰어나다. KV-cache per token과 32k 토큰 생성을 포함한 엔드-투-엔드 메모리 측정에서 MELT은 약 3–4× 메모리 절감을 달성한다.
Stability and ablations
게이트 구조의 ablation에서 element-wise gating이 가장 좋은 성능을 보였고, chunk-wise training 없이는 성능 저하가 크게 발생한다. 두 학습 단계의 공조로 MELT의 표현 drift를 억제한다.
핵심 아이디어 이해하기
출발점: LoopLM 등 루프형 Transformer는 반복 루프마다 KV-cache를 확장하므로 깊이가 깊어질수록 메모리 증가가 심각하다. 해결 원리: MELT는 latent state h_t^l를 도입하고, k_t^l, v_t^l를 h_t^l에서 projection으로 얻어 각 토큰의 KV를 구성한다. 이 latent state는 gating으로 업데이트되며, KV-cache는 각 토큰별로 한 줄씩 누적되되 루프 간 확장은 없다. thus, memory grows 선형이 아닌 고정에 가깝다.
방법론
수학적 흐름: x_t^(l) → z_t^(l) = σ(W_z x_t^(l) + U_z h_{t-1}^(l) + b_z) → h_t^(l) = z_t^(l) ∘ h_{t-1}^(l) + (1 - z_t^(l)) ∘ x_t^(l).
관련 Figure
MELT의 inference 시나리오에서 latent state가 KV 캐시에 어떻게 기여하는지 시각적으로 제시한다. 이 다이어그램은 MELT의 핵심 아이디어인 latent-state로부터 KV를 생성하고, 각 토큰에서 레이어를 재사용하는 방식의 메모리 관리와 추론 흐름의 결합을 보여준다.
MELT의 추론(inference) 흐름 다이어그램(캐시 업데이트 및 latent state 흐름 포함)
Chunk-wise training과 interpolated transition의 작동 방식을 시각화하며, KV base와 MELT KV를 혼합하는 방식이 학습 중 어떤 방식으로 진행되는지 보여준다. MELT의 학습이 LoopLM에서 MELT로 연속적으로 전이되는 과정을 이해하는 데 직접적으로 기여한다.
Phase 1 학습(Chunk-wise + Interpolated Transition) 구조 및 KV 캐시 흐름 다이어그램
주요 결과
주요 벤치마크에서 MELT-1.6B는 비-looped 모델보다 우수하거나 유사한 성능을 보이며, 메모리 footprint는 Ouro 대비 대략 3–4× 감소한다. KV-cache per token: MELT-1.6B = 0.196608 MB/token, Ouro-1.4B-Thinking = 0.786432 MB/token. 32k 토큰 세대의 KV-cache 총합: MELT = 3.272 GB, Ouro = 25.17 GB; 전체 메모리 합계: MELT = 9.49 GB, Ouro = 27.97 GB.
관련 Figure
MELT-1.6B가 비슷한 메모리 footprint에서 비-looped 모델들을 능가하거나 근접한 성능을 보임을 시각적으로 보여준다. 동일한 32k 토큰 세대에서 MELT의 KV-cache 총합은 약 3.3 GB로 Ouro의 약 25 GB에 비해 크게 낮다. 이는 MELT가 메모리 효율을 유지하면서도 추론 성능을 확보한다는 핵심 메시지와 직결된다.
AIME26 벤치마크에서 MELT-1.6B와 Ouro-1.4B의 메모리 대 성능 비교 차트
기술 상세
아키텍처 개요: MELT는 각 레이어에 대해 고정 크기의 latent state h_t^l를 도입하고, 이를 통해 K^l_t, V^l_t를 h_t^l에서 선형 변환으로 얻는다. KV-cache는 기존처럼 매 토큰마다 확장하지 않고, 새로운 KV 엔트리를 추가하는 대신 토큰 간 게이트 업데이트를 통해 누적된 정보를 반영한다. 업데이트 공식: z_t^(l) = σ(x_t^(l) W_z + h_{t-1}^(l) U_z + b_z), h_t^(l) = z_t^(l) ∘ h_{t-1}^(l) + (1 - z_t^(l)) ∘ x_t^(l). 이후 k_t^(l) = h_t^(l) W_K^(l), v_t^(l) = h_t^(l) W_V^(l). 이로써 M^(l)MELT ∝ O(L)로, 깊이에 비례한 KV 증가를 제거한다. 탐색적 파라미터 업데이트는 gating 파라미터와 W_K, W_V를 학습한다. 학습 전략: Phase 1은 chunk-wise + interpolated transition으로 LoopLM에서 MELT로의 부드러운 전이를 시도하고, 두 KV 세트를 α 비율로 혼합( KV = α KVMELT + (1-α) KVbase)하여 inference 동작을 근사한다. Phase 2는 attention-aligned distillation을 도입해 LoopLM과 MELT의 layer/loop 간 표현 정렬을 강제한다. 손실 함수 조합: L = L_KD + β NT^{-1} ∑{l,t} || o^(l,t)(MELT) − sg(o^(l,t)(LoopLM)) ||^2. 구현 세부: Ouro-1.4B-Thinking으로 초기화하되 gating 파라미터는 무작위 초기화. 데이터로는 AceReason-1.1-SFT + OpenThoughts3를 사용. 총 학습 시간은 130시간(8x H100 80GB) + 추가 60시간의 ablation, 총 약 1,440 GPU-hours.
관련 Figure
Phase 2의 attention-aligned distillation이 MELT의 post-attention 표현을 frozen LoopLM의 대응 표현과 계층/루프별로 정렬하도록 정규화하는 방식을 시각적으로 보여준다. 이 손실은 MELT의 표현 drift를 줄이고 전체 성능 차이를 더 좁히는 데 중요한 역할을 한다.
Attention-Aligned Distillation의 보조 정렬 손실(사진 내 그림으로 나타낸 loss 경로)
한계점
추론 시 고정된 루프 깊이를 사용하므로 적응적 루프 깊이 적용은 아직 탐구 중이다. 또한 현재 구현은 GQA와 같은 KV-캐시 공유 기법의 탐색을 추가로 수행하지 않았다. 트레이닝의 sequential KV 업데이트는 더 큰 모델로 확장 시 병렬성을 제한한다.
실무 활용
MELT는 메모리 한계가 있는 환경에서도 루프형 추론을 더 깊게 수행할 수 있게 해주는 아키텍처다. Chunk-wise training과 두 단계 학습으로 안정적으로 파인튜닝 가능하며, Ouro 기반 파인튜닝의 메모리 제약 없이 유사한 성능을 낼 수 있다.
- 긴 컨텍스트를 요구하는 수학/코딩 추론에서 메모리 제약 하의 루프형 추론 적용
- 제한된 GPU 메모리 환경에서의 루프형 LLM 파인튜닝 및 배포
- 메모리-효율성과 추론 성능 간의 균형이 중요한 대화형 에이전트의 심층 추론
- GQA/다중 단계 추론 벤치마크에서의 효율적 루프 구조 연구
코드 공개 여부: 미확인
키워드
AI 요약 · 북마크 · 개인 피드 설정 — 무료
출처 · 인용 안내
인용 시 "요약 출처: AI Trends (aitrends.kr)"를 표기하고, 사실 확인은 원문 보기 기준으로 진행해 주세요. 자세한 기준은 운영 정책을 참고해 주세요.