기술 요약: Non-IID FedAvg에서 보존되었으나 정렬이 어긋난 표현(Representations)에 대한 기계론적 증거

1. 문제 정의

연합 평균화(Federated Averaging, FedAvg)는 연합 학습(Federated Learning, FL)의 표준 최적화 절차이지만, 특히 레이블 스큐(label skew)와 같은 non-IID(통계적 이질성) 클라이언트 데이터 환경에서 심각한 성능 저하를 겪습니다. 기존의 설명들은 이러한 실패를 클라이언트 드리프트(client drift), 목적 함수 불일치(objective inconsistency), 가중치 발산(weight divergence)의 관점에서 다루며 최적화 과정을 특징짓지만, 집합(aggregation)의 내부적 효과는 거의 관찰되지 않은 채로 남아 있습니다.

핵심적인 미해결 질문은 다음과 같습니다: 글로벌 정확도의 저하가 클라이언트가 학습한 표현의 **삭제(erasure)**를 반영하는 것인가, 아니면 내부 구조는 보존되었으나 최종 예측 경로와 효과적으로 연결되지 못하는 **정렬 불량(misalignment)**을 반영하는 것인가? 본 논문은 기계론적 해석 가능성(Mechanistic Interpretability, MI) 도구를 사용하여, non-IID FedAvg가 학습된 지식을 파괴하는지, 아니면 단지 분류기와의 연결을 끊어버리는지를 조사합니다.

2. 방법론

저자들은 CIFAR-10 및 Fashion-MNIST 데이터셋을 대상으로 희소하게 훈련된 비전 모델(SimpleCNN 및 ResNet10)을 분석합니다. 해석 가능성을 보장하기 위해, 클라이언트는 크기 기반 희소도 스케줄(99.9% 희소도로 점진적 감소) 하에서 훈련하며, 덴스 모델(dense-model) 대조군을 통해 연구 결과가 희소성으로 인한 인위적인 결과가 아님을 확인합니다. 본 연구는 네 가지 기계론적 진단 도구를 사용합니다:

2.1. 회로 발견 (Circuit Discovery)

미분 가능한 이진 게이팅(differentiable binary gating)을 사용하여 특정 클래스를 예측하는 데 충분한 최소 서브 네트워크(회로)를 식별합니다. 다음을 측정합니다:

• 클라이언트 내 안정성 (Intra-client stability): 통신 라운드 전반에 걸친 회로의 중첩도.
• 클라이언트 간 일관성 (Inter-client consistency): 서로 다른 클라이언트 간의 회로 유사도.
• 로컬-글로벌 보존 (Local-to-global preservation): 클라이언트의 로컬 회로와 집합된 글로벌 모델의 대응하는 회로 사이의 IoU(Intersection-over-Union).
• 중재 (Interventions): 충분성(Sufficiency) (회로만 활성화했을 때의 정확도) 및 필요성(Necessity) (회로를 제거했을 때의 정확도 저하).

2.2. 선형 프로빙 (Linear Probing)

글로벌 모델의 동결된 마지막 층 직전 활성화값(penultimate-layer activations)에 대해 선형 분류기를 훈련합니다. 이는 모델의 고유 헤드가 해당 정보를 활용하지 못하더라도, 백본(backbone) 내에 클래스 판별 정보가 선형적으로 접근 가능한 상태로 남아 있는지를 테스트합니다.

2.3. 헤드 전용 미세 조정 (Head-Only Finetuning)

백본을 동결하고, 소수의 균형 잡힌 IID 서브셋(200개 샘플)에 대해서만 최종 분류 헤드를 미세 조정합니다. 이를 통해 정확도 손실이 읽아내기 층(readout layer)에 국한된 것인지, 아니면 특징 백본(feature backbone)의 문제인지를 격리하여 확인합니다.

2.4. 범용 희소 오토인코더 (Universal Sparse Autoencoders, USAE)

IID 및 non-IID 글로벌 모델의 활성화값으로부터 공유된 특징 사전(feature dictionary)을 학습하기 위해 USAE를 공동으로 훈련합니다. 이는 두 모델이 공유된 기저(basis) 상에서 특징을 구성하는지 테스트하며, 교차 모델 재구성을 통해 특징의 호환성을 평가할 수 있게 합니다.

3. 주요 결과

3.1. 회로 안정성 및 발산

• 로컬 안정성: 데이터 이질성과 관계없이 로컬 훈련 역학은 안정적입니다. 클라이언트는 극단적인 스큐($\alpha=0.05$) 상황에서도 매우 안정적인 클래스 특이적 회로(IoU $\approx$ 0.90)로 수렴합니다.
• 클라이언트 간 발산: 이질성이 증가함에 따라 클라이언트는 동일한 클래스에 대해 서로 다른 회로를 학습합니다. 극단적인 스큐 하에서 클라이언트 간 회로 중첩도는 IoU $\approx$ 0.45로 떨어집니다.
• 부분적 보존: 글로벌 모델은 회로를 노이즈로 평균화하지 않습니다. 로컬-글로벌 보존(IoU $\approx$ 0.55)은 클라이언트 간 중첩보다 높게 유지되며, 이는 글로벌 모델이 단일 솔루션으로 수렴하기보다는 클라이언트로부터 부분적인 구조를 상속받음을 나타냅니다.

