기술 요약: MILES (Modular Instruction Memory with Learnable Selection)

문제 정의

대규모 언어 모델(LLM)은 Chain-of-Thought (CoT), Self-Consistency, Tree Search와 같은 테스트 타임 연산(test-time computation) 기술을 통해 상당한 추론 능력을 입증했습니다. 그러나 기존의 접근 방식은 각 문제를 개별적으로 처리하는 경로를 취합니다. 실제 시나리오에서 문제가 순차적으로 발생하는 경우, 모델은 관련 쿼리 전반에 걸쳐 경험을 축적하고 재사용하는 데 실패합니다. 매 새로운 문제마다 모델 파라미터를 업데이트하는 것은 특히 폐쇄형 또는 고정된(frozen) 모델의 경우 실질적으로 불가능합니다.

외부 메모리 방법이 존재하기는 하지만, 다음과 같은 결정적인 한계가 있습니다:

1. 낮은 일반화 능력: 전체 솔루션 템플릿이나 전략을 저장하는 방식은 새로운 문제가 새로운 구조를 가질 때 실패하는 경우가 많습니다.
2. 휴리스틱 선택: 단계별(step-level) 메모리 접근 방식은 일반적으로 최종 정답의 정확성을 최적화하기보다는 휴리스틱한 선택(예: 유사도 검색 또는 프롬프팅)에 의존합니다.
3. 학습 제약: 선택 정책을 학습하려면 대규모 레이블 데이터셋, 고정된 액션 공간 또는 파라미터 업데이트가 필요한데, 이는 메모리가 점진적으로 확장되고 감독(supervision)이 제한되는 테스트 타임 설정과 호환되지 않습니다.

핵심 과제는 밑바탕이 되는 LLM 파라미터를 업데이트하지 않으면서, 경험이 축적됨에 따라 추론 정확도를 향상시키는 외부 메모리를 조직, 유지 및 재사용하는 것입니다.

방법론: MILES 프레임워크

MILES(Modular Instruction Memory with LEarnable Selection)는 재사용 가능한 단계별 추론 모듈의 축적과 구성을 통해 테스트 타임 자기 개선(self-improvement)을 가능하게 함으로써 이러한 과제를 해결합니다. 이 프레임워크는 고정된(frozen) LLM 위에서 작동하며 네 가지 핵심 구성 요소로 이루어져 있습니다.

1. 모듈형 비대칭 메모리 구조 (Modular Asymmetric Memory Structure)

MILES는 메모리를 비대칭 쌍 $(g_i, u_i)$로 정의된 모듈형 단위의 집합으로 저장합니다:

• 하위 목표 임베딩 ($g_i$): 추론 단계의 목적을 요약하는 밀집 벡터(dense vector)입니다. 이는 검색 키 역할을 합니다.
• 하위 지시문 ($u_i$): 해당 유닛이 적용될 때 LLM의 생성을 조건화하는 데 사용되는 자연어 지시문입니다.
  보조 모델(selection head)은 현재 추론 상태가 주어졌을 때 해당 유닛의 적용 가능성을 예측하는 가볍고 학습 가능한 선택 헤드($\pi_{\theta_a}$)와 결합됩니다.

이 구조는 점진적 확장을 가능하게 합니다. 기존 구성 요소를 재학습하지 않고도 새로운 유닛과 그에 대응하는 헤드를 추가할 수 있습니다.

2. 조대-정밀(Coarse-to-Fine) 검색 메커니즘

검색 효율성과 선택 정확도의 균형을 맞추기 위해, MILES는 두 단계의 선택 프로세스를 채택합니다:

• Layer 1 (조대 검색 - Coarse Retrieval): 하위 목표 유사성을 기반으로 합니다. LLM은 후보 단계 목표를 생성하고, 이들의 임베딩은 저장된 하위 목표들과 코사인 유사도를 사용하여 비교됩니다. 유닛당 유효 임계값(effective threshold)을 통해 무관한 유닛을 필터링합니다. 이 레이어는 메모리 확장과 학습 데이터 수집을 용이하게 합니다.
• Layer 2 (학습된 적용 가능성 점수 산출 - Learned Applicability Scoring): 학습된 선택 헤드는 전체 추론 상태를 기반으로 검색된 후보들의 적용 가능성을 점수화합니다. 가변 길이의 상태를 처리하기 위해, MILES는 컨텍스트 코드북(과거 단계 목표 및 조건의 클러스터링된 임베딩)을 사용하여 max-similarity pooling을 통해 고정 길이 특징 벡터를 생성합니다.
• 결합 점수: 최종 선택은 Layer 1의 의미적 유사성과 Layer 2의 학습된 적용 가능성 점수를 결합합니다.

3. 학습 가능한 메모리 구성 (Learnable Memory Composition)

MILES는 메모리 선택을 순차적 의사결정 문제로 취급합니다. 목표는 올바른 최종 정답의 확률을 최대화하는 것입니다.

• 학습 데이터 수집: "확신도가 높은(confident)" 샘플(여러 LLM 롤아웃이 동일한 답에 동의하는 경우)에 대해, 시스템은 Layer 1 샘핑을 사용하여 몬테카를로 트리 탐색(MCTS)을 실행합니다. 이를 통해 하위 지시문이 적용된 궤적(trajectory)을 생성하며, 결과적인 정확성이 의사 레이블(pseudo-label)로 사용됩니다.
• 최적화: 선택 헤드는 수집된 이러한 궤적들을 사용하여 이진 교차 엔트로피(binary cross-entropy)로 학습됩니다. 이 목적 함수는 KL-divergence 제약 조건 하에서 도출된 최적의 정책을 근사하며, 헤드가 정답으로 이어지는 유닛을 선택하도록 유도하는 동시에 one-vs-rest 분류 및 앙상블 방법을 통해 과적합을 방지합니다.

