기술 요약: 선택적 레이어 복원을 통한 다양성 회복

문제 정의

지시어 튜닝(Instruction Tuning) 및 인간 피드백 기반 강화 학습(RLHF)과 같은 사후 학습(Post-training) 기법은 대규모 언어 모델(LLM)의 지시 이행 능력과 유용성을 크게 향상시킨다. 그러나 이러한 방법들은 모드 붕괴(Mode Collapse) 현상을 유발하는 경우가 많은데, 이는 모델이 좁은 범위의 응답에 불균형적으로 높은 확률 질량을 할당함으로써, 다양하지만 유효한 대안들을 억제하는 현상을 의미한다. 이는 창의적 글쓰기, 개방형 질의응답(QA), 추론과 같이 다양한 응답을 가치 있게 여기는 애플리케이션에서 모델의 유용성을 제한하며 반복적인 출력을 생성하게 만든다.

기존의 완화 전략은 세 가지 범주로 나뉘며, 각각 한계를 지닌다:

1. 디코딩 기반 방법 (예: 온도 스케일링, min-p 샘플링): 무작위성을 높이지만, 사후 학습 과정 중 파라미터 변화로 인해 손실된 다양성을 복구할 수는 없다.
2. 프롬프트 기반 방법 (예: 여러 개의 답변 요청): 추론 비용을 증가시키고 프롬프트 문구에 민밀하게 반응한다.
3. 학습 기반 방법 (예: 목적 함수 재설계): 비용이 많이 드는 재학습과 데이터 수집을 요구한다.

저자들은 다양성을 효과적으로 복구하면서도, 적응 과정에서의 계산 효율성이 높고, 추론 시 추가적인 비용을 발생시키지 않는 방법을 찾고자 한다.

방법론

제안된 접근 방식인 **선택적 레이어 복원(Selective Layer Restoration, SLR)**은 두 가지 경험적 관찰에 기반한, 학습이 필요 없는(training-free) 가중치 공간 개입 방식이다:

1. LLM 레이어는 뚜렷한 기능적 역할(예: 사실적 연관성은 종종 중간 깊이의 레이어에 국한됨)을 가진다.
2. 사후 학습은 주로 어떤 사전 학습된 행동을 표현할지를 조절(이로 인해 발생하는 파생적 망각/catastrophic forgetting 포함)할 뿐, 완전히 새로운 지식을 습득하는 것이 아니라 기존의 사전 학습된 행동을 억제하거나 유도하는 것이다.

저자들은 모드 붕괴가 특정 레이어에 국한되어 있으며, 사후 학습된 모델($M_{post}$)의 특정 레이어 범위를 사전 학습된 가중치($M_{pre}$)로 복원함으로써 품질을 유지하면서도 다양성을 회복할 수 있다는 가설을 세웠다.

SLR 프레임워크

$N$개의 레이어를 가진 모델에 대해, SLR은 $M_{post}$의 연속된 레이어 구간 $[i, j]$를 그에 대응하는 $M_{pre}$의 레이어로 교체하여 하이브리드 모델 $M_{i:j}$를 구축한다. 목표는 최소 품질 임계값($Q \ge q_{min}$)을 준수하면서 다양성($D$)을 극대화하는 구간 $(i, j)$를 찾는 것이다.

CRC 대리 작업 (Proxy Task)

복잡한 다운스트림 작업에서 복원 구간을 직접 최적화하는 것은 계산 비용이 많이 들고, "품질"의 주관성으로 인해 평가가 어렵다. 이를 해결하기 위해 저자들은 **제약된 무작위 문자(Constrained Random Character, CRC)**를 대리 작업으로 도입했다.

• 메커래즘: CRC는 엄격하게 정의된 유효 집합 내에서 무작위 문자(숫자 또는 알파벳)를 생성하는 것을 포함한다 (예: "0에서 5 사이의 무작위 정수를 생성하라").
• 지표:
  • 품질 ($Q_{CRC}$): 유효한 출력에 할당된 확률 질량 (지시 이행 능력 측정).
  • 다양성 ($D_{CRC}$): 유효한 출력 분포에 대한 엔트로피 (다양성 측정).
• 선택 과정: 저자들은 모든 가능한 연속 레이어 구간에 대해 CRC를 사용하여 그리드 탐색을 수행한다. 이들은 $Q_{CRC} \ge 0.9 \times Q_{CRC}(M_{post})$를 만족하면서 $D_{CRC}$를 극대화하는 구간을 선택한다. 선택된 구간은 추가적인 튜닝 없이 다운스트림 작업에 적용된다.

주요 기여

