기술 요약: 모든 클라이언트는 하나의 환경이다: 시공간 예측을 위한 연합 디컨파운딩(Federated De-confounding)

1. 문제 정의

시공간 예측(Spatio-temporal forecasting, STF)은 교통량 예측 및 대기 질 모니터링과 같은 도시 컴퓨팅 애플리케이션에서 매우 중요합니다. 딥러닝 모델은 높은 성능을 달 achievement 했지만, 일반적으로 중앙 집중식 데이터 수집을 필요로 하며, 이는 개인정보 보호 규정, 통신 제약 및 데이터 소유권 문제로 인해 불가능한 경우가 많습니다. 연합 학습(Federated Learning, FL)은 원시 데이터를 공유하지 않고 분산된 클라이언트를 통해 모델을 학습함으로써 이에 대한 해결책을 제공합니다.

그러나 기존의 연합 STF 방법들은 주로 **클라이언트 간 이질성(cross-client heterogeneity)**을 개인화된 모델을 통해 완화해야 할 최적화 과제로 취급합니다. 저자들은 이러한 관점이 결함이 있다고 주장하는데, 그 이유는 이질성이 근본적으로 각 클라이언트의 위치에 특화된 다양한 환경 조건(예: 날씨, 특수 이벤트, 출퇴근 패턴)에서 기인하기 때문입니다. 결과적으로, 로컬 모델은 클라이언트 특유의 컨파운딩 효과(confounding effects)와 얽힌 환경 특화 예측 패턴을 포착하게 됩니다. 이러한 모델은 환경 조건이 위치나 시간에 따라 변할 때 일반화에 실패합니다. 또한, 기존의 인과 추론 방법들은 환경적 컨파운더(environmental confounder)를 분리하는 데 필요한 중앙 집중식 원시 데이터에 접근할 수 없기 때문에 연합 환경에서의 STF에 어려움을 겪습니다.

2. 방법론: Atlas 프레임워크

본 논문은 Atlas라는 새로운 연합 디컨파운딩 프레임워크를 제안합니다. 핵심 통찰은 각 연합 클라이언트를 방해 요소가 아니라, 동일한 기저의 시공간 시스템에 대한 상호 보완적인 관측을 제공하는 별개의 **인과적 환경(causal environment)**으로 취급하는 것입니다.

2.1 인과적 정식화

저자들은 연합 STF를 위한 **구조적 인과 모델(Structural Causal Model, SCM)**을 정식화합니다:

• 변수: 과거 관측값($X_k$), 미래 지평($Y_k$), 그리고 클라이언트별 환경 조건($E_k$).
• 컨파운딩(Confounding): $E_k$는 백도어 경로(back-door path) $X_k \leftarrow E_k \rightarrow Y_k$를 통해 가짜 상관관계를 도입합니다.
• 로컬 조정의 한계: 백도어 조정을 통해 컨파운딩을 제거하려면 $P(Y|do(X)) = \sum P(Y|X, E=e)P(E=e)$를 추정해야 합니다. 연합 설정에서 개별 클라이언트는 편향된 환경 층위의 부분 집합($P_k(E)$)만을 관측하며 전역 마진 분포(global marginal distribution) $P(E)$가 결여되어 있습니다. 이는 두 가지 구체적인 편향을 초래합니다:
  1. 누락된 층위(Missing Strata): 클라이언트는 로컬 데이터에 존재하지 않는 환경 레짐(regime)을 놓칩니다.
  2. 왜곡된 가중치(Skewed Weights): 로컬 추정치는 전역 분포가 아닌 로컬 분포에 의해 가중치가 부여됩니다.

2.2 Atlas 메커니즘

Atlas는 연합 집계(federated aggregation)를 **암시적 백도어 조정(implicit back-door adjustment)**으로 활용하여 이러한 문제를 해결합니다.

