기술 요약: MoSE – 계층적 자기 증류(Self-Distillation)를 통한 초기 레이어 임베딩 강화

1. 문제 정의

생산 환경에 대규모 언어 모델(LLM)을 배포할 때는 엄격한 지연 시간(latency) 및 자원 제약을 충족하기 위해 정확도와 성능 사이의 트레이드오프를 조율해야 하는 경우가 많습니다. 전통적인 모델 증류(distillation) 방식은 모델 크기를 줄여주지만, "스튜던트(student)" 모델이 "티처(teacher)" 모델을 모방하도록 별도의 학습 과정을 거쳐야 하므로 상당한 비용이 발생합니다. 또한, 기존 방식들은 대개 트랜스포머의 마지막 레이어가 가장 유용한 표현(representation)을 생성한다고 가정하지만, 최근의 통찰에 따르면 특정 다운스트림 태스크를 위해서는 중간 레이어가 가장 풍부한 의미론적 정보를 보유하고 있을 수 있습니다. 따라서 성능 저하나 별도의 증류 파이프라인 없이도 유연한 배포(다양한 모델 부분 선택)가 가능한 통합된 아키텍처가 필요합니다.

2. 방법론

저자들은 StarCoder-2 아키텍처를 기반으로 구축된 10억 파라미터 규모의 멀티 엑시트(multi-exit) 인코더인 **MODULARSTARENCODER (MOSE)**를 소개합니다. 핵심 혁신은 단일 모델 스택 내에서 여러 엑시트 포인트(exit points)를 동시에 학습시켜 하위 레이어의 표현력을 높이는 자기 증류(Self-Distillation) 메커니즘입니다.

아키텍처 및 학습 전략

• 멀티 엑시트 설계: 이 모델은 36개의 은닉 레이어를 특징으로 하며, 레이어 4, 9, 18, 27, 36에 엑시트 헤드(exit heads)를 배치했습니다. 추론 시 사용자는 계산 제약 조건에 따라 엑시트 포인트를 선택할 수 있습니다 (예: 속도를 위해 레이어 4에서 중단하거나, 최대 정확도를 위해 레이어 36까지 사용).
• 자기 증류 손실(Self-Distillation Loss): 학습 목적 함수는 선택된 모든 엑시트 레이어의 손실을 집계합니다. 상위 레이어가 하위 레이어를 가이드하여, 하위 레이어가 견고한 표현을 학습하도록 유도합니다. 전체 손실은 가중 합계로 계산됩니다: $L = \sum_{i \in \iota} L_i \cdot \alpha_i$, 여기서 $\alpha_i$는 레이어 깊이에 따라 증가하여 더 깊은 레이어에 우선순위를 둡니다.
• 사전 학습 목적 함수:
  • 마스크드 언어 모델링 (MLM): 표준적인 토큰 예측입니다.
  • 인컨텍스트 분류 (In-Context Classification, ICC): 전통적인 다음 문장 예측(NSP) 손실을 대체합니다. ICC는 연결된 코드 스니펫들이 동일한 저장소(repository)에서 유래했는지 여부를 분류합니다. 이는 입력 밀도를 높여 패딩을 줄이고, 컨텍스트 윈도우 활용도를 확장하며, 의미론적 검색에 필수적인 청크 수준(chunk-level)의 이해를 촉진합니다.
• 컨텍스트 강화: 모델은 Grouped Query Attention (GQA), Rotary Positional Encoding (RoPE), 그리고 FlashAttention V2를 활용합니다. 양방향 어텐션을 가능하게 하기 위해 StarCoder-2의 인과적 마스킹(causal masking)을 제거하였으며, 수용 영역(receptive field)의 제한을 피하기 위해 슬라이딩 윈도우 어텐션을 전체 어텐션(full attention)으로 교체하였습니다.

데이터 및 미세 조정(Fine-tuning)

• SYNTHCoNL 데이터셋: 자연어-코드-코드 삼중항(triplet) 107만 개로 구성된 새로운 데이터셋입니다. 이 데이터셋은 CodeSearchNet의 코드 스니펫을 Qwen2.5-Coder-7B-Instruct를 사용하여 여러 언어(Python, Java, Go, PHP 등)로 번역하여 구축되었습니다. 이를 통해 동일한 자연어 설명을 공유하는 교차 언어 코드-투-코드 쌍을 생성함으로써, 텍스트-투-코드 벤치마크를 확장했습니다.
• 미세 조정: MOSE는 멀티 레이어 프로젝션 헤드를 갖춘 CLIP 스타일의 손실 함수를 사용하여 텍스트-투-코드 및 코드-투-코드 검색을 위해 SYNTHCoNL에서 미세 조정되었습니다. 또한 코드 클론 탐지를 위해 BigCloneBench에서도 미세 조정되었습니다.