1. 선택적 레이어 복원 (SLR): 사후 학습된 LLM의 특정 연속 레이어 구간을 사전 학습된 파라미터로 복원하는 학습이 필요 없는 방법이다. 이는 원래의 아키텍처와 파라미터 수를 보존하며, 추가적인 추론 비용을 발생시키지 않는다. 또한 디코딩 및 프롬프트 기반 다양성 방법과 상보적임이 입증되었다.
2. 제약된 무작위 문자 (CRC): 명시적인 유효 집합과 자연스러운 다양성 목적을 가진 대리 작업이다. CRC는 효율적이고 저비용인 복원 구간 탐색을 가능하게 하며, 레이어 국소화 가설에 대한 경험적 근거를 제공한다.
3. 경험적 검증: 세 가지 모델 제품군(Llama, Qwen, Gemma)과 세 가지 다운스트림 작업(창의적 글쓰기, 개방형 QA, 다단계 추론)에 대한 포괄적인 평가.

실험 결과

저자들은 SLR을 사후 학습된 베이스라인 및 "모델 수프(Model Soup)" 베이스라인(사전 학습 모델과 사후 학습 모델의 가중치 평균화)과 비교 평가하였다.

• 창의적 글쓰기: SLR은 최소한의 품질 손실(평균 -2.1% LLM-judge 점수)과 함께 일관되게 다양성을 개선했다(임베딩 기반 비유사도 평균 +42.5% 증가). 반면, 모델 수프는 더 작은 다양성 이득(+27.3%)과 더 심각한 품질 저하(-6.0%)를 보였다.
• 개방형 QA: SLR은 답변 다양성을 실질적으로 증가시켰으며, 엔트로피 +112.7% 증가 및 커버리지(고유한 정답의 비율) +169.5% 증가를 보여주는 동시에 높은 정밀도(평균 -1.8% 하락)를 유지했다.
• 추론 (GSM8K): 그리디 디코딩(greedy decoding) 하에서 사후 학습된 모델이 실패하는 문제들에 대해, SLR은 모든 샘플링 예산($k \in \{1, \dots, 64\}$)에서 Pass@k를 개선했다. 이는 확률 질량이 다양하지만 올바른 추론 경로 쪽으로 이동했음을 나타낸다.
• 결합 가능성 (Composability): SLR은 높은 온도 디코딩($T=1.5$) 및 언어화 샘플링(verbalized sampling, 프롬프트 기반 방법)과 결합했을 때도 효과가 유지되었다. 이는 SLR이 밑바닥의 모델 분포를 수정하는 반면, 다른 방법들은 샘플링 과정이나 입력을 조작한다는 것을 시사한다.
• 절제 연구 (Ablation Studies): CRC 대리 작업에 의해 가이드된 레이어 복원이 임의의 구간(예: 가장 초기 또는 가장 마지막 레이어)을 복원하는 것보다 우수한 결과를 냈다. 임의의 레이어를 복원하는 것은 다양성을 개선하지 못하거나 품질을 저하시키는 경우가 많았으며, 이는 특정 구간 선택이 매우 중요하다는 것을 확인시켜 준다.

의의 및 주장

본 논문은 SLR이 재학습의 계산 오버헤드나 복잡한 프롬프팅의 추론 비용 없이 모드 붕합을 완화할 수 있는 단순하면서도 효과적인 경로를 제공한다고 주장한다. 특정 가중치 공간 개입을 통해 다양성을 회복할 수 있음을 입증함으로써, 본 연구는 사후 학습된 LLM에 대한 모듈형 관점을 뒷받침한다:

1. 기능과 행동(모드 붕합 포함)은 모델 전체에 균일하게 분포되는 것이 아니라 특정 레이어에 국한되어 있다.
2. 사후 학습은 근본적인 핵심 지식 베이스를 변경하기보다는 사전 학습된 행동의 표현을 조절하는 것이다.

저자들은 결론적으로, SLR과 같은 미세한 가중치 공간 개입이 품질-다양성 트레이드오프를 제어하는 유망한 방안을 제공하며, 이를 통해 사후 학습의 이점(지시 이행 등)을 유지하면서도 바람직한 사전 학습 특성(다양성 등)을 회복할 수 있다고 밝히고 있다. 본 논문은 현재 약 7~9B 파라미터 모델에 국한되어 있다는 점을 언급하며 확장성에 대해 겸허한 태도를 유지하였으며, 이러한 결과가 더 큰 규모의 모델이나 더 세밀한 개입(예: 서브 레이어 복원)에서도 유효한지 판단하기 위해서는 향후 연구가 필요하다고 명시했다.