기술 요약: 연방 중첩 학습 (FedNL)

문제 제기

대규모 언어 모델 (LLM) 을 위한 연방 학습 (FL) 은 **데이터 이질성 (Non-IID)**과 롱테일 분포라는 두 가지 지속적이고 밀접하게 얽힌 과제에 직면해 있습니다. 현실적인 시나리오에서 클라이언트 데이터 분포는 매우 편향되어 있습니다 (예: 의료 데이터 대 코딩 데이터). 또한, 중요한 지식은 종종 롱테일에 존재하며, 전역 집계 과정에서 헤드 클래스에 의해 쉽게 가려집니다.

기존 접근법은 주로 정적 모델 가중치를 정규화하거나 개인화하는 데 중점을 둡니다 (예: FedProx, FedSSI). 이러한 방법들은 치명적인 망각을 완화하는 데 효과적이지만, 전역 모델을 정적 지식의 컨테이너로 간주합니다. 그 결과, 보지 못한 매우 이질적인 데이터를 가진 클라이언트에 배포될 때, 모델은 경사 하강 업데이트 없이 현재 맥락에 적응할 수 있는 **테스트 시간 가소성 (test-time plasticity)**이 부족합니다. 이는 맥락에 따라 미묘한 차이가 결정적인 도메인별 작업에서 최적이지 않은 성능으로 이어집니다.

방법론: 연방 중첩 학습 (FedNL)

본 논문은 정적 지식의 집계를 넘어 최적화 규칙을 협력적으로 학습하는 프레임워크인 **연방 중첩 학습 (FedNL)**을 제안합니다. 중첩 학습 (NL) 이론에서 영감을 받은 FedNL 은 추론을 수동적인 검색 과정이 아닌, 모델이 맥락을 일시적인 메모리 상태로 능동적으로 압축하는 고주파수 '내부 루프' 훈련 과정으로 재해석합니다.

핵심 아키텍처

FedNL 은 FL 을 3 단계 중첩 최적화 시스템으로 재형식화합니다:

1. 레벨 0 (외부 루프 - 서버): 협력적 일반화. 서버는 메모리 내용이 아닌 메타 규칙 (매개변수 $\theta$) 을 집계합니다. 이러한 규칙은 메모리가 어떻게 구성되고 업데이트되는지를 지배합니다.
2. 레벨 1 (중간 루프 - 클라이언트): 규칙 학습. 로컬 훈련 동안 클라이언트는 내부 루프가 하류 작업에 유용한 메모리 상태를 생성하도록 $\theta$(예: LoRA 어댑터, 게이트 메커니즘) 를 최적화합니다. 여기에는 메모리 궤적을 통한 역전파가 포함됩니다.
3. 레벨 2 (내부 루프 - 클라이언트): 테스트 시간 적응. 추론 중 클라이언트는 글로벌 규칙 $\theta$를 사용하여 **델타 규칙 (Delta Rule)**을 통해 개인적이고 맥락을 인식하는 메모리 상태 $S_t$를 온라인으로 업데이트합니다. 이 상태는 일시적이며 절대 장치 밖으로 나가지 않습니다.

기술적 구현

• Titans 기반 선형 어텐션: FedNL 은 선형화된 어텐션 메커니즘을 활용하는 Titans 아키텍처를 활용합니다. 메모리 상태 $S_t$는 델타 규칙을 통해 업데이트됩니다:
  $$S_t = S_{t-1} + \eta(v_t - S_{t-1}k_t)k_t^\top$$
  이 업데이트는 연관성 메모리 목적 함수에 대한 경사 하강의 단일 단계와 수학적으로 동일하며, 추론을 온라인 최적화 과정으로 효과적으로 처리합니다.
• 지식과 규칙의 분리: 모델은 일반 언어 능력을 제공하는 **동결된 백본 (frozen backbone)**과 학습 가능한 메타 매개변수 ($\theta$) 로 구성됩니다. 메타 매개변수에는 쿼리/키/값 생성을 위한 저랭크 프로젝션 (LoRA) 과 정적 및 동적 어텐션을 균형 있게 조절하는 메모리 게이팅 메커니즘 ($\alpha$) 이 포함됩니다.
• 효율성: 전체 모델 가중치나 메모리 상태가 아닌 소규모 메타 규칙 (어댑터) 만을 집계함으로써 FedNL 은 통신 오버헤드를 극적으로 줄입니다. 또한, 선형 어텐션 메커니즘은 시퀀스 길이에 관계없이 일정한 추론 메모리 복잡도 ($O(1)$) 를 유지하며, $O(T)$ KV 캐시가 필요한 표준 어텐션과 대조됩니다.

주요 기여

