기술 요약: PEP-FedPT

문제 정의

연방 학습 (FL) 은 데이터 이질성 하에서 대규모 사전 학습된 비전 트랜스포머 (ViT) 를 하위 작업에 적응시킬 때 상당한 도전에 직면합니다. 시각적 프롬프트 튜닝 (VPT) 은 미세 조정을 위한 매개변수 효율적인 솔루션을 제공하지만, 기존 접근 방식은 일반화와 개인화 사이의 균형을 맞추는 데 어려움을 겪습니다.

• 전역 프롬프트 튜닝: FedPR 와 같은 방법은 모든 클라이언트 간에 공유된 프롬프트를 사용합니다. 이러한 방법들은 일반화 성능이 뛰어나지만, 매우 이질적인 클라이언트 데이터 분포 (비 IID) 에 적응하지 못하여 '클라이언트 드리프트'와 비최적의 성능을 초래합니다.
• 개인화된 프롬프트 튜닝: pFedPG 와 같은 방법은 클라이언트별 프롬프트를 생성합니다. 이러한 방법들은 로컬 데이터에 잘 적응하지만, 종종 과적합되어 보지 못한 클라이언트나 참여하지 않는 클라이언트에 대한 일반화 성능이 떨어집니다. 또한, 많은 개인화 접근 방식은 매 라운드마다 모든 클라이언트의 참여를 가정하는데, 이는 실제 FL 시스템에서는 비현실적입니다.
• 기존 하이브리드 접근 방식: SGPT 와 같은 방법은 그룹별 프롬프트를 사용하여 이 격차를 해소하려 시도하지만, 비미분 가능한 선택 메커니즘과 2 단계 학습 과정을 통해 최적화 복잡성을 도입합니다.

해결하려는 핵심 질문은 다음과 같습니다: 오직 전역적으로 공유된 프롬프트에만 의존하면서도 효과적인 개인화를 달성할 수 있을까요?

방법론: PEP-FedPT

저자들은 클라이언트 의존적 학습 가능 매개변수를 저장하지 않고 개인화를 달성하는 통합 프레임워크인 PEP-FedPT(연방 프롬프트 튜닝을 위한 프로토타입 기반 프롬프트 추정) 를 제안합니다. 핵심 혁신은 **클래스-맥락 혼합 프롬프트 (CCMP)**입니다.

1. 프롬프트 아키텍처

이 프레임워크는 모든 클라이언트 간에 전역적으로 공유되는 두 가지 학습 가능 매개변수 세트를 활용합니다.

• 공유 프롬프트 ($P_S$): ViT 의 가장 첫 번째 레이어에 삽입됩니다. 이는 매니폴드 전반에 공통적으로 존재하는 저수준의 클래스 무관 표현을 포착합니다.
• 클래스별 프롬프트 ($P_C$): 각 클래스에 해당하는 프롬프트 집합으로, 전역적으로 공유됩니다.

2. 클래스-맥락 혼합 프롬프트 (CCMP)

단일 정적 프롬프트를 삽입하는 대신, PEP-FedPT 는 중간 레이어 (실험에서는 구체적으로 5, 6, 7 레이어) 에서 각 입력 샘플에 대해 동적으로 혼합 프롬프트 ($m$) 를 생성합니다. 이 혼합 프롬프트는 다음에서 도출된 점수로 가중치가 부여된 전역 클래스별 프롬프트의 부드러운 조합입니다.

• 전역 클래스 프로토타입 ($\mu$): 클라이언트 간 클래스 내 데이터 포인트의 CLS 토큰을 평균화하여 계산된 각 클래스에 대한 CLS 토큰 표현의 집계된 중심점입니다. 이러한 값은 모멘텀을 가진 연방 평균화를 통해 주기적으로 업데이트됩니다.
• 클라이언트 클래스 사전 확률 ($\delta$): 각 클라이언트의 로컬 데이터에 고유한 경험적 레이블 분포 (클래스 빈도) 입니다.

3. 프롬프트 혼합 메커니즘

클라이언트 $k$의 주어진 입력 $x$에 대해, 시스템은 클래스별 프롬프트에 대한 부드러운 가중치 벡터 $s$를 계산합니다.

1. 유사도 점수: $l$ 레이어의 입력 CLS 토큰 표현 ($cls_{l-1}$) 과 전역 클래스 프로토타입 ($\mu$) 간의 코사인 유사도가 계산됩니다.
2. 가중치 점수 산출: 이러한 유사도 값들은 클라이언트의 로컬 클래스 사전 확률 ($\delta$) 과 온도 매개변수 ($\tau$) 로 스케일링됩니다.
3. 정규화: 점수들은 확률 분포를 형성하도록 정규화됩니다.
4. 혼합: 최종 CCMP 는 가중치 합으로 계산됩니다: $m = P_C \times s$.

이를 통해 각 클라이언트는 로컬 데이터 분포와 현재 입력에 기반하여 프롬프트를 동적으로 개인화할 수 있으며, 이는 원시 데이터나 클라이언트별 모델 가중치를 전송하지 않고 전역 정보 (프롬프트 및 프로토타입) 와 로컬 통계 (사전 확률) 만을 사용하여 가능합니다.

