모델 학습 가능성 인식 데이터 할당 기반 연합 추론 증류 프레임워크

Federated Reasoning Distillation Framework with Model Learnability-Aware Data Allocation

데이터 할당은 연합 대형 언어 모델(LLM) 및 소형 언어 모델(SLM)의 추론 협업에서 중요한 역할을 한다. 그럼에도 불구하고 기존의 데이터 할당 방법들은 협업 시 충분히 탐구되지 않은 문제인 양방향 모델 학습 가능성 격차(bidirectional model learnability gap)를 해결하지 못한다. 이로 인해 클라이언트 측 SLM은 LLM으로부터의 효과적인 지식 전이를 위해 자신의 학습 제약에 맞는 고보상 샘플을 식별하지 못하며, 반대로 LLM은 기존 데이터를 넘어선 새로운 지식을 제공하는 샘플을 선택하는 데 어려움을 겪는다. 더욱이 이러한 협업 프레임워크들은 도메인 불가지론적 추론 전이(domain-agnostic reasoning transfer)라는 또 다른 주요 과제에 직면해 있다. 기존 추론 전이 방법은 로컬 도메인 데이터에 유연하게 적응하지 못하여, SLM이 범용 LLM으로부터 단계적(step-by-step) 추론 능력을 효과적으로 습득하는 것을 방해한다. 이러한 문제를 해결하기 위해, 본 논문에서는 모델 학습 가능성을 인식하는 데이터 할당을 적용한 연합 추론 증류 프레임워크인 LaDa를 제안한다. LaDa는 각 SLM과 LLM 쌍 간의 학습 가능성 격차를 기반으로 고보상 샘플을 적응적으로 할당하는 모델 학습 가능성 인식 데이터 필터를 도입하여 양방향 지식 전이를 효과적으로 촉진한다. 나아가, SLM과 LLM 간의 대조 증류 학습(contrastive distillation learning)을 통해 필터링된 고보상 샘플 상에서 추론 경로의 결합 확률을 정렬하는 도메인 적응형 추론 증류 기법을 설계함으로써, SLM이 로컬 데이터 분포 하에 내재된 추론 패턴을 포착할 수 있도록 한다. LaDa는 기존 협업 프레임워크의 플러그인 모듈로 작동하며, 모델 학습 가능성 격차에 따라 지식 전이를 적응시킨다.

Data allocation plays a critical role in federated large language model (LLM) and small language models (SLMs) reasoning collaboration. Nevertheless, existing data allocation methods fail to address an under-explored challenge in collaboration: bidirectional model learnability gap, where client-side SLMs cannot identify high-reward samples matching their learnability constraints for effective knowledge transfer from LLMs, while LLMs struggle to select samples contributing novel knowledge beyond their existing data. Furthermore, these collaboration frameworks face another key challenge: domain-agnostic reasoning transfer, where existing reasoning transfer methods fail to flexibly adapt to the local domain data, preventing SLMs from effectively acquiring step-by-step reasoning abilities within from general LLM. To address these challenges, we propose LaDa, a federated reasoning distillation framework with model learnability-aware data allocation. It introduces a model learnability-aware data filter that adaptively allocates high-reward samples based on the learnability gap between each SLM and LLM pair, effectively facilitating bidirectional knowledge transfer. We further design a domain adaptive reasoning distillation method that aligns joint probabilities of reasoning paths on filtered high-reward samples through contrastive distillation learning between SLM and LLM, enabling SLM to capture underlying reasoning patterns under local data distribution. LaDa operates as a plug-in module for existing collaboration frameworks, adapting knowledge transfer based on model learnability gaps.