기술 요약: 레이어를 건너뛸 것인가, 반복할 것인가? LLM의 레이어 프로그램 학습(Learning Program-of-Layers)

1. 문제 정의

현재 대규모 언어 모델(LLM)은 입력의 복잡도와 관계없이 모든 트랜스포머 레이어를 순차적으로 실행하는 **고정된 깊이와 고정된 순서의 순전파(forward pass)**를 통해 추론을 수행합니다. 이러한 정적 아키텍처는 입력 난이도의 높은 편차를 반영하지 못하며, 다음과 같은 두 가지 비효율성을 초래합니다:

1. 과잉 계산(Over-computation): 단순한 입력의 경우 더 적은 레이어로도 올바르게 해결될 수 있음에도 불구하고, 모델은 전체 깊이를 모두 실행합니다.
2. 계산 부족(Under-computation): 표준 순전파 방식에서 실패한 복잡한 입력이라도, 모델이 실행 경로(예: 특정 레이어를 반복하여 정교화하거나 불필요한 레이어를 건너뛰는 방식)를 동적으로 변경할 수 있다면 해결 가능할 수 있습니다.

본 논문은 표준 순전파가 LLM의 잠재적 추론 능력 중 극히 일부만을 포착한다고 주장합니다. 핵심 질문은 일반적인 모델이 베이스 모델을 재학습하지 않고도 각 특정 입력에 대해 자신의 "프로그램"(레이어 실행 시퀀스)을 최적화할 수 있는가 하는 점입니다.

2. 방법론

2.1. 몬테카를로 트리 탐색(MCTS)을 통한 경험적 발견

제안된 학습 기반 솔루션에 앞서, 저자들은 MCTS를 사용하여 **레이어 프로그램(Program-of-Layers, POLAR)**의 공간을 탐색하는 경험적 연구를 수행합니다.

