LIFT가 사전학습에 잠재 상태 피드백을 넣어 추론 성능을 높인다

무엇인가

Transformer 언어모델은 순전파 구조라 깊은 층의 표현이 얕은 층으로 되돌아가지 않는다. 생성 단계 사이에서 정보가 아래로 흐르는 유일한 통로는 디코딩된 토큰 하나뿐이고, 어휘 10만 기준으로 최대 약 17비트밖에 실을 수 없다. 저자들은 이 좁은 통로 때문에 모델이 이미 계산한 중간 결과를 다시 계산하거나, 샘플링되지 않은 대안 연속들을 버려야 한다고 지적한다. 또 고정 깊이 Transformer는 시퀀스 길이에 따라 계산 단계가 늘어나는 상태 추적 과제를 풀 수 없다는 이론적 결과(S5 순열 합성은 NC1-완전이며 층 수가 Ω(log N)만큼 필요)를 근거로 든다. 기존 연구는 사후학습에서 잠재 표현을 되먹이려 했지만 순차적 롤아웃 비용이 컸고, 사전학습 단계에서 Jacobi 반복으로 근사하는 방법은 반복 횟수와 학습 비용 사이의 트레이드오프를 안고 있었다.

어떻게 동작하나

LIFT는 다음 토큰과 함께 다음 상태를 예측하도록 모델을 확장한다. 상태는 출력 로짓을 top-k(k=1024)로 자르고 온도 τ=1.5의 소프트맥스로 재정규화한 분포로, 최대 k개의 0이 아닌 항을 갖는다. 토큰 하나보다 훨씬 많은 정보를 담으면서 대안 연속들의 순위를 유지한다. 상태는 토큰 임베딩 행렬 E를 통해 투영되어 상태가 유도하는 임베딩의 가중평균이 되고, 토큰 임베딩과 연결된 뒤 RMSNorm과 SwiGLU 블록, 잔차 연결로 융합되어 표준 Transformer 스택에 들어간다. 첫 위치에서는 학습된 벡터가 대신 들어간다. 이 융합 층이 LIFT가 추가하는 모든 파라미터(11d² 가중치와 초기 상태 벡터)로, 1B 모델의 3%에 해당하며 깊이가 늘수록 비중이 줄어든다. 저자들은 은닉 상태 대신 분포를 상태로 쓴 이유로, 토크나이저를 공유하는 두 모델 사이에서 분포는 같은 좌표계(토큰 확률)를 갖지만 은닉 상태는 공유 좌표가 없어 회귀가 더 어렵다는 점을 든다.

무엇과 다른가

학습의 핵심은 상태를 교사 강제(teacher forcing)로 미리 계산해 병렬성을 지키는 것이다. 동결된 외부 사전학습 LM(학생과 같은 크기, Chinchilla 예산의 10배로 학습)이 같은 텍스트를 병렬로 처리해 로짓을 내고, 이를 같은 top-k 변환으로 상태로 바꾼다. (x_{i+1}, s*_{i+1}) 쌍이 융합 층에 병렬로 들어가고 교사로는 그래디언트가 흐르지 않는다. 목적함수는 언어모델링 손실에 top-k 분포용으로 변형한 전방 KL 발산 항을 λ=1.5로 더한 형태다. 교사 상태로 학습하고 추론에서는 자기 상태를 쓰기 때문에 분포 불일치가 누적될 수 있는데, 이를 상태 정렬 항과 마지막 10% 구간의 자기 상태 적응(교사와 정렬 항을 제거하고, 50% 확률로 1~2회 초기 순전파를 돌려 자기 상태를 만들어 사용하며 그 초기 반복으로는 그래디언트를 보내지 않음)으로 완화한다. 병렬 프리필을 위해 학습 중 각 시퀀스의 상태를 확률 p로 학습된 편향 b로 대체하는 prefix state dropout을 쓴다. 추론 시에는 모델이 예측한 자기 상태를 되먹이며, 이 오버헤드는 작고 모델이 커질수록 줄어든다.

어떻게 쓰나

상태 추적 능력은 S5 순열 합성 과제로 검증했다. 2층·2M 파라미터 모델을 처음부터 학습시키고 LIFT는 융합 층(파라미터의 7%)을 더했다. 학습 길이는 N∈[1,16], 12개의 고정 생성기와 무작위 초기 상태를 사용했고, 길이별 2,000개 홀드아웃 시퀀스에 대해 N=64까지 정확도를 3개 시드로 평균했다. 고정 깊이 Transformer는 이론 예측대로 무너져 8배 긴 학습(40,000스텝)으로도 N=6까지만 풀고 N=12부터는 우연 수준이었다. 순차적으로 잠재 피드백을 학습한 LIFT-RNN은 학습 범위의 모든 길이를 풀었고, 이 모델의 상태로 병렬 학습한 LIFT는 학습 범위를 넘어 N=32에서 96%, N=64에서 65%까지 외삽했다. 흥미롭게도 깊이가 제한된 비순환 교사(8배 더 학습한 동일 크기 Transformer)로 학습해도 LIFT는 N=12까지 전부, N=16의 절반을 풀었는데, 그 교사 자체는 N=8을 넘으면 우연 수준(N=12에서 3%)이었다. 다만 이렇게 학습한 LIFT는 무한 깊이를 얻지는 못하고 학습 범위를 넘으면 실패한다.

