기술 요약: 베이지안 모델 병합 (BMM)

문제 정의

모델 병합은 여러 작업별 전문가 모델을 공동 재학습 없이 또는 원래 학습 데이터에 접근하지 않고 단일 아키텍처로 결합하는 것을 목표로 합니다. 이 패러다임은 별도의 전문가를 유지하는 데 따른 저장 및 배포 오버헤드와 데이터 실로 (data silos) 나 프라이버시 제약 하에서 다중 작업 학습의 prohibitive 한 비용을 해결합니다.

기존 병합 방법들은 두 가지 주요 한계에 직면해 있습니다:

1. 유도 편향 (Inductive Bias) 의 미활용: 대부분의 방법은 처음부터 병합된 가중치를 추정하여, 강력한 "앵커" 모델 (예: 기존 병합 솔루션인 Task Arithmetic 또는 TSV) 에 내재된 귀중한 사전 지식을 활용하지 못합니다.
2. 전역 최적화의 부재: 고급 방법들은 종종 모듈 단위로 고립되어 작동하거나 이질적인 네트워크 모듈 간에 공유되는 하이퍼파라미터에 의존하며, 종단 간 성능을 극대화하기 위해 정규화 강도를 전역적으로 조정하는 전략이 부족합니다.

또한, 많은 방법들은 보조 보정 데이터 (data-assisted) 를 필요로 하는데, 이는 프라이버시나 저장 공간 제약으로 인해 종종 사용할 수 없으며, 데이터 프리 (data-free) 방법들은 활성화 통계를 효과적으로 근사하기 위한 이론적 근거가 부족한 경우가 많습니다.

방법론

본 논문은 플러그 앤 플레이 (plug-and-play) 형태의 이계 (bi-level) 최적화 프레임워크인 **베이지안 모델 병합 (BMM)**을 소개합니다.

1. 활성화 기반 베이지안 회귀 (내부 루프)

BMM 은 모듈 단위의 병합을 베이지안 선형 회귀 문제로 공식화합니다.

• 사전 유도: 처음부터 가중치를 추정하는 대신, BMM 은 강력한 앵커 모델 (예: TA, TIES, TSV) 을 사용하여 사전 분포를 정의합니다. 병합된 작업 벡터 $U$는 요소별 독립 가우시안 사전 분포를 통해 앵커 작업 벡터 $U^{(0)}$에 의해 정규화됩니다.
• 가능도 (Likelihood): 데이터 보조 설정에서 가능도는 작은 보정 세트에서 수집된 경험적 활성화 통계량 ($X$) 과 잔차 출력 ($Y$) 에서 유도됩니다.
• 폐쇄형 해 (Closed-Form Solution): 병합된 작업 벡터의 최대 사후 확률 (MAP) 추정은 분석적으로 유도됩니다:
  $$U_{MAP} = (YX^\top + \lambda U^{(0)})(XX^\top + \lambda I)^{-1}$$
  여기서 $\lambda$는 경험적 적합도와 앵커 사전 분포 사이의 균형을 맞추는 정규화 가중치입니다. 이는 모듈 단위 병합을 위한 효율적인 폐쇄형 해를 제공합니다.

2. 전역 조정을 위한 베이지안 최적화 (외부 루프)

네트워크 모듈의 이질성을 해결하기 위해, BMM 은 단일 공유 하이퍼파라미터를 사용하지 않습니다. 대신 **베이지안 최적화 (BO)**를 사용하여 모듈별 정규화 강도 ($\lambda$) 를 전역적으로 탐색합니다.

• 목적: 외부 루프는 하이퍼파라미터 $\lambda$에 대한 예측 증거 (검증 세트 성능으로 근사됨) 를 최대화합니다.
• 효율성: 가우시안 프로세스 (GP) 대리 모델이 탐색을 안내하여 탐색과 활용 사이의 균형을 맞춥니다. 대규모 모델 (예: ViT-L/14) 에서 계산 복잡성을 관리하기 위해, 저자들은 기능별 (예: attention-in, MLP-out) 로 모듈을 그룹화하여 탐색 공간의 차원을 줄이는 블록 단위 파라미터 결속 (block-wise parameter tying) 전략을 제안합니다.

3. 이론적 정렬을 통한 데이터 프리 변형

보조 보정 데이터 없이 병합을 가능하게 하기 위해, 저자들은 활성화 통계량과 작업 벡터 간의 이론적 정렬에 기반한 데이터 프리 변형을 유도했습니다.