3.2. 기능적 회로는 정확도가 0임에도 지속됨

극단적인 non-IID 설정에서 글로벌 모델은 특정 클래스(예: "고양이" 또는 "사슴")에 대해 거의 0%에 가까운 정확도를 보일 수 있습니다. 그러나:

• 충분성: 추출된 해당 클래스의 회로는 격리되었을 때 거의 완벽한 충분성(정확도 100%)을 달성합니다.
• 필요성: 이러한 회로를 제거(ablate)하면 정확도가 거의 0으로 급락합니다.
• 결론: 해당 클래스를 위한 계산 경로(computational pathways)는 존재하며 예측에 인과적 책임을 지고 있지만, 모델은 표준 추론 과정에서 입력을 해당 경로로 라우팅하는 데 실패합니다.

3.3. 선형 분리 가능성 및 헤드 불일치

• 선형 프로브: 스큐 하에서 글로벌 정확도가 크게 떨어지더라도(예: 68.3%에서 43.1%로), 선형 프로브 정확도는 높고 안정적으로 유지됩니다(64.7%로 하락). 프로브와 글로벌 정확도 사이의 격차는 이질성이 증가함에 따라 단조롭게 넓어집니다 (+$\alpha=0.05$에서 21.7 포인트).
• 헤드 미세 조정: 소량의 균형 잡힌 세트에 대해 분류 헤드만 미세 조정하면 상당한 정확도를 회복하지만(+$\alpha=0.05$에서 10.2 포인트), IID 성능을 완전히 복구하지는 못합니다. 이는 백본이 유용한 특징을 보유하고 있으나, 기본 헤드가 정렬 불량 상태임을 확인시켜 줍니다.

3.4. 공유 특징 기저 (USAE)

• 개념 중첩: 2,048개의 USAE 사전 개념 중 98.7%가 IID 및 non-IID 모델 모두에서 비자명하게 활성화됩니다. 이는 두 모델이 대체로 공유된 특징 기저를 사용함을 나타냅니다.
• 교차 모델 전이: non-IID 활성화값을 IID 특징 공간으로 디코딩하면 IID 헤드가 잘 작동합니다. 반대로, non-IID 헤드는 IID 활성화값을 입력받았을 때 실패합니다. 이는 특징은 존재하고 호환되지만, non-IID 읽기 메커니즘(readout mechanism)이 실패의 지점임을 확인해 줍니다.

4. 의의 및 주장

본 논문은 non-IID FedAvg의 성능 저하가 완전히 표현의 삭제(representational erasure)로 설명되지 않는다고 주장합니다. 대신, 주요 메커니즘은 **정렬 불량(misalignment)**입니다:

1. 보존: 클라이언트가 학습한 내부 구조(회로)와 특징 기저는 집합된 글로벌 모델에서도 대체로 보존됩니다.
2. 정렬 불량: 최종 예측 경로(분류 헤드)가 이러한 보존된 표현들과 정렬이 어긋납니다. 모델은 클래스를 "알고" 있지만, 입력을 올바른 회로로 라우팅하거나 특징을 효과적으로 활용하지 못합니다.

시사점:

• 문제의 재정의: 문제는 반드시 클라이언트가 "망각"하거나 가중치가 지식을 파괴하는 호환되지 않는 영역으로 발산하는 것이 아니라, 집합 과정이 백본과 헤드 사이의 구조적 관계를 방해하는 것입니다.
• 알고리즘 방향:
  • 집합 후 헤드 재정렬 (Post-aggregation Head Realignment): 연합 프로토콜을 변경하지 않고도 소량의 균형 잡힌 데이터로 헤드를 미세 조정함으로써 상당한 성능을 회복할 수 있습니다.
  • 회로 인식 집합 (Circuit-Aware Aggregation): 향후 방법론은 (예: 순열 대칭성 방법을 통해) 클라이언트 간에 유사한 기능적 회로를 구현하는 뉴런을 집합하기 전에 정렬함으로써, 근원지에서의 구조적 변화를 줄일 수 있습니다.

한계점:
저자들은 헤드 재정렬이 상당한 성능을 회로하지만 IID 성능과의 잔여 격차가 남아 있음을 언급하며, 이는 백본의 기하학적 구조 또한 다소 저하됨을 시사합니다. 현재 연구 결과는 이미지 데이터셋에 대한 CNN/ResNet 아키텍처로 제한되어 있으며, 분석은 서브 채널 단위가 아닌 채널 수준에서 수행되었습니다.