재귀 언어모델의 고정점 최적화로 학습·추론 비용 2배 절감
원제 Towards Looped Models Done Right, Part II: Rethinking at Fixed Points
추천 20댓글 2
Key Point
재귀 언어모델의 학습·추론 비용을 획기적으로 줄이는 구체적 기법(KV 캐시 축소, 역전파 생략, RL 속도 개선)과 이론적 기반(고정점 특성)을 제시하며, 실험 결과로 입증된 개선 방법을 다룬다.
핵심 요약
- 루프된 언어모델의 매 반복은 학습·디코딩·프리필·강화학습에서 비용을 증가시키는데, 재귀 상태가 고정점에 가까워질수록 그 경로의 중요도는 감소한다.
- 고정점 근처에서의 수렴 성질을 활용하면 학습 중 절단된 역전파(truncated backpropagation), 거의 정확도 손실 없는 터미널 KV 공유 디코딩, 1.79배 빠른 프리필을 위한 증류 학생 모델, 저장된 롤아웃 상태에서 그래디언트를 계산하는 RL 업데이트(재생 경로 역전파보다 2배 빠름)가 가능하다.
- 고정점을 형성하는 두 핵심 요소인 깊이 사전(depth prior)과 입력 주입(input injection)을 개선했다.
- 고정 깊이 학습은 KV 공유를 방해하고, Huginn의 넓은 깊이 사전은 공유를 지원하지만 목표 깊이에서의 감독 신호를 공유보다 더 희석시킨다.
- 새로운 접근은 예측 피드백에서 깊이 사전을 학습하되, 엔트로피 항으로 사전을 넓게 유지한다.
- 기존 주입 방식은 상태의 입력 방향 성분이 주입을 증폭하거나 상쇄할 수 있으므로, 직교 주입(orthogonal injection)으로 이 성분을 제거한다.
- 100M~1.6B 파라미터 범위에서 학습된 사전과 직교 주입이 기존 Huginn 사전 및 주입 방식 대비 모든 규모에서 낮은 퍼플렉시티를 달성한다.
- 1.6B 규모에서 학습된 사전과 3배 작은 KV 캐시 조합이 전체 캐시를 사용한 고정 깊이 학습의 다운스트림 평균 성능과 맞춘다.