[논문 리뷰] Think-at-Hard: Selective Latent Iterations to Improve Reasoning Language Models

TL;DR

본 논문에서는 대형 언어 모델(LLM)의 추론 능력을 향상시키기 위한 새로운 방법론인 Think-at-Hard (TaH)를 제안한다. TaH는 LLM이 모든 토큰에 대해 고정된 횟수의 반복 연산을 수행하는 대신, 첫 번째 순방향 패스 이후 '어려운' 토큰으로 식별된 토큰에 대해서만 추가적인 반복을 수행함으로써 계산 효율성을 높인다. 이를 위해 경량의 신경 결정기(Neural Decider)를 사용하여 각 토큰의 난이도를 판단하고, 어려운 토큰에 대해서만 Low-Rank Adaptation (LoRA) 모듈을 적용하여 LLM의 학습 목표를 어려운 토큰의 정교화로 전환한다. 또한, 듀오 인과적 주의 메커니즘(Duo-Causal Attention Mechanism)을 도입하여 토큰 시퀀스 차원뿐만 아니라 반복 깊이 차원으로 주의 메커니즘을 확장하여 반복 간 정보 흐름을 원활하게 한다. 실험 결과, TaH는 다양한 추론 벤치마크에서 기존 방법 대비 성능 향상을 보였으며, 특히 파라미터 효율적인 방식으로 LLM의 추론 능력을 개선할 수 있음을 입증했다. TaH는 LLM의 추론 능력 개선과 계산 자원 효율성이라는 두 가지 중요한 목표를 동시에 달성할 수 있는 잠재력을 지닌다.

연구 배경 및 동기

최근 몇 년 동안 대형 언어 모델(LLM)은 자연어 처리(NLP) 분야에서 혁명적인 발전을 이루어냈다. LLM은 텍스트 생성, 번역, 질의 응답 등 다양한 작업에서 인간에 필적하는 성능을 보여주며, 그 활용 범위가 빠르게 확장되고 있다. 그러나 LLM의 추론 능력은 여전히 개선해야 할 여지가 많다. 특히, 복잡한 논리적 추론, 수학적 추론, 상식적 추론 등은 LLM에게 여전히 어려운 과제이다.

기존의 LLM은 일반적으로 고정된 횟수의 순방향 패스(forward pass)를 통해 텍스트를 생성한다. 즉, 각 토큰에 대해 동일한 연산을 수행하며, 이는 계산 자원의 낭비로 이어질 수 있다. 특히, 쉬운 토큰의 경우, 한 번의 순방향 패스로도 충분히 정확한 예측이 가능하지만, 추가적인 반복 연산은 오히려 예측 정확도를 떨어뜨리는 잠재적 과잉 사고(latent overthinking) 현상을 유발할 수 있다.

반복적 변환기(recurrent transformer)와 같은 이전 연구에서는 토큰당 고정된 횟수의 추가 반복을 할당하여 생성 품질을 향상시켰다. 그러나 이러한 접근 방식은 모든 토큰에 대해 동일한 연산을 반복하므로, 계산 효율성이 떨어진다는 단점이 있다. 또한, 쉬운 토큰에 대한 불필요한 반복 연산은 잠재적 과잉 사고 현상을 심화시킬 수 있다.

따라서, 본 연구에서는 LLM의 추론 능력을 향상시키면서도 계산 자원 효율성을 높일 수 있는 새로운 방법론이 필요하다는 문제 의식에서 출발한다. 구체적으로, 다음과 같은 연구 질문에 답하고자 한다.

1. LLM의 추론 과정에서 잠재적 과잉 사고 현상은 실제로 발생하는가?
2. 어려운 토큰과 쉬운 토큰을 구분하여 선택적으로 반복 연산을 수행하는 것이 LLM의 추론 능력 향상에 도움이 되는가?
3. 어려운 토큰을 식별하고, 해당 토큰에 대한 반복 연산을 수행하는 효율적인 방법은 무엇인가?
4. 선택적 반복 연산을 통해 LLM의 추론 능력을 향상시키면서도 계산 자원 효율성을 높일 수 있는가?

본 연구는 이러한 연구 질문에 답하기 위해 Think-at-Hard (TaH)라는 새로운 방법론을 제안한다. TaH는 LLM이 모든 토큰을 반복적으로 처리하는 대신, '어려운' 토큰에 집중함으로써 효율성과 정확성을 높이는 것을 목표로 한다. 이를 통해 LLM의 추론 능력을 향상시키고, 다양한 실제 응용 분야에서 LLM의 활용도를 높이는 데 기여하고자 한다.

관련 연구

LLM의 추론 능력을 향상시키기 위한 다양한 연구가 진행되어 왔다. 본 섹션에서는 주요 선행 연구를 분석하고, 본 논문과의 차별점을 설명한다.

1. Chain-of-Thought Prompting (CoT): CoT는 LLM에게 문제 해결 과정을 단계별로 설명하도록 유도하여 추론 능력을 향상시키는 방법이다. Wei et al. (2022)은 CoT가 LLM의 추론 능력을 크게 향상시킬 수 있음을 보여주었다. 그러나 CoT는 프롬프트 엔지니어링에 의존하며, 최적의 프롬프트를 찾는 것이 어려울 수 있다는 단점이 있다.

2. Self-Consistency Decoding: Self-Consistency Decoding은 LLM이 생성한 여러 개의 답변을 비교하여 가장 일관성 있는 답변을 선택하는 방법이다. Wang et al. (2022)은 Self-Consistency Decoding이 CoT와 함께 사용될 때 LLM의 추론 능력을 더욱 향상시킬 수 있음을 보여주었다. 그러나 Self-Consistency Decoding은 여러 개의 답변을 생성해야 하므로, 계산 비용이 높다는 단점이 있다.

3. Recurrent Transformers: Recurrent Transformers는 LLM의 마지막 레이어의 hidden state를 다시 입력으로 사용하여 반복적으로 텍스트를 생성하는 방법이다. 이러한 방법은 LLM이 더 깊이 생각하고, 더 나은 추론 결과를 얻도록 돕는다. 그러나 Recurrent Transformers는 모든 토큰에 대해 동일한 연산을 반복하므로, 계산 효율성이 떨어진다는 단점이 있다.

4. Active Inference: Active Inference는 LLM이 환경과 상호작용하면서 학습하는 방법이다. 이러한 방법은 LLM이 더 능동적으로 추론하고, 더 나은 결정을 내리도록 돕는다. 그러나 Active Inference는 복잡한 환경 모델링이 필요하며, 학습 과정이 불안정할 수 있다는 단점이 있다.

5. Knowledge Distillation: Knowledge Distillation은 큰 모델(teacher model)의 지식을 작은 모델(student model)에게 전달하는 방법이다. 이러한 방법은 작은 모델의 추론 능력을 향상시키면서도 모델 크기를 줄일 수 있다. 그러나 Knowledge Distillation은 teacher model의 성능에 크게 의존하며, student model이 teacher model의 모든 지식을 학습하기 어려울 수 있다는 단점이 있다.

본 논문과의 차별점