[AI System] Federated Learning - Part 1

"Federated Learning - Part 1"

[Basic of FL(Federated Learning)]

· Edge Device로 데이터를 직접 생성 후, Privacy와 Communication Cost에 의해 Server로 전송 X → Device에서 학습

· But, 각 사용자 데이터가 분산되어 데이터가 부족 & Privacy에 의해 서로 데이터 공유 불가

▣ FL(Federated Learning) - 연합 학습

· 각 Client는 자체 Local Data를 보유하며, 이 데이터는 비공개로 유지

· Server와 Client 간의 Model 교환을 반복적으로 수행하여 데이터 공유 없이 Model을 공동으로 학습 (Model을 공유)

→ Server가 Model을 보내면 각 Edge Device가 학습 후 Update만 Server로 보내고, Server가 이를 합침

▣ Objective Function

· F_k(w) : Client k의 Local Loss 함수 / F(w) : Global Objective Function / L(w, x) : 개별 Sample x에 대한 Loss

· Client마다 데이터 개수가 다르므로 각 Cleint 내부에서는 평균 Loss로 정리하고, 전체 Objective에서 p_k로 Weight 부여

· 전체 데이터를 한 곳에 모으면 Centralized Training(모든 데이터가 Server에 모여있고, 해당 데이터로 학습) Objective와 동일

∴ FL도 결국 전체 데이터 평균 Loss를 최소화 하려는 것, 다만 데이터를 직접 모으지 못해서 분산 방식으로 학습

· 각 Client는 다른 Class를 가지고 Server가 이를 통합하여 전체 Class를 Cover → 데이터 다양성 확보

▣ FL Process
> Step 0 : Model Initialization

· Server가 초기 Model 생성

> Step 1 : Model Download

· Client가 동일한 Model Download (각자 다른 Model에서 출발 시, 이후 평균을 내도 의미가 약함)

> Step 2 : Model Update with Local Data

· 각 Client k는 Local Dataset D를 사용하여 Model Update → 각 Client는 Local Dataset에 따라 고유한 Model을 얻음

> Step 3 : Model Upload & Aggregation

· Server가 Upload된 각 Client의 Model을 가지고 Model 평균 생성

※ FedAvg : Weighted Average(가중 평균)를 계산하여 새로운 Global Model 생성 w₁ = ∑p_k · w₁^k

    → Client 1의 데이터가 1,000개, Client 2의 데이터가 10개 → 데이터 수가 많은 Client의 Update를 더 크게 반영

w₀ → w₁ → w₂ → … → w_T-1 → w_T

· Multiple Global Rounds(or Communication Rounds) =  원하는 성능 달성을 위해 Step 1~3 반복

    → Centralized Training에서 Epochs를 사용하여 ML Model을 학습하는 과정과 유사

· Round가 반복될수록 Global Model이 점점 개선되며 최종적으로 얻는 w_T가 배포할 Model이 되고, 성능 평가의 대상이 됨

▣ Key Characteristics in FL

> 여러 번의 Local Update 후 전송

· Client가 작은 단위의 학습(Mini-Batch)을 할 때마다 Server로 보내면 통신 횟수 기하급수적으로 증가

∴ 각 Client에서 여러 번의 Local Update 수행 후 전송 → 통신 빈도 감소

> Non-IID(Independent, Identically Distributed) 데이터

· 각 Client가 가지고 있는 학습 데이터의 특성이 서로 매우 다름 = Heterogeneous(이질적)

▣ FedAvg Algorithm

· 매 라운드마다 모든 Client가 참여하는 것이 아니라, 일부만 참여 가능(Active Clients) → 전체 Client 동시 참여 X / 부분 참여 O

    ∵ 배터리 부족 / 네트워크 불안정 / Device 사용 중 / 자원 제한

· Client k는 Server에서 받은 w_t를 시작으로 E번의 Local Update 수행 후 얻은  w_{t+1, E}^k Model을 Server로 전송

    → 한 Round 안에서 각 Client가 자기 손실함수 F_k를 줄이는 방향으로 여러 번 수행

    → E가 크면 통신 횟수↓, But, Local Data에 Overfitting 될 가능성↑

** 사실 Server는 "얼마나 바뀌었는지"만 알면 충분

∴ 최종 Local Model(w_{t+1, E}^k)이 아닌, Client가 이번 Round에서 만든 Update 양(g_t^k = w_t+1^k - w_t)을 전송해도 됨

→ Server는 이것을 합쳐서 새로운 Global Model(w_t+1 = w_t + ∑p_k · g_t^k)을 얻고, Client에게 Model(w_t+1)을 전송

▣ FedAvg 성능

· IID Case : 각 Client가 비슷한 Class 분포를 가짐 (Client1 : 0~9 골고루 분포 / Client2 : 0~9 골고루 분포)
· Non-IID Case : 각 Client가 다른 Class나 다른 특성을 가짐 (Client1 : 1, 4 위주 / Client2 : 2, 7 위주)

