FLaRe는 잠재 추론의 다섯 조건을 흐름 매칭으로 충족한다

무엇인가

대형 언어모델이 답을 내기 전에 길게 추론하는 방식은 대개 명시적 사고사슬(CoT)을 토큰 단위로 자기회귀 생성하는 것이다. 논문은 이 방식이 세 가지 한계를 갖는다고 본다. 비용이 사고사슬 길이에 비례해 늘고, 그 비용의 상당 부분이 추론에 거의 기여하지 않는 서술문에 쓰이며, 한 번의 롤아웃이 초기 토큰에 확정적으로 커밋되므로 다른 경로를 탐색하려면 추가 생성이 필요하다. 잠재 추론은 중간 계산을 연속 상태로 옮기고 답만 언어화해 이 한계를 겨냥한다. 저자들은 효과적인 잠재 사고가 다섯 조건을 만족해야 한다고 주장한다. 답을 바꾸는 데 그치지 않고 정답에 이르게 하는 유용성, 재샘플링이 서로 다른 추론 경로를 낳는 다양성, 디코딩된 CoT가 답이 실제로 따르는 추론을 반영하는 설명 가능성, 더 많은 추론 연산으로 개선되는 정련 가능성, 그리고 비슷한 정확도에서 명시적 CoT보다 저렴한 효율성이다. 기존 잠재 방법들은 질문에서 지름길을 학습하거나, 명시적 CoT를 가중치로 증류하거나, 토큰 하나씩 흉내 내는 데 그친다고 지적한다.

어떻게 동작하나

제안 방법 FLaRe는 학습된 잠재 공간에서의 흐름 매칭을 쓴다. 먼저 VAE가 명시적 CoT를 잠재 코드로 압축하도록 학습된다. 인코더는 CoT 뒤에 M개의 학습된 슬롯 토큰을 붙여 읽고, 슬롯 위치의 최종 은닉 상태에서 평균과 로그분산을 얻어 z = μ + σ⊙ε로 코드를 샘플링한다. 디코더는 코드만 보고 CoT를 복원하며, β-VAE 목적함수로 학습한다. 그다음 흐름 모델이 질문을 조건으로 t=0의 가우시안 노이즈에서 t=1의 정제된 VAE 코드까지 선형 보간한 경로에서 상수 속도 z̄−ε를 회귀하도록 학습된다. 추론 시에는 질문 하나로 노이즈에서 시작해 S번의 오일러 스텝과 classifier-free guidance를 거쳐 잠재 사고를 만들고, 같은 공유 모델이 그 사고로부터 답을 자기회귀 생성한다. 기본 설정은 슬롯 M=8, 차원 d=512, 인코더 Llama-3.2-1B, 디코더 Llama-3.2-3B, 흐름 모델 Llama-3.2-1B-Instruct이며, GSM8K-Aug의 385K 질문으로 학습한다.

무엇과 다른가

잠재 공간 설계에서 가장 눈에 띄는 발견은 재구성만 높이면 안 된다는 것이다. d=512에서 슬롯을 8개에서 32개로 늘리면 재구성이 96.8%에서 99.4%로 오르지만 직접 정확도는 47.1%에서 25.4%로, 디코딩 정확도는 53.2%에서 32.3%로 떨어진다. 용량·손상·목적함수·인코더·디코더·데이터를 바꾼 아홉 번의 스윕에서 나온 66개 VAE를 보면 재구성과 디코딩 정확도의 상관은 −0.12이고, 재구성이 가장 좋은 VAE 중 여럿이 하류 성능은 최악에 가깝다. 코드가 정밀해질수록 흐름이 생성해야 할 세부가 늘어나기 때문이다. 그래서 저자들은 β를 10⁻⁵로 낮게 유지하고, 재구성을 조이는 대신 학습 중 입력과 코드를 손상시켜 잠재 공간을 매끄럽게 만든다. CoT 토큰을 확률 0.3으로 무작위 토큰으로 치환하고, 샘플링한 코드의 절반에 δ=0.7의 분산 보존 노이즈를 더하며, 확률 0.4의 잠재 드롭아웃을 적용한다. 세 손상은 재구성을 낮추면서도 두 정확도 지표를 모두 끌어올린다. 또한 자연어 CoT 대신 추론 단계와 중간 결과만 담은 기호 CoT를 인코더 입력으로 쓰고, 디코더는 코드만으로 기호 CoT를 쓰는 경로와 질문을 함께 받아 자연어 CoT를 쓰는 경로를 동시에 학습한다. 자연어 CoT를 입력으로 쓰는 VAE는 기본 설정보다 직접 정확도 약 10점, 디코딩 정확도 약 17점이 낮다. 1B 인코더를 고정했을 때 파인튜닝한 3B 디코더가 가장 강한 결과를 냈다.