전제와 한계

사전학습 실험은 OLMo 2 아키텍처·코드베이스·데이터로 135M, 350M, 1B 파라미터 모델을 각각 Chinchilla 예산의 5배(1B 규모에서 최대 100B 토큰)만큼 학습했다. 135M은 시드 3개, 더 큰 모델은 연산 제약으로 시드 1개다. 교사는 각 규모에서 같은 크기로 Chinchilla 예산 10배를 학습한 모델이며, 더 강한 교사는 결과를 거의 바꾸지 않았다. 토큰을 맞춘 비교에서 LIFT는 모든 규모의 모든 지표에서 앞서 퍼플렉서티가 5~5.5% 낮아졌고, 36개 다운스트림 벤치마크 비교 중 33개에서 이겼다. 연산을 맞춘 비교에서는 기준 Transformer가 40% 넘게 더 많은 데이터를 학습하는데도 350M과 1B에서 LIFT가 모든 지표 우위를 지켰고 135M에서는 동등하거나 앞섰다. 같은 가중치에서 채널만 끈 LIFT w/o states와, 같은 상태 정렬 목적함수만 쓰고 채널은 없는 증류 기준선은 모든 규모에서 LIFT에 뒤졌다. 채널을 좁히거나 상태 정렬 항을 빼거나 교사를 아예 없앤 절제 실험에서도 이득이 크게 줄었다. 학습이 진행될수록 LIFT의 토큰 효율 이점은 커져서, 기준 Transformer가 같은 손실에 도달하는 데 필요한 추가 토큰 수가 계속 늘어났다.

가장 큰 이득은 절차적 과제에서 나왔다. 한 자리 수 2~10개를 더하고 빼는 자체 산술 과제(사고 사슬 없이 잠재 상태로 다단계 계산)에서 LIFT는 모든 규모에서 기준선을 앞섰고, 2배 많은 토큰으로 학습한 자기 교사보다도 정답 비트가 평균 1~5% 적었다. 규칙을 따르는 짧은 시퀀스의 다음 원소를 맞히는 패턴 완성 과제에서는 토큰을 맞춘 Transformer를 모든 규모에서 BPB 기준 6~10% 이겼고, 2배 토큰 Transformer는 350M과 1B에서 4~5% 이겼다(135M은 동률). Jacobi 반복으로 상태를 근사하는 피드백 Transformer(T2MLR, Multi-pass Transformer)와의 비교에서는 SmolLM2-135M 백본과 10B 토큰 FineWeb-Edu 샘플을 공유했는데, 토큰을 맞추면 LIFT가 모든 지표에서 평균 최고였고 생성 과제에서는 T2MLR과 동률이었지만 이 설정에서는 모든 방법이 노이즈 범위 안이었다. 연산을 맞추면 LIFT는 기준 Transformer와 동등하거나 앞서면서 잠재 피드백 기준선들을 앞섰고, 그 기준선들은 연산이 맞춰지는 순간 산술을 제외한 모든 지표에서 기준 Transformer에 뒤졌다. LIFT는 측정한 배치 크기에서 두 피드백 기준선보다 피크 메모리도 적게 썼다.

개발자 관점에서 이 논문이 주는 실무적 시사점은 세 가지다. 첫째, 되먹임 채널을 사후학습이 아니라 사전학습에 넣고, 그 입력을 외부 교사 모델의 출력 분포로 미리 계산해 두면 순차 롤아웃 없이 기존과 같은 병렬 학습 파이프라인을 유지할 수 있다. 둘째, 추가 비용이 융합 층 파라미터(1B에서 3%)와 추론 시 상태 예측 오버헤드로 제한되고, 그 비중이 모델이 커질수록 줄어든다. 셋째, 프롬프트를 병렬로 처리하는 프리필 경로를 유지하려면 prefix state dropout 같은 설계가 필요하며, 논문은 프롬프트를 반복 정제하는 방식이 순차 처리 이득의 상당 부분을 몇 번의 반복 안에 회복한다고 보고한다. 다만 교사 모델의 토크나이저가 학생과 같아야 상태 공간이 일치하므로, 교사를 고를 때 어휘 호환성을 먼저 확인해야 한다.

저자들이 밝힌 한계는 분명하다. 이 연구는 학생과 같은 아키텍처를 공유하는 교사를 쓰는 선택에 집중했고, 다른 크기나 다른 어휘를 가진 교사의 효과는 향후 과제로 남긴다. 또한 LIFT를 교사의 학습 예산 훨씬 너머까지 학습시키는 영역은 성능과 학습 동학이 아직 열린 질문이다. 깊이가 제한된 비순환 교사로 병렬 학습한 LIFT는 무한한 계산 깊이를 얻지 못하고 학습 길이 범위를 넘으면 실패한다는 점도 명시했다.