** 실제 현실은 대부분 Non-IID에 가까움

· FedAvg가 IID에서는 비교적 안정적으로 정확도 상승 / Non-IID에서는 느리거나 덜 안정적

→ 기본 FedAvg는 간단하고 효과적이지만 데이터가 Non-IID인 경우 한계 존재

[Non-IID Data in FL]

· Non-IID & 여러 번의 Local Update → 평균 Model이 Global Optimum으로 느리게 수렴하거나, 아예 수렴하지 않을 수 있음

∵ 각 Client의 Local Optimum이 다름(같은 Global Model에서 출발해도 각자 다른 방향으로 많이 이동하려고 함)

→ Local 학습을 오래 하면 할수록 Client는 Global Optimum보다 자기 Local Optimum 쪽으로 더 끌려감

※ Local Optimum : Client 각자의 데이터에 가장 좋은 지점 / Global Optimum : 전체 데이터를 고려했을 때 가장 좋은 지점

※ Client Drift : Client가 Global Optimum에서 벗어나 자기 Local Data에 맞는 방향으로 너무 멀리 Drift(이탈)하는 현상

    → 평균을 내도 Global 최적점 가까이 못 감 + 수렴이 느려지거나 불안정해짐

Q) 데이터가 IID여도 이 문제가 생길까?

A) 큰 문제 발생 X (∵ IID면 각 Client Local Optimum이 Global Optimum 근처에 존재 → 각자 다른 방향으로 크게 벗어나지 않음)

Q) Client가 여러 번이 아니라 한 번만 Update하면 이 문제가 생길까?

A) No. (이는 중앙 집중식 학습에서 Mini-Batch SGD와 같아짐)

∴ 문제의 핵심 = Non-IID & Multiple Local Updates

▣ FedProx

· FedAvg : 각 Client가 자기 Loss만 최소화

· FedProx : 현재 Global Model(w_t)에서 멀어질수록 Penalty 증가

· μ = 0 → FedAvg와 동일 / μ↑ → Global Model에 더 강하게 묶임

① Client Update (Loss 함수만 변경)

② Server로 Upload

③ 평균

→ Loss 함수 변경 외에는 FedAvg와 동일, But, Client Drift 감소로 인해 FedAvg와 달리 Non-IID 상황에서 안정성 증가

▣ SCAFFOLD

· Gradient를 보정하여 Client가 Local Optimum 방향으로 이동하는 것을 Global Optimum 방향으로 가도록 보정

※ ∇F_k(w_t) : 기존 Gradient / c_t : Global 보정 / c_t^k : 각각의 Client 보정

· FedAvg : Local Update를 여러 번 할수록(E↑) Client Drift 증가(= Local Data의 성향이 진하게 반영)

· SCAFFOLD : 각 Step마다 보정 Term이 같이 들어가므로 Local Step 수(E)가 커져도 Client Drift를 보정

△ E↑ → 통신 횟수↓(통신이 제일 비싼 자원)

▼ 추가 변수 저장 필요

▼ 통신해야 할 정보 증가

▼ 구현 복잡

> FedAvg

△ 단순

▼ Non-IID에 취약

> FedProx

· Global Optimum에서 멀리 가지 못하도록 제약

△ 추가 통신 X

△ 구현 쉬움

▼ μ 튜닝 필요

> SCAFFOLD

· 방향 자체를 보정

△ 추가 Hyperparameter 튜닝이 상대적으로 덜 중요

▼ Control Variate를 보내야 하므로 통신/저장 부담 추가

[Communication/Computation in FL]

▣ Straggler

· Model이 크면 통신/연산 지연 크게 발생

· Client마다 연산 능력과 통신 기능 다름 + 이질적인 Resource → 일부 Client는 다른 Client보다 훨씬 느려질 수 있음

※ Straggler : 가장 느린 Client가 전체 진행 속도 결정 (원인 : System Failure / Battery Outage / Mobility / Packet Loss)

▣ Handling Straggler Issue

> Wait for Stragglers

· 모든 Client가 다 끝날 때까지 대기

△ 해당 Round에 선택된 모든 Client 정보 사용 가능

▼ 가장 느린 Client 때문에 전체 속도 감소

> Timeout Threshold based Aggregation(Ignore Stragglers)

· 일정 시간까지만 기다리고 늦은 Client는 무시

△ 속도 향상

▼ Timeout 안에 들어온 Client 수가 너무 적을 수 있음 → 항상 같은 빠른 Client만 반영될 수 있음

> Wait for X% of Stragglers

· 전체의 X%가 오면 Aggregation

△ 최소 참여 수 확보 가능

▼ 여전히 느린 Client를 기다림

▼ 느린 Client가 계속 배제될 가능성 존재