• 정리 1 (Theorem 1): 미세 조정의 말기 단계 (신경 붕괴 및 표현 기하학) 에 관한 약간의 가정 하에서, 입력 활성화의 그람 행렬 ($E[xx^\top]$) 은 작업 벡터의 그람 행렬 ($U^\top U$) 과 양의 상관관계를 가집니다.
• 근사: 이 정렬을 통해 경험적 활성화 통계량을 전문가 가중치만으로 유도된 대리치로 대체할 수 있습니다: $XX^\top \approx c \sum (U^{(t)})^\top U^{(t)}$.
• 결과: 이는 보정 데이터가 필요하지 않고 전문가 및 앵커의 작업 벡터만 필요한 데이터 프리 MAP 추정을 산출하며, 폐쇄형 해를 유지합니다.

주요 기여

1. 이계 최적화 프레임워크: BMM 은 앵커 사전 분포를 활용하는 폐쇄형 내부 루프 MAP 추정기와 전역 하이퍼파라미터 조정을 위한 외부 루프 베이지안 최적화를 결합합니다.
2. 이론적 데이터 프리 유도: 본 논문은 활성화 통계량과 작업 벡터 간의 이론적 연결 (정리 1) 을 확립하여, 폐쇄형 해를 유지하는 데이터 프리 베이지안 병합 변형을 가능하게 합니다.
3. 플러그 앤 플레이 유연성: 이 프레임워크는 모듈화되어 있어 기존 병합 솔루션을 앵커 사전 분포로 사용할 수 있으며, 데이터 보조 및 데이터 프리 설정 모두를 지원합니다.

실험 결과

저자들은 비전 (ViT-B/32, ViT-L/14) 및 언어 (Llama-3.2-3B, Llama-3.1-8B) 벤치마크에서 8 개, 14 개, 20 개 작업 병합 시나리오에 걸쳐 BMM 을 평가했습니다.

• 성능 향상: BMM 은 데이터 보조 및 데이터 프리 설정 모두에서 모든 플러그 앤 플레이 앵커 베이스라인 (TA, TIES, RegMean, TSV, WUDI-Merging, ISO-CTS) 보다 일관되게 우수한 성능을 보였습니다.
  • ViT-L/14 8 개 작업 벤치마크에서 단일 BMM 병합 모델은 **95.1%**의 정확도 (데이터 보조) 와 95.0% (데이터 프리) 를 달성하여, 여덟 개의 개별 작업별 전문가의 평균 성능 (95.8%) 과 근접했습니다.
  • 약한 앵커 (예: TA, TIES) 대비 상대적 개선폭은 최대 **27%**에 달했습니다.
• 확장성: BMM 은 병합된 작업 수가 증가함에 따라 (최대 20 개 작업까지) 효과성을 유지하며 작업 간 간섭을 성공적으로 완화했습니다.
• 효율성: BO 탐색에도 불구하고 BMM 은 계산적으로 경쟁력 있게 유지되었습니다. 예를 들어, ViT-B/32 에서 20 회 BO 시도만으로도 225 구성을 요구하는 포괄적인 그리드 검색이 필요한 베이스라인보다 우수한 성능을 내면서도 실제 소요 시간을 크게 단축했습니다.
• 애블레이션 연구:
  • 전역 조정: BO 를 공유-$\lambda$ 또는 무작위 탐색으로 대체하면 성능이 저하되어, 모듈별 전역 최적화의 필요성을 확인했습니다.
  • 데이터 프리 유효성: 데이터 프리 변형은 데이터 보조 변형과 유사한 성능을 보여, 활성화 통계량과 작업 벡터 간의 이론적 정렬을 검증했습니다.
  • 불확실성 보정: 전체 가우시안 사후 분포를 사용하는 샘플링 기반 BMM 변형은 정확도를 희생하지 않으면서 포인트 추정치에 비해 기대 보정 오차 (ECE) 가 개선됨을 보여주었습니다.

중요성과 주장

본 논문은 보조 보정 데이터가 제한되거나 사용할 수 없는 상황에서 모델 병합을 위한 실용적이고 확장 가능한 경로를 BMM 이 제공한다고 주장합니다. 정보력 있는 앵커 사전 분포와 전역 하이퍼파라미터 조정을 결합함으로써, BMM 은 개별 작업별 전문가의 성능과 통합된 다중 작업 모델 간의 격차를 해소합니다.

저자들은 그들의 접근 방식이 공동 재학습이나 원래 학습 데이터에 대한 접근을 요구하지 않으므로, 데이터 실로 나 프라이버시 제약이 있는 환경에 적합하다고 강조합니다. 데이터 프리 변형의 유도는 작업 벡터의 기하학적 구조가 경험적 활성화 통계를 효과적으로 대체할 수 있음을 보여주는 핵심 이론적 기여로 강조됩니다. 결과는 제로샷 조합으로 병합을 취급하는 대신 기존 병합 솔루션을 사전 분포로 활용하는 것이 최첨단 성능을 달성하는 데 결정적인 요소임을 시사합니다.