어떻게 쓰나

흐름 학습에서는 네 가지 선택이 결정적이다. 첫째, 시간 분포를 노이즈 쪽으로 크게 치우쳐야 한다. 균등 분포가 가장 나쁘고, 추론이 시작되는 노이즈 근처에서 학습이 집중될 때만 잘 작동한다. 잠재 차원이 클수록 필요한 치우침이 커지므로 t=sigmoid(s), s~N(−2, 0.8²)인 로짓노멀 분포를 쓴다. 둘째, 답을 읽는 패스가 모델 자신의 불완전한 사고를 읽을 수 있어야 한다. 정확한 코드나 노이즈 낀 코드만으로 답 패스를 학습하면 생성된 사고를 잘못 읽으므로, 노이즈 낀 코드와 모델의 원스텝 예측을 50대 50으로 섞어 읽되 예측은 그래디언트에서 분리한다. 셋째, 질문당 여러 개의 CoT가 가장 큰 이득을 준다. 질문당 CoT가 하나면 각 질문이 단일 목표 코드에 묶여 확산 모델이 이를 암기하는 경향이 생기지만, 여러 CoT가 있으면 흐름이 서로 다른 노이즈를 서로 다른 유효 해로 매핑해야 한다. 같은 스타일의 새 질문은 이득이 작고, 다른 스타일의 질문은 정확도를 떨어뜨린다. 넷째, 2단계 자기학습이 필요하다. EMA 모델이 새 질문마다 20번의 가이드 스텝으로 K=8개의 사고를 생성하고, 동결된 VAE 디코더가 이를 기호 CoT로 읽어 참조 답과 일치하는 것만 남겨 다시 코드로 인코딩한다. 이후 흐름 손실과, 노이즈에서 10스텝 롤아웃 전체와 두 가이던스 분기를 통과해 역전파되는 답 손실을 함께 학습한다. 이 단계는 직접 정확도를 3.8점, 디코딩 정확도를 1.3점 올리며, 롤아웃 전체를 통과하는 답 손실을 제거하거나 단일 스텝으로 줄이면 이득 대부분이 사라진다.

전제와 한계

평가는 GSM8K 테스트 1,319문항과 GSM8K-Hard, SVAMP 1,000문항, MultiArith 180문항에서 이뤄진다. FLaRe는 산술 벤치마크에서 기존 잠재 방법보다 우수하며, 명시적 CoT 정확도의 97%를 4분의 1 지연으로 달성한다. 지연 측정에서 오일러 1스텝은 22ms, 명시적 CoT 토큰 하나는 9.7ms이므로 20스텝은 약 45토큰짜리 CoT와 맞먹는다. 직접 읽기 기준으로 1스텝에서 56.7점을 47ms에 내고, 2스텝에서 57.4점을 68ms에 내는데 이는 CODI의 절반, 명시적 CoT의 4분의 1 지연이다. 20스텝에서는 명시적 CoT에 거의 근접하지만 1.7배 지연이 된다.