• 정의: 레이어는 원자적 함수($f_0, \dots, f_{D-1}$)의 라이브러리로 취급됩니다. 프로그램 $\pi$는 건너뛰기(레이어 생략)와 재귀(레이어 반복)를 허용하는 레이어 인덱스의 시퀀스입니다.
• 발견 사항:
  • 결합 최적화: 건너뛰기와 반복 연산을 결합하는 것이 각각을 단독으로 사용하는 것보다 유의미하게 더 나은 정확도를 보였습니다.
  • 오컴의 면도날(Occam's Razor): 많은 입력(초기에 오답을 냈던 입력 포함)에 대해, 유효한 실행 프로그램은 표준 순전파보다 더 짧은 형태로 존재합니다.
  • 테스트 타임 스케일링(Test-Time Scaling): (재귀를 통한) 깊이 예산을 늘리면 어려운 작업에 대한 정확도가 체계적으로 향상되며, 이는 유효한 잠재 계산의 연속체가 존재함을 보여줍니다.
  • 구조적 편향(Structural Bias): 유효한 프로그램은 주로 연속적인 레이어 세그먼트로 구성되며, 한 세그먼트 내에서 단일 재귀 이상을 요구하는 경우는 드뭅니다. 장거리 점프나 깊은 반복 사용은 흔치 않습니다.

2.2. POLAR 프레임워크

MCTS의 발견에 착안하여, 저자들은 고정된 사전 학습된 모델의 파라미터를 수정하지 않고 추론 시점에 입력별 실행 프로그램을 예측하는 경량 프레임워크인 POLAR를 제안합니다.

• 프로그램 표현:

  • 세그먼트화(Segmentation): $D$개의 레이어를 연속적인 세그먼트 $[s_j, s_{j+1})$로 분할합니다.
  • 연산: 각 세그먼트에는 세 가지 연산 중 하나가 할당됩니다: Skip(생략), Keep(표준 실행), Repeat(두 번 실행).
  • 제약 조건: 유효한 프로그램이 국소적 구조를 선호한다는 발견에 근거하여, 세그먼트 길이는 경험적으로 $K_{max}=4$로 제한됩니다.

• 예측 네트워크 아키텍처:

  • 입력 인코딩: 입력 $x$는 고정된 임베딩 모델을 통해 토큰 임베딩으로 인코딩됩니다.
  • 레이어 쿼리: 학습 가능한 임베딩이 각 레이어 인덱스를 나타내는 쿼리 역할을 합니다.
  • 교차 주의 집중(Cross-Attention): 교차 주의 집중 메커니즘은 레이어 쿼리를 토큰 임베딩과 결합하여 각 레이어에 대한 입력 조건부 표현을 생성합니다.
  • 교차 레이어 인코더: 가벼운 트랜스포머 인코더가 깊이 차원 간의 의존성을 모델링하여, 한 레이어에서의 결정이 전역적 문맥에 의존할 수 있도록 합니다.
  • 예측 헤드: 두 개의 헤드가 (1) 세그먼트 경계(이진 마스크)와 (2) 각 세그먼트 시작점의 연산 라벨(skip/keep/repeat)을 예측합니다.

• 학습 및 추론:

  • 감독(Supervision): 네트워크는 오프라인에서 MCTS를 통해 발견된 정답 프로그램(ground-truth programs)을 바탕으로 학습됩니다.
  • 추론: 네트워크는 세그먼트 경계와 연산을 예측합니다. 고정된 LLM을 실행하기 전, 전역적 일관성을 보장하기 위해 세그먼트 수준의 연산에 대해 빔 서치(beam search)가 적용됩니다.

3. 주요 결과

3.1. 분포 내 성능 (DART-Math)

네 가지 모델(LLaMA-3.2-3B, Qwen1.5-MoE, Qwen2.5-3B, Qwen3-8B)에 대해 수학적 추론 벤치마크(DM-1 ~ DM-5)에서 평가되었습니다:

• 정확도 향상: POLAR는 표준 추론(Base) 및 기존의 동적 깊이 방법들(예: MindSkip, FlexiDepth, DR.LLM)보다 일관되게 우수한 성능을 보였습니다.
  • 예시: LLaMA-3.2-3B의 경우, POLAR는 DM-1에서 **+46.8%**의 절대적 이득을, 가장 어려운 DM-5에서 베이스 모델 대비 **+32.0%**의 이득을 달고했습니다.
• 테스트 타임 스케일링: 빔 서치를 통해 후보 프로그램의 수($k$)를 늘리면 정확도가 단조 증가하며, 이는 표준 확률적 샘플링(Base sampling)의 수확 체감 현상을 훨씬 상회합니다.
• 효율성: POLAR는 종종 더 적은 고유 레이어를 실행하면서도 더 높은 정확도를 달성합니다.

3.2. 분포 외(OOD) 일반화

DART-Math로 학습된 POLAR는 ASDiv, MAWPS, MMLU-Pro(수학, 물리, 화학, 법률 등)에 대해 제로샷(zero-shot) 평가를 받았습니다:

• 교차 도메인 전이: POLAR는 산술 문장제 및 일반 지식 과제를 포함한 다양한 도메인에서 베이스 순전파 및 다른 동적 방법들보다 일관되게 우수한 성능을 보였습니다.
• 메커니즘: 저자들은 이를 예측 헤드가 입력을 공유된 의미 공간으로 매핑하고, 단순한 구조적 패턴(연속 세그먼트)으로 제약을 둔 덕분이라고 설명합니다. 이는 데이터셋 특화된 휴리스틱 대신 재사용 가능한 전략을 장려합니다.

3.3. 효율성 및 오버헤드

• 지연 시간(Latency): 예측 네트워크는 표준 순전파의 약 0.8% 수준의 미미한 오버헤드(약 0.23개 레이어 분량)만을 추가합니다.
• 순이익: 쉬운 입력에 대해 레이어를 건너뜀으로써, POLAR는 더 높은 정확도를 달성하는 동시에 엔드 투 엔드 실행 시간을 단축합니다(예: 쉬운 입력에서 베이스 런타임의 0.83×).
• 파라미터 수: 예측기는 약 2.1M 개의 파라미터만을 추가하며, 이는 베이스 모델 크기의 0.01%~0.06%에 불과합니다.

4. 의의 및 주장

본 논문은 고정된 깊이의 실행이 범용 파운데이션 모델에 있어 최적이 아님을 주장합니다. 레이어를 재사용 가능한 함수의 라이브러리로 취급함으로써, POLAR는 다음을 입증했습니다:

1. 잠재 역량(Latent Capacity): LLM은 표준 순전파보다 오류를 수정하거나 작업을 더 효율적으로 해결할 수 있는 방대하고 유효한 대안적 계산(프로그램) 공간을 보유하고 있습니다.
2. 실용성: 이러한 잠재 역량은 값비싼 탐색을 단일 단계 예측으로 대체하는 경량화된, 베이스 모델의 학습이 필요 없는(training-free) 예측 메커니즘을 통해 실현될 수 있습니다.
3. 일반화: 실행 깊이와 순서를 동적으로 재구성하는 능력은 아키텍처의 재설계나 파인튜닝 없이도 다양한 도메인에서 성능을 향상시키는 전이 가능한 기술입니다.

이 연구는 미래의 추론 전략이 정적 깊이를 넘어, 입력 난이도에 맞춰 계산을 동적으로 재배분함으로써 더 효율적이고 표현력이 풍부한 파운데이션 모델 배포의 길을 제시해야 함을 시사합니다.