3. 주요 기여

1. 자기 증류 방법론: 단일 레이어 스택 내에서 여러 모델을 학습시키는 참신한 접근 방식으로, 별도의 티처/스튜던트 모델이 필요하지 않으며 중복성을 제거합니다.
2. MOSE 아키텍처: 5개의 엑시트 포인트를 가진 사전 학습 및 미세 조정된 10억 파라미터 인코더로, 사용자가 메모리 및 계산 제약에 따라 모델 크기를 동적으로 선택할 수 있게 하여 경쟁력 있는 성능을 유지합니다.
3. SYNTHCoNL 데이터셋: 코드 번역을 통해 생성된 삼중항 1,071,367개를 포함하는 새로운 벤치마크 리소스로, 교차 언어 코드-투-코드 검색 연구를 용이하게 합니다.

4. 결과

평가는 CodeXGLUE 벤치마크 제품군(CodeSearchNet, Code Translation, POJ104)과 BigCloneBench를 대상으로 수행되었습니다.

• 텍스트-투-코드 검색 (CodeSearchNet): MOSE는 평균 MRR(Mean Reciprocal Rank) 81.0을 달와 달성하여, CodeT5+ (770M), UniXcoder, ModernBERT-Large를 포함한 모든 오픈 소스 모델을 능가했습니다. 또한 폐쇄형 모델인 OpenAI Text-Embedding-3-Large와의 성능 격차를 좁혔습니다.
• 코드-투-코드 검색:
  • POJ104 데이터셋(언어 내 의미론적 검색)에서 MOSE는 제로샷 설정에서 75.9의 mAP(mean Average Precision)를 기록하며 오픈 소스 모델 중 가장 높은 성적을 거두었습니다.
  • Code Translation (CT) 데이터셋에서 MOSE는 OpenAI의 임베딩 모델과 대등한 성능을 보였으며, 다른 오픈 모델들을 앞질렀습니다.
• 코드 클론 탐지 (BigCloneBench): MOSE는 레이어 18에서 94.2의 F1 점수를 기록하여, 더 깊은 레이어의 성능과 일치하며 CodeT5+ (95.1) 및 UniXcoder (95.2)와 경쟁할 만한 수준을 보여주었습니다.
• 효율성-정확도 트레이드오프: 본 논문은 모델을 36개 레이어에서 4개 레이어로 줄였을 때, 텍스트-투-코드 검색에서 MRR의 절대적 하락은 단 **6.4%**에 불과하면서도 부동 소수점 연산량(GFLOPs)은 90% 감소함을 입증했습니다. 특히 코드-투-코드 태스크의 경우, 초기 레이어(예: 레이어 4)에서 최적의 성능이 관찰되는 경우가 많아, 더 깊은 레이어가 항상 우월하다는 가설에 의문을 제기했습니다.

5. 의의 및 주장

본 논문은 MOSE가 추론 비용과 정확도 사이의 트레이드오프를 다루는 원칙적인 접근 방식을 제공한다고 주장합니다. 자기 증류를 활용함으로써 모델은 중간 표현을 강화하며, 이를 통해 전통적인 증류 과정 없이도 초기 레이어를 배포 가능한 수준으로 만듭니다.

• 배포 유연성: 이 아키텍처는 사용자가 지연 시간 및 에너지 제약에 가장 적합한 엑시트 포인트를 선택할 수 있는 자연스러운 모델 선택 메커니즘을 제공합니다.
• 표현의 분기(Representation Divergence): 결과는 최적의 레이어가 입력 모달리티(텍스트-투-코드 vs 코드-투-코드)에 따라 달라진다는 점을 강조하며, 초기 레이어가 코드-투-코드 태스크에서 더 나은 성능을 보이는 경우가 많음을 보여줍니다.
• 효율성: 이 접근 방식은 성능의 큰 저하 없이 더 작고 에너지 효율적인 모델 구성을 가능하게 함으로써 지속 가능한 배포를 지원합니다.

저자들은 광범위한 하이퍼파라미터 튜닝을 제한한 계산적 제약과, 미세 조정 시 번역을 통해 생성된 합성 코드를 사용한 것에 따른 영향이 불분명하다는 점을 포함한 한계점을 인정했습니다. 향후 연구에서는 태스크 유형과 레이어 깊이 간의 상호작용을 더욱 탐구할 것을 제안했습니다.