1. 프로토타입 코드북(Prototype Codebook): 프레임워크는 공유된 전역 프로토타입 코드북 $Q = [q_1, \dots, q_I]$를 학습하며, 여기서 각 프로토타입 $q_i$는 잠재적 환경 층위에 대응합니다.
2. 로컬 인코딩: 각 클라이언트 $k$는 개인화된 백본을 사용하여 입력 $X_k$를 잠재 표현 $z_k$로 인코딩합니다. 그 다음 소프트 할당 메커니즘(식 5)을 통해 $z_k$를 코드북에 매핑하여 환경 조건이 반영된 표현 $\hat{z}_k$를 생성합니다. 이는 로컬 백도어 조정을 파라미터 수준에서 근사하는 역할을 합니다.
3. 정렬 및 집계(Alignment and Aggregation): 프로토타입 슬롯은 교환 가능하므로(예: 클라이언트 A의 슬로트 $i$가 클라이언트 B의 슬로트 $i$와 다른 레짐에 대응할 수 있음), 서버는 집계 전에 프로토타입 정렬을 수행합니다.
   • 서버는 클라이언트 프로토타입과 전역 앵커 프로토타입 사이의 코사인 거리에 기반한 비용 행렬을 계산합니다.
   • 클라이언트 슬롯을 전역 앵커 슬롯에 정렬하는 최적의 순열 $\Pi_k$를 찾기 위해 선형 할당 문제(헝가리안 알고리즘 사용)를 해결합니다.
   • 서버는 정렬된 코드북을 집계합니다: $Q \leftarrow \sum w_k \Pi_k Q_k$.
4. 이론적 보장: 저자들은 **연합 디컨파운딩 바운드(federated de-confounding bound)**를 도출하여, 오차가 개별 클라이언트의 강도가 아닌 클라이언트 간의 평균된 컨파운딩 강도에 의해 선형적으로 제어됨을 증명합니다. 이는 다양한 클라이언트 환경이 집계될 때 특정 편향을 상쇄함을 의미합니다.

2.3 구현

• 아키텍처: 클라이언트는 로컬 인코더(윈도우 레벨 컨볼루션 + ResMLP)와 디코더를 사용합니다. 전체 모델 가중치가 아닌 경량화된 프로토타입 코드북($I \times D$ 파라미터)만 교환됩니다.
• 학습: 각 라운드에서 클라이언트는 로컬 인코더, 디코더 및 로컬 코드북을 업데이트합니다. 업데이트된 코드북을 서버에 업로드하면, 서버는 이를 정렬하고 집계하여 새로운 전역 코드북을 형성합니다.

3. 주요 기여

1. 연합 디컨파운딩 프레임워크: 본 논문은 클라이언트를 분산된 인과적 환경으로 취급하는 새로운 인과적 관점을 도입합니다. Atlas는 이질적인 클라이언트를 가로질러 공유된 환경적 컨파운딩 효과를 학습하여 일반화 성능을 향상시킵니다.
2. 암시적 백도어 조정으로서의 연합: 이론적으로, 본 논문은 프로토타입 코드북을 집계하는 것이 환경적 컨파운딩을 완화하기 위한 조건을 확립함을 밝힙니다. 또한 디컨파운딩 오차가 개별 클라이언트의 편향이 아닌 평균 컨파운딩 강도에 의존한다는 바운드를 도출합니다.
3. 실증적 검증: 5개의 실제 데이터셋(METR-LA, PEMS-BAY, PEMS03, PEMS04, KnowAir)에 대한 광범위한 실험을 통해 Atlas가 기존의 연합 및 중앙 집중식 베이스라인들을 능가함을 입증했습니다. 이 프레임워크는 전이 가능하고, 해석 가능하며, 통신 효율적임이 확인되었습니다.

4. 실험 결과

4.1 예측 성능

Atlas는 5개 데이터셋 모두에서 모든 연합 베이스라인(FedAvg, FedProx, FedPer, FedGRU, FedGTP 및 Fuels, FedDis와 같은 특화된 방법 포함)을 일관되게 능가합니다.

