[FL] FedAvg (Federated Averaging)

FedAvg, Federated Averaging, FL FedAvg, FedSGD, 연합 평균, communication-efficient FL, McMahan FedAvg

정의

FedAvg (Federated Averaging) 는 McMahan et al. (2017) 이 제안한 연합 학습의 기본 알고리즘 입니다. 원본 데이터를 중앙 서버로 보내지 않고 각 클라이언트가 로컬 데이터로 모델을 학습한 뒤 weight 을 가중 평균 하여 글로벌 모델을 갱신합니다.

“local 여러 step, 서버 평균 한 번” 이라는 단순 구조로 FedSGD 대비 통신 라운드 수를 10-100 배 줄여, 실용적 연합 학습의 문을 열었습니다.

FedSGD 와의 차이

FedSGD (naive baseline)

• 각 라운드에서 클라이언트가 1 step gradient 만 계산해 서버로 전송
• 서버가 gradient 를 평균해 global weight 갱신
• 각 SGD step 마다 통신 -> 매우 비효율

FedAvg (핵심 개선)

• 각 클라이언트가 로컬에서 여러 epoch 을 돌리고
• 최종 weight 를 서버에 전송
• 서버는 weight 을 가중 평균
• 한 라운드에 로컬 여러 step 이 들어가므로 통신 횟수 대폭 감소

핵심 통찰: 데이터가 IID 에 가깝다면 로컬 여러 step 이 낭비가 아니라 이득 (서버 라운드 = 통신 = 병목이므로).

알고리즘

Server:

Initialize w_0
For each round t = 0, 1, 2, ...:
    S_t = random subset of K clients (fraction C of N clients)
    For each client k in S_t (in parallel):
        w_k^{t+1} = ClientUpdate(k, w_t)
    w_{t+1} = sum over k in S_t of (n_k / n) * w_k^{t+1}

Client (ClientUpdate(k, w)):

Split D_k into batches of size B
For each local epoch e = 1..E:
    For each batch b:
        w <- w - eta * grad_l(w; b)
Return w

하이퍼파라미터

기호	의미	관용 값
$N$	전체 클라이언트 수	100 ~ 10^6+ (cross-device)
$C$	라운드당 참여 fraction	0.001 ~ 0.1 (cross-device), 1.0 (cross-silo)
$K$	라운드당 실제 참여 클라이언트 수	$C\cdot N$
$E$	로컬 epoch 수	1 ~ 20
$B$	로컬 minibatch 크기	10 ~ 200
$\eta$	로컬 learning rate	0.001 ~ 0.1
$T$	총 라운드 수	수백 ~ 수만

서버 집계 공식

가중 평균 (weighted average):

$$w_{\{}t+1\}=\sum _{\{}k\in S_t\}{\frac{\{}{n}}_k\}\{n\}\cdot w_k^{\{}t+1\}$$

• $n_k$: 클라이언트 $k$ 의 데이터 개수
• $n={\sum }_{\{}k\in S_t\}{n}_k$: 이번 라운드 참여 클라이언트 총 데이터 개수

왜 개수 비율로 가중: 데이터가 많은 클라이언트의 로컬 optimum 이 실제 loss 표면에 대해 더 많은 정보를 담고 있다고 가정 (empirical risk minimization 관점의 자연스러운 가중).

주의: 이 가중은 개수 편향 을 유발할 수 있습니다. 대형 클라이언트가 지배적이면 소수 클라이언트의 분포가 무시됩니다. 문제가 있으면 uniform 또는 클래스 균형 가중으로 대체.

왜 로컬 여러 epoch 이 통하는가

IID 데이터 (모든 클라이언트가 같은 분포에서 샘플링) 라면 각 클라이언트의 로컬 optimum 이 글로벌 optimum 과 근접합니다. 여러 step 이 낭비가 아니라 서버 라운드 사이의 진전.

Non-IID 라면 각 클라이언트가 자기 로컬 optimum 으로 수렴하려 하고, 서로 다른 방향으로 이동합니다 (client drift). 평균이 잘못된 방향으로 이동해 수렴 속도 저하 또는 성능 하락. 이 문제는 FL Non-IID 위키에서 상세 다룹니다.

Convergence 이론

Li et al. (2019) “On the Convergence of FedAvg on Non-IID Data” 는 convex loss 가정 하에 다음을 보였습니다.