4. 이론적 기반

본 논문은 CCMP 설계에 대한 두 가지 이론적 근거를 제시합니다.

• 2 차 상한 최소화: 저자들은 손실 함수가 Lipschitz 매끄럽고 프롬프트 공간이 컴팩트하다고 가정할 때, CCMP 가 전역 손실 함수의 2 차 상한을 최소화함을 증명합니다. 최적의 혼합 전략은 클래스 사전 확률에 기반한 클래스 프롬프트의 가중 평균과 일치합니다.
• MMSE 최적성: CCMP 는 관측된 CLS 토큰 표현이 주어졌을 때 실제 클래스 프롬프트의 최소 평균 제곱 오차 (MMSE) 추정기임이 입증되었습니다.

주요 기여

1. 통합 프레임워크: ViT 연방 학습의 데이터 이질성을 해결하기 위해 공유 프롬프트와 클래스별 프롬프트를 공동으로 최적화하는 PEP-FedPT 를 도입하여 일반화와 개인화 사이의 균형을 맞춥니다.
2. 새로운 혼합 전략: 로컬 클래스 사전 확률과 전역 프로토타입에 의해 조절되는 전역 클래스별 프롬프트의 함수로서 샘플별 프롬프트를 생성하는 CCMP 를 설계했습니다. 이는 로컬 프롬프트 저장이나 서버 측 클라이언트별 프롬프트 생성의 필요성을 제거합니다.
3. 이론적 통찰: CCMP 가 손실의 2 차 상한을 최소화하고 MMSE 관점에서 최적임을 입증했습니다.
4. 경험적 우위: 다양한 이질성 시나리오에서 최첨단 베이스라인을 능가하는 방법을 광범위한 검증을 통해 확인했습니다.

실험 결과

이 방법은 다양한 비 IID 설정 (병리적 분할 및 디리클레 분포) 하에서 CIFAR-100, TinyImageNet, DomainNet, iNaturalist에 대해 평가되었습니다.

• 성능: PEP-FedPT 는 평균 정확도와 최악의 로컬 정확도 모두에서 베이스라인 (FedVPT, FedPR, SGPT, pFedPG 포함) 을 일관되게 능가했습니다.
  • CIFAR-100(병리적) 에서 **95.46%**의 평균 정확도를 달성하여 pFedPG(92.96%) 와 FedVPT-D(85.15%) 를 상회했습니다.
  • TinyImageNet 에서 pFedPG 대비 평균 정확도를 8.59% 향상시켰습니다.
• 일반화: 훈련에 참여하지 않은 클라이언트에 대한 테스트 (홀드아웃 평가) 에서 PEP-FedPT 는 높은 정확도 (CIFAR-100 에서 93.71%) 를 유지한 반면, pFedPG 와 같은 개인화 방법들은 일반화에 실패했습니다 (NA 표시 또는 현저히 낮음).
• 효율성: 이 방법은 SGPT(13M) 및 FedVPT-D 에 비해 훨씬 적은 통신 라운드 (CIFAR-100 에서 83% 정확도 도달까지 12 라운드) 와 낮은 통신 오버헤드 (4.6M 매개변수 통신) 를 필요로 했습니다.
• 절대 실험: 실험을 통해 공유 프롬프트 단독 사용에 비해 공유 프롬프트와 CCMP 를 결합했을 때 상당한 이득을 얻는 것이 확인되었습니다. 클래스 사전 확률의 포함은 특히 병리적 설정에서 성능을 더욱 향상시켰습니다.

중요성과 주장

본 논문은 PEP-FedPT 가 비전 트랜스포머의 효율적이고 일반화 가능한 연방 프롬프트 튜닝을 위한 강력한 기반을 확립한다고 주장합니다. 그 주요 중요성은 FL 에서 개인화와 일반화 사이의 트레이드오프를 해결하는 데 있습니다.

• 로컬 데이터 분포 (사전 확률을 통해) 를 활용하여 프롬프트를 동적으로 혼합함으로써 개인화를 달성합니다.
• 완전히 개인화된 모델의 과적합 문제를 피하기 위해 전역적으로 공유된 프롬프트와 프로토타입에만 의존함으로써 일반화를 유지합니다.
• 클라이언트별 매개변수 저장이나 전체 클라이언트 참여를 요구하지 않으므로 리소스가 제한된 엣지 장치에 적합하여 확장성을 제공합니다.

저자들은 방법론이 클래스 사전 확률과 중심점에 대한 경험적 추정에 의존한다는 한계를 지적하며, 이는 반지도 또는 비지도 설정에서는 신뢰할 수 없을 수 있으며, 엄격한 IID 데이터 분포에서는 이 방법이 표준 Fed-VPT 로 축소되어 이점이 감소한다고 언급합니다. 그러나 비-IID 연방 학습의 맥락 내에서 제안된 접근 방식은 견고하고 이론적으로 근거 있으며 경험적으로 우월한 솔루션을 제공합니다.