다섯 조건 각각에 대한 프로브 결과도 제시된다. 유용성은 숫자 하나를 바꾼 쌍둥이 질문의 사고를 원래 질문에 주입해 답이 사고를 따르는지 보는 검사로 측정한다. 명시적 CoT는 97% 따르지만 CODI 32%, PCCoT 28%, Coconut 7%에 그친다. FLaRe는 전체 쌍에서 40%(1단계)와 37%(2단계)로 낮아 보이지만, 주입한 사고 자체가 틀린 경우가 많아 모델이 둘 중 어느 답도 내지 않는 경우가 대부분이며, 모델이 주입 없이도 정답을 내는 깨끗한 쌍에서는 91%와 76%를 따른다. 다양성은 질문당 16개 샘플로 측정해, 디코딩 읽기에서 1단계 모델이 greedy 61.2에 대해 pass@16 77.8, 질문당 서로 다른 답 3.4개, 투표 이득 3.3점, 모든 샘플이 반복하는 확정 오류 6%를 기록한다. 기존 방법들은 노이즈를 더하면 답이 5분의 1에서 절반까지 바뀌는데도 pass@16이 greedy보다 최대 5점 높은 데 그치고 투표가 greedy를 이기지 못한다. 설명 가능성은 참조 CoT의 중간 결과가 디코딩된 사고에 순서대로 나타나는지로 재는데, FLaRe가 46%와 49%로 CODI 39%, Coconut 19%, PCCoT 8%를 앞서고 사고 없이 읽는 대조군 3.8%와 우연 수준 0.3%를 크게 웃돈다. 정련 가능성에서는 Coconut이 평평하고 CODI가 예산 두 배에서 약 1점 떨어지며 PCCoT가 3회 반복 54.1점에서 12회 42.8점으로 하락하는 반면, FLaRe는 디코딩 읽기에서 1스텝 40.6점이 2스텝 58.2점, 20스텝 61.3점으로 오르고 40스텝에서도 1점 이상 잃지 않는다.

개발자 관점에서 이 논문의 실용적 메시지는 잠재 추론을 붙일 때 무엇을 측정해야 하는가에 있다. 재구성 손실이나 코드 복원률 같은 겉보기 지표는 하류 정확도와 음의 상관을 보일 수 있으므로, 잠재 표현을 도입한다면 생성된 사고를 그대로 읽었을 때의 정확도와 깨끗한 코드로 읽었을 때의 정확도를 분리해 측정해야 한다. 두 값의 격차는 생성된 사고가 VAE 공간에 얼마나 깔끔하게 안착하는지를 드러내며, 이 격차를 좁히는 것이 2단계 롤아웃 학습의 목적이다. 또한 추론 지연을 줄이려는 경우 1~2 오일러 스텝만으로도 명시적 CoT의 절반 이하 지연에서 근접한 정확도를 얻을 수 있다는 점, 그리고 그 대가로 정련 여지를 남겨두는 스텝 수 선택이 실질적인 튜닝 포인트라는 점을 시사한다.

저자들이 밝힌 한계도 분명하다. 정렬은 부분적이어서, 쌍둥이 쌍의 약 40%에서 디코딩된 사고가 쌍둥이의 값을 보여주는데도 답은 쌍둥이의 답이 되지 않는다. 대부분 그 값 이후에 쌍둥이 사고 자체가 잘못되어 답이 어느 쪽도 아닌 값으로 떨어지기 때문이다. 직접 읽기는 1스텝 이후 3점 미만만 개선되는데, 이는 1단계가 답 패스를 모델 자신의 원스텝 끝점으로 학습해 거친 사고에서도 답을 읽어내기 때문이다. 디코딩 읽기는 디코더가 토큰 단위로 CoT를 생성하므로 유효 코드 근처의 잠재를 잘 형성된 CoT로 바꿔줄 수 있어 더 관대하며, 그 자체로 명시적 CoT 생성에 의존한다. 2단계는 추가 CoT 주석이나 보상 모델 없이 참조 답만 필요로 하지만 검증 절차는 필수다. VAE 학습 데이터를 다양화하는 것은 재구성이 이미 포화 상태라 이득이 매우 제한적이며, 자연어 CoT만 있는 데이터셋을 기호 CoT로 바꾸는 파이프라인은 별도의 오프라인 LLM 단계로 수행되며 그 품질 검증에 의존한다.