> Asynchronous Aggregation (FedAsync)

· Model이 도착하는 즉시 Server가 반영

△ 가장 빠름

△ 기다릴 필요가 거의 없음

▼ Stale Update 문제 발생 가능

▼ 원래 FL Objective로 수렴하지 않을 수 있음

    ∵ Asynchronous에서는 각 Client가 서로 다른 시점의 Global Model을 기준을 학습

        → Server가 모으는 Update들이 같은 기준점에서 발생한 것이 아님    ∴ FedAvg가 상정한 "Synchronous 평균"과 달라짐

** 위 Straggler 대책은 "기다리는 방식 변경" 또는 "Aggregation Timing 변경" → 실제로 전송해야 하는 데이터 양을 줄이지 않음

▣ Gradient Compression for FL

g_t^k = w_t+1^k - w_t

· Server가 필요한 것은 Client의 최종 Local Model(w_t+1^k)이 아닌 Update 양 → 변화량을 Sparsify/Compress 가능

∴ 통신/연산 자체를 감소

· Magnitude가 큰 원소만 남기고 나머지는 0으로 변경 (top_Q([1, 2, 3, -4]) & Q = 2 → [0, 0, 3, -4])

· 각 Client가 압축된 Update를 보내면, Server는 그 Sparse Update들을 평균내서 Global Model 갱신

→ 구조는 여전히 FL인데, 보내는 메세지만 Sparse해짐

∵ Gradient의 절댓값↑ = Parameter 방향의 변화가 현재 학습에서 상대적으로 중요

∴ 작은 성분은 버리고 큰 성분만 전송 → 적은 통신량으로도 주요 Update 방향은 전달 가능

** 작은 Gradient가 여러 차원에서 누적되면 의미↑

    → 계속 버려지면 Bias 발생 가능 = Global Optimum 경로가 원래  Gradient Descent와 달라짐

▼ 수렴 속도 저하

▼ 수렴 실패 가능성 존재

▼ 정확도 저하 발생 가능

▣ Error Feedback

· 이번에 못 보낸 정보를 Local에 저장해뒀다가, 다음번 Gradient에 다시 더함

> Step 1 - Gradient Computation

· Client가 현재 Global Model(w_t)에서 시작해서 Local Update 후 Gradient/Update (g_t^k) 계산

> Step 2 - Error Feedback

· 이전 압축 때문에 못 보낸 부분(e_t^k)을 현재 Gradient에 더함

> Step 3 - Compression & Transmission

· 보정된 신호에 top_Q를 적용해서 Server로 전송

> Step 4 - Error Update

· 보정된 신호 중에서 이번에도 못 보낸 나머지를 새로운 Error Buffer에 저장 → 버려진 정보는 사라지지 않고 누적

∴ 작은 Gradient 값이 Round마다 누적되어 추후에는 중요한 값으로 커져 전송될 가능성 존재

** Error 보상은 Delay되어 입력 → Full Gradient와 완전히 같지는 않고 성능이 약간 감소 가능

Q) Model에 top_Q 적용 = Gradient에 top_Q 적용 ?

A) 일반적으로 같지 않음.

▣ Binary Sparsification

top_Q-bin([1.1, 2.2, 3.5, -4.7]) & Q = 2 → [0, 0, 3.5, -4.7] → [0, 0, +1, -1]

· 단순히 몇 개만 남기는 것뿐만 아니라, 남긴 값도 더 적은 Bit로 표현 → Gradient를 더 강하게 압축

(Like 위치 Index만 보내고, 값은 부호 또는 제한된 표현만 전송 → 통신량 감소)

△ Spare + Low-Bit Representation을 통해 통신 효율↑

▼ 정보 손실 위험↑

Q) Gradient Compression 방법들이 Client 쪽 계산량 또는 메모리 사용량을 줄여주나?

A) No. ∵ Gradient를 보내는 양은 줄일 수 있어도 Client는 여전히 큰 Model에 대해 Forward/Backward 하고 Gradient도 계산

→ 통신 최적화 ≠ 계산 최적화

▣ Reducing Computation Burden(부담) in FL

· LoRA : 작은 Low Rank 행렬 A, B만 학습

→ Trainable Parameter 수 + Backward 계산량 + Optimizer State + 저장해야 할 Gradient/Parameter 감소

∴ LoRA 사용 시 전체 Model이 아닌 A, B 행렬만 전송하여 계산량, 저장공간, 통신량을 동시에 감소

▣ Freeze Matrix A During Federated LoRA

· 기본 LoRA도 이미 가볍지만, A와 B 모두 Update하면 비용 존재하므로 A 고정 후 B만 Update하여 학습

→ 기존 LoRA : (d₁ + d₂)r / A 고정 후 B만 Update : d₂r

[Reference]

· Federated Learning Part 1 (AI System : Week 5) - Dongjun Han