• METR-LA에서 Atlas는 MAE를 3.30으로 낮추어, 가장 우수한 연합 베상라인인 FedGTP(3.54)보다 뛰어난 성능을 보였으며, 중앙 집중식 방법(DCRNN, 3.05)과의 격차를 4.9%~6.7% 이내로 좁혔습니다.
• PEMS-BAY에서 Atlas는 MAE 1.68을 달성하여, 최상의 연합 베이스라인인 FedGTP(1.73)를 앞질렀습니다.

4.2 절제 연구(Ablation Studies)

• 연합 집계: 집계를 제거하면(로컬 학습만 수행) 성능이 대폭 하락하여(METR-LA에서 MAE 77% 증가), 단일 클라이언트가 안정적인 환경 레짐을 포착할 수 없음을 확인했습니다.
• 공유 아키텍처: 클라이언트 간 통합된 인코더/디코더를 강제하는 경우(A.2, A.3) 성능이 저하되었으며, 이는 클라이언트 특유의 시계열 모델링이 보존되어야 함을 나타냅니다.
• 정렬(Alignment): 프로토타입 정렬 단계를 제거하면(A.4) 일관된 성능 저하가 발생하여, 집계 전 환경적 의미를 정렬하는 것의 필요성을 입증했습니다.
• 소프트 할당(Soft Assignment): 소프트 할당을 하드 할당으로 교체하면(A.5) 성능이 저하되어, 레짐 간의 매끄러운 전이를 모델링할 필요가 있음을 보여주었습니다.

4.3 전이 가능성 및 효율성

• 전이 가능성: 학습된 코드북은 **보지 못한 클라이언트(unseen clients)**에 효과적으로 전이됩니다. 5%의 클라이언트로만 사전 학습했을 때, 고정된 코드북은 동일한 서브셋에서 처음부터 학습된 모델과 유사한 성능을 보이며, 무작위 초기화보다 훨씬 뛰어난 성능을 보입니다.
• 효율성: Atlas는 계산량과 통신량 사이의 유리한 트레이드오프를 달성합니다. 클라이언트는 전체 모델 가중치를 업로드하는 FedAvg나 FedGRU와 같은 방법보다 두 자릿수 더 작은(예: 640개 파라미터) 코드북만을 업로드합니다.
• 강건성: 코드북은 가우시안 및 Rademacher 노이즈에 대해 강건하며 성능 저하가 미미하여, 안정적인 환경적 의미를 포착하고 있음을 시사합니다.

4.4 해석 가능성

KnowAir 데이터셋에 대한 시각화 결과, 프로토타입 할당이 타겟 변수(PM2.5)뿐만 아니라 뚜렷한 환경 레짐(예: 깨끗한 시기 vs 오염된 시기) 및 기상 조건(온도, 기압, 습도)과 상관관계가 있음을 보여줍니다. 이는 코드북이 잠재적인 환경적 맥락을 포착하고 있음을 확인시켜 줍니다.

5. 의의 및 주장

본 논문은 Atlas가 연합 STF의 패러다임 전환을 나타낸다고 주장합니다:

• 방해 요소에서 자원으로: 클라이언트 간 이질성을 억제해야 할 최적화 장애물으로 보는 대신, Atlas는 이를 환경 레짐에 대한 분산된 증거로 활용합니다.
• 인과적 일반화: 연합을 백도어 조정을 위한 메커니즘으로 이론적으로 프레임화함으로써, 전통적인 개인화 FL이 실패하는 환경 변화 상황에서도 모델이 일반화할 수 있도록 합니다.
• 실용적 생존 가능성: 이 프레임워크는 개인정보 보호가 중요한 시공간 예측에 대한 실질적인 솔루션을 제공하며, 최소한의 통신 오버헤드로 높은 성능과 해석 가능한 환경 표현을 제공합니다.

저자들은 이산적인 프로토타입이 연속적이거나 극도로 고차원적인 변화에는 어려움을 겪을 수 있다는 점과, $O(I^3)$의 정렬 복잡도가 매우 큰 프로토타입 규모에서는 병목이 될 수 있다는 점을 언급하며 한계를 인정했습니다. 향-후 연구는 확장 가능한 정렬 전략과 동적 코드북을 탐구하는 것을 목표로 합니다.