1. 이론적 재형식화: 본 논문은 FL 에 중첩 학습 관점을 도입하여, 정적 표현이 아닌 최적화 규칙의 협력적 훈련으로 문제를 공식화합니다. 이는 테스트 시간 적응을 통해 Non-IID 문제를 해결하기 위한 이론적 기반을 제공합니다.
2. FedNL 프레임워크: Titans 기반 선형 어텐션을 FL 파이프라인에 통합한 실용적인 알고리즘입니다. 델타 규칙을 온라인 경사 하강 단계로 간주함으로써, 무거운 계산 작업 없이 보지 못한 도메인에서의 제로샷 테스트 시간 적응을 가능하게 합니다.
3. 실증적 우위성: 실험 결과에 따르면 FedNL 은 짧은 컨텍스트 추론에서 경쟁력 있는 성능을 달성하고, 강력한 연방 베이스라인 대비 긴 컨텍스트 검색 및 스트리밍 진단에서 우수한 향상을 보였습니다.

실험 결과

저자들은 두 가지 연방 설정에서 FedNL 을 평가했습니다: 5 클라이언트 Non-IID MMLU 분할과 7 클라이언트 Non-IID Needle In A Haystack (NIAH) 긴 컨텍스트 벤치마크입니다.

• Non-IID MMLU (짧은 컨텍스트 추론):

  • Qwen2.5-1.5B 백본에서 FedNL 은 최고 평균 정확도 (58.88%) 를 달성하여 최상의 베이스라인인 FedSSI(58.70%) 를 약간 앞섰습니다.
  • STEM(+2.0 pp) 및 수학/컴퓨터 과학 (+3.8 pp) 도메인에서 유의미한 향상이 관찰되었습니다.
  • Llama-3.2-1B 백본에서 FedNL 은 **42.64%**를 기록하여 FedSSI 의 42.10% 를 능가했습니다.
  • 특히 FedNL 은 로컬 모델에서 연방 모델로 이동할 때 발생하는 '집계 감소 (aggregation drop)'를 최소화하여 도메인 전반에 걸쳐 거의 0 에 가깝게 유지한 반면, 정적 베이스라인은 상당한 성능 저하를 보였습니다.

• 긴 컨텍스트 검색 (NIAH):

  • FedNL 은 7 가지 검색 템플릿에서 평균 정확도 **29.7%**를 달성하여 FedALA(28.6%) 를 능가했습니다.
  • MK-NIAH, MV-NIAH, UUID 코드 작업에서 향상이 가장 두드러졌습니다.
  • 스트리밍 교차 엔트로피 (CE) 진단: 16K 토큰 스트리밍 프로브에서 FedAvg 의 정규화된 CE 는 컨텍스트가 전개됨에 따라 3.1% 증가하여 (불확실성 누적) 불안정성을 보인 반면, FedNL 은 2.1% 감소하여 재귀적 메모리 상태가 긴 스트림을 통해 유용한 컨텍스트를 성공적으로 흡수했음을 나타냈습니다.
  • 메모리 효율성: FedNL 은 일정한 추론 메모리를 유지하여 정적 어텐션 베이스라인에 영향을 미쳤던 16K 토큰에서의 메모리 부족 (OOM) 오류를 방지했습니다.

• 통신 효율성:

  • FedNL 은 테스트된 설정에서 학습 가능한 매개변수에 대해 FedAvg 대비 라운드당 통신을 약 350 배 줄였습니다 (0.26 MB 대 22.5 MB).

중요성 및 주장

본 논문은 FedNL 이 정적 지식의 저장소에서 지속적이고 맥락을 인식하는 학습의 패러다임으로 진화하는 연방 학습의 새로운 방향을 확립한다고 주장합니다.

• 패러다임 전환: 일반 언어 능력 (동결된 백본) 과 메모리 구성 규칙 (학습 가능한 어댑터) 을 분리함으로써, FedNL 은 다양한 클라이언트가 상충되는 '지식 상태'가 아닌 '학습 능력'을 공유할 수 있게 합니다.
• 이질성에 대한 강건성: 글로벌 메타 매개변수는 서로 다른 도메인 (예: 코드 대 의료) 간의 충돌을 인코딩할 필요가 없습니다. 대신, "맥락이 X 를 포함하면 메모리를 Y 로 저장하도록 업데이트한다"는 보편적인 규칙을 인코딩합니다. 이는 매개변수 공간에서의 Non-IID 충돌을 완화합니다.
• 제로샷 적응: 이 프레임워크는 모델이 테스트 시간 훈련 (TTT) 방법이 종종 직면하는 '콜드 스타트' 문제를 해결하기 위해, 학습된 메모리 업데이트 규칙을 실행함으로써 추론 초기 몇 토큰 동안 특정 로컬 분포에 적응할 수 있게 합니다.

저자들은 1B~1.5B 규모에서 결과가 일관되지만, 메모리 규칙 집계의 일반성을 검증하기 위해 FedNL 을 더 큰 파운데이션 모델로 확장하는 것이 필수적인 다음 단계라고 지적합니다.