4. 추론 파이프십 (Inference Pipeline)

시스템은 신뢰도에 따라 테스트 샘플을 라우팅합니다:

• 확신도가 높은 샘플: 메모리를 확장하고 선택 헤드를 학습하는 데 사용됩니다. LLM이 응답을 생성하며, 일관성이 있는 경우 해당 궤적은 MCTS 데이터 수집을 위해 사용됩니다.
• 불확실한 샘플: 메모리 가이드 트리 탐색으로 라우팅됩니다. 시스템은 Layer 1을 통해 후보를 검색하고, Layer 2를 통해 이를 재순위화(rerank)하며, 최적의 하위 지시문을 적용하여 다음 추론 단계를 가이드함으로써 최종 답변에 도달할 때까지 반복합니다.

주요 기여

1. MILES 프레임워크: 파라미터 업데이트, 외부 감독 또는 고정된 액션 공간 없이도, 확신 있는 추론 궤적으로부터 모듈형 지시 메모리를 점진적으로 구축하고 학습된 메모리를 불확실한 사례에 적용하는 새로운 테스트 타임 자기 개선 프레임워크를 제시합니다.
2. 모듈형 비대칭 메모리: (하위 목표, 하위 지시문) 쌍과 유닛당 선택 헤드를 갖춘 구조로, 점진적인 메모리 확장과 학습 가능한 구성을 모두 지원하는 조대-정밀 메커니즘을 가능하게 합니다.
3. 학습 가능한 메모리 구성: 확신 있는 샘플의 롤아웃 결과로부터 가벼운 선택 헤드를 학습하는 방법을 제공하여, 제한된 테스트 타임 데이터 하에서도 자기 지도 방식의, 정확도에 최적화된 메모리 구성을 가능하게 합니다.
4. 실증적 검증: 6개의 추론 벤치마크와 4개의 백본 모델에 대해 MILES를 평가하여, 기존 방법론들을 능가하거나 대등한 성능을 보여주는 광범한 실험을 수행했습니다.

실험 결과

저자들은 네 가지 백본(GPT-4.1, GPT-4.1-mini, GPT-OSS-20B, Qwen3-30B-Instruct)을 사용하여 6개의 벤치마크(MATH-500, AIME 2024/2025, GPQA-Diamond, MMLU-Pro Physics/Engineering)에서 MILES를 평가했습니다.

• 정확도: MILES는 기존 베이스라인 방법들(Zero-shot CoT, Self-Consistency, Buffer-of-Thoughts, Dynamic CheatSheet 포함)과 대등하거나 이를 능가하는 성능을 일관되게 보여주었습니다. 특히 다른 메모리 기반 방법들이 zero-shot CoT보다 나은 성능을 내지 못했던 작은 규모의 오픈 웨이트 백본에서 주목할 만한 개선이 관찰되었습니다.
• 효율성: MILES는 정확도 대비 토큰 프런티어(accuracy-vs-token frontier)에서 우위를 점합니다. 계산량이 증가함에 따라 성능이 정체되는 베이스라인들과 달리, MILES는 추가적인 메모리 유닛을 구축함으로써 토큰 예산이 늘어남에 따라 지속적으로 향상됩니다.
• 전이 가능성: 이 프레임워크는 교차 모델 전이를 지원합니다. 보조 모델(예: Qwen3-30B)의 롤아웃으로부터 구축된 메모리는 다른 더 강력한 고정 백본(예: GPT-4.1-mini)의 성능을 효과적으로 향상시킵니다.
• 절제 연구(Ablation Studies): 조대 검색 레이어나 학습된 선택 헤드 중 하나라도 제거하면 성능이 크게 저하되어, 두 레이어 메커니즘의 필요성을 입증했습니다. 학습된 헤드는 상태 유사성이나 LLM 프롬프팅에 기반한 대안적인 재순위화 모델보다 우수한 성능을 보였습니다.

의의 및 주장

본 논문은 모듈형 메모리 구성이 고정된 LLM의 테스트 타임 자기 개선을 위한 실행 가능하고 효과적인 메커니즘임을 입증한다고 주장합니다.

주요 의의는 다음과 같습니다:

• 상보성: 단계별 메모리 구성 접근 방식은 기존의 프롬프팅 및 테스트 타임 탐색 방법과 상보적이며, 축적된 경험을 통해 개선을 위한 추가적인 경로를 제공합니다.
• 강건성: 이 방법은 다양한 모델 제품군과 데이터 분포에 대해 강건하며, 외부 프로세스 보상 모델(PRM)이나 정제된 학습 데이터셋에 의존하지 않습니다.
• 확장성: 메모리 선택을 파라미터 업데이트와 분리함으로써, MILES는 모델 가중치를 수정할 수 없는 환경에서 지속적인 개선을 위한 경로를 제공합니다.

저자들은 학습 가능한 재사용 메모리 선택 지식을 배우는 것이 지속적으로 발전하는 추론 시스템을 구축하는 유망한 방향이라고 결론지었으나, 롤아웃 기반 데이터 수집의 계산 오버헤드와 기초 LLM의 지시 이행 능력에 대한 의존성이라는 한계점 또한 인정하였습니다.