• IID + full participation: FedSGD 와 유사한 rate 로 수렴
• Non-IID + partial participation: decay learning rate 가 필요, 아니면 수렴하지 않을 수 있음
• 통신 라운드 $T$ 관점에서 $O(1/T)$ (strongly convex), $O(1/\sqrt{\{}T\})$ (convex)

Non-convex (딥러닝) 는 정확한 이론이 없고 empirical validation 에 의존합니다.

FedSGD 로부터의 도출

E = 1, B = 무한 (전체 배치) 이면 FedAvg = FedSGD. 즉 FedSGD 는 FedAvg 의 특수 케이스.

E 를 늘리면:

• 통신 감소 (좋음)
• Client drift 증가 (나쁨, non-IID 에서)
• Batch normalization 통계 로컬 편향 심화

Sweet spot: E = 1-5, B = 클라이언트 데이터의 1/10 정도.

통신 비용 분석

한 라운드 통신량:

$$\text{총 통신량}=2\cdot K\cdot |w|$$

• 서버 -> 클라이언트: global weight 다운로드
• 클라이언트 -> 서버: 로컬 학습 후 weight 업로드

$|w|$ 는 모델 파라미터 총 크기 (float32 기준 4 bytes/param).

최적화 방법:

• Model compression: quantization (8-bit, 4-bit), sparsification (top-k)
• Structured updates: low-rank, random mask
• FetchSGD: Count sketch 로 gradient 압축
• Federated Dropout: 부분 모델만 학습/전송

배치 정규화의 함정

Batch Normalization 은 배치 통계에 의존합니다. 로컬 데이터가 non-IID 이면 로컬 BN 통계가 글로벌 분포를 대표하지 않아 성능 저하.

해결:

• GroupNorm: 배치 무관 (BN 대체 자주 사용)
• LayerNorm
• FedBN: BN 파라미터는 로컬 유지, 다른 파라미터만 aggregate

실전 구현 (Flower)

import flwr as fl
import torch
from torch.utils.data import DataLoader

class Client(fl.client.NumPyClient):
    def __init__(self, model, train_ds, val_ds):
        self.model = model
        self.train_dl = DataLoader(train_ds, batch_size=32, shuffle=True)
        self.val_dl = DataLoader(val_ds, batch_size=64)

    def get_parameters(self, config):
        return [p.detach().cpu().numpy() for p in self.model.parameters()]

    def set_parameters(self, parameters):
        for p, new in zip(self.model.parameters(), parameters):
            p.data = torch.tensor(new)

    def fit(self, parameters, config):
        self.set_parameters(parameters)
        # 로컬 E epoch 학습
        opt = torch.optim.SGD(self.model.parameters(), lr=config["lr"])
        for epoch in range(config["local_epochs"]):
            for x, y in self.train_dl:
                opt.zero_grad()
                loss = self.model.loss(x, y)
                loss.backward()
                opt.step()
        return (
            self.get_parameters({}),
            len(self.train_dl.dataset),  # n_k, 서버가 가중 평균에 사용
            {}
        )

    def evaluate(self, parameters, config):
        self.set_parameters(parameters)
        loss, acc = self.eval()
        return float(loss), len(self.val_dl.dataset), {"acc": acc}

# 서버
strategy = fl.server.strategy.FedAvg(
    fraction_fit=0.1,           # C = 0.1
    min_fit_clients=10,
    min_available_clients=100,
    on_fit_config_fn=lambda rnd: {"lr": 0.01, "local_epochs": 1},
)

fl.server.start_server(
    server_address="0.0.0.0:8080",
    config=fl.server.ServerConfig(num_rounds=200),
    strategy=strategy,
)

FedAvg 의 개선/변형

변형	개선점
FedProx	Non-IID client drift 완화 (proximal term)
SCAFFOLD	Control variate 로 drift correction
FedYogi / FedAdam	서버측 adaptive optimizer
FedNova	Local step 수 불균형 정규화
FedBN	Batch Norm 통계 로컬 유지
Personalized FedAvg	클라이언트별 로컬 fine-tuning

각 변형은 특정 조건 (non-IID 극심, system heterogeneity 등) 에 특화. 자세한 알고리즘 유도는 원 논문 참조.

함정

관련 위키

• Federated Learning - 상위 개념
• Non-IID Data in FL - FedAvg 의 주요 한계
• Secure Aggregation - 프라이버시 강화 집계
• Personalized FL - Local fine-tuning 결합
• FL Frameworks - Flower/TFF/NVFlare 비교
• Differential Privacy - DP-SGD 결합