LoopUS: Pretrained LLM을 Looped Latent Refinement 형태로 재구성하는 포스트-트레이닝 프레임워크

프리트레인된 LLM의 파라미터를 늘리지 않고도 테스트 시간 계산(TTC)을 증가시켜 추론 품질을 향상시킬 수 있다. middle 레이어를 재활용하는 구조적 분해와 게이팅, sparse 감독, 신뢰도 기반 중단으로 긴 루프를 안정적으로 학습할 수 있다. 결과적으로 다양한 백본에서 perplexity를 감소시키고 downstream 정확도를 향상시키며, 적은 학습 자원으로도 적응 효율성을 높인다.

출발점은 pretrained Transformer의 깊이에 따른 representational geology다. 초기 레이어는 로컬 단어 처리에서 시작해 중간 레이어에서 점진적으로 추상화된 예측 표현으로 이동하고, 최종 레이어에서 출력 공간으로 급격히 전이된다. LoopUS는 middle block을 재사용 가능한 latent workspace로 보고, loop body로 삼아 M(h) 업데이트를 연속적으로 적용하되 Gate를 통해 각 업데이트의 폭을 제어한다. 학습은 random depth에서의 샘플링과 monotonicity, confidence 기반 손실로 안정적인 루프를 유도하고, inference 시 adaptive stopping으로 TTC를 효율적으로 관리한다. 결과적으로 루프는 깊이 확장 없이도 “정제적 추론”을 수행하는 제어된 성능 향상을 제공한다.

3단 구성: (1) Recasting: LLM을 Encoder E, Reasoning Block M, Decoder D로 분할하고 middle Block M만 루프 바디로 재사용한다. (2) Selective gating: 각 루프에서 M(h(b))와 h(b) 간의 convex 조합을 통해 h(b+1) = α(b) ⊙ M(h(b)) + (1−α(b)) ⊙ h(b)로 drift를 억제한다. α(b)는 Δδ(b)에서 추출되는 채널별 스케일로 결정되며, A를 곱해 exp를 적용해 (0,1) 사이의 값을 갖는다. (3) Random Deep Supervision: B 단의 루프를 펼치되 S 깊이만 학습하고 나머지는 detach, memory 효율성 확보. (4) Confidence Head 및 Early Stopping: q̃(b+1) = qφ(h(b+1)), q(b+1) = σ(q̃(b+1))로 루프 중지 여부를 판단한다. (5) Training Objective: L(b) = LLM + Lmono + LQ; Lmono는 예측 손실의 비약적 악화를 방지하기 위해 이전 손실 차이를 통해 완만한 감소를 유도한다. (6) KV 캐시 및 인퍼런스: 루프의 각 깊이별 KV 캐시를 독립적으로 유지해.prefix 재계산을 피하고 속도 향상을 얻는다.

벤치마크에서 LoopUS는 프리트레인 백본의 성능을 개선한다. 예를 들어 Qwen3-1.7B의 평균 정확도는 w/o LoopUS 53.7에서 w/ LoopUS 55.3으로 상승(+1.6). Qwen3-4B는 60.3에서 62.1로 상승(+1.8). Qwen3-8B는 63.2에서 65.4로 상승(+2.2). Phi-4 14B는 67.0에서 68.6으로 상승(+1.7). TinyLlama는 평균 성능에서 더 큰 향상을 보이며, Adapted 49.4(+6.3)로 증가한다. WikiText/Lambada perplexity 역시 감소하며, LAMBADA와 ARC 관련 지표에서도 개선이 관찰된다. 또한 KV caching을 도입하면 1024 토큰 생성 시 Qwen3-1.7B에서 1.64x, Qwen3-4B에서 2.31x, Qwen3-8B에서 2.49x의 속도 향상을 보인다. Moreover, ablation 연구는 selective gate 제거 시 drift 억제가 약화되고, encoder–decoder decomposition 제거 시 루프의 안정성이 저하되며, random deep supervision 미적용 시 최적화가 불안정해지는 것을 보인다. Monotonicity 손실은 SiLU가 가장 안정적이며 TBPTT 대비 효율이 높다.

1. Architecture: Encoder E, Reasoning Block M, Decoder D로 구분하고 middle Block M을 루프 바디로 재활용한다. 2) 업데이트 규칙: h(b+1) = α(b) ⊙ M(h(b)) + (1−α(b)) ⊙ h(b)이며 α(b) = exp(∆(b) ⊙ A)로 각 채널에 대해 0~1 사이의 값을 갖는다. 3) 학습 손실: L(b) = LLM + Lmono + LQ; Lmono = SiLU(LLM(b) − LLM(b−1))로 업데이트의 질적 향상을 Soft하게 규제한다. LQ는 confidence head의 이진 cross-entropy로 적절한 중단을 지도한다. 4) Random Deep Supervision: B 루프 중 K 깊이만 supervisor로 학습하고, 나머지는 detach한다. 5) Halting: threshold qth 기반 stopping과 convergence 기반 stopping, CDF 기반 stopping을 비교한다. 6) Inference: KV 캐시를 층별 깊이별로 분리 관리하고, prefill와 decode 두 단계로 구성한다.