이 논문에선 기존 분산 학습 방법의 하나인 FedAvg의 느린 업데이트 집계와 대역폭 병목 문제를 지적하고 이를 해결하기 위해 over-the-air computation (AirComp)방법을 제안한다. AirComp는 무선 다중 접속 채널의 신호 중첩 성질을 활용해 곧바로 업데이트를 계산하여 통신과 집계를 동시에 빠르게 수행한다. 이 방법은 sparse, low-rank, difference-of-convex(DC)등의 수학적 방법을 통해 이 모델이 이론적으로 안정된 알고리즘(global convergence)임을 증명한다. 실험 결과, 제안한 방법은 기존 방법보다 더 많은 디바이스 수를 유지해 지연을 줄이고 학습 손실을 더 빨리 줄이며 예측 정확도를 높였다.
논문 링크:
Federated Learning via Over-the-Air Computation
The stringent requirements for low-latency and privacy of the emerging high-stake applications with intelligent devices such as drones and smart vehicles make the cloud computing inapplicable in these scenarios. Instead, <italic xmlns:mml="http://www.w3.or
ieeexplore.ieee.org
출처:
K. Yang, T. Jiang, Y. Shi and Z. Ding, "Federated Learning via Over-the-Air Computation," in IEEE Transactions on Wireless Communications, vol. 19, no. 3, pp. 2022-2035, March 2020.
요약
드론, 자율 주행 자동차, 증강 현실과 같은 다음 세대 기술은 저지연과 프라이버시가 요구된다. 클라우드 컴퓨팅은 이러한 요구사항을 충족하기 어렵기에 엣지 머신러닝이 주목받고 있다. 이러한 엣지 머신 러닝 중 Federated Learning(FL)은 프라이버시를 보장하며 글로벌 모델을 학습할 수 있다.
그러나 FL은 불균형하고 non-IID한 데이터, 단말의 계산, 저장, 에너지, 대역폭 자원의 제약과 같은 문제를 갖는다. 또 각 단말의 계산 능력, 저장 용량, 통신 성능이 다른 이질성(Heterogeneity) 문제를 가지며 악의적인 공격으로 성능이 저하될 수 있으며 실제 환경에서 시스템을 구현하는데 여러 어려움이 따른다. 이에 따른 해결책으로 Federated Averaging이 제안되었다.
FedAvg는 여러 단말이 각자의 로컬 데이터셋으로 모델을 학습한 뒤 중앙 서버가 이들을 가중 평균하여 글로벌 모델을 갱신하는 방식이다. 이를 통해 통신 라운드 수를 줄여 효율성을 높이고 글로벌 모델 집계를 가능하게 한다. 그러나 로컬에서의 업데이트를 전송하는데 필요한 대역폭 병목이 일어날 수 있다.

그리하여 이 논문에서는 무선 다중접속 채널의 신호 중첩 성질을 활용하는 over-the-air computation (AirComp) 방법을 제안하여 FedAvg가 빠른 글로벌 모델 학습을 가능하게 한다. AirComp는 각 단말이 순차적으로 데이터를 보내는 orthogonal 방식과 달리 동시 전송을 허용하여 통신 라운드 수가 줄고 집계 과정이 빠르게 이루어져 저지연 학습이 가능해진다.
AirComp를 설계할 땐 각 학습 라운드의 단말 수를 최대한 크게 하여 수렴 속도를 빠르게 해야 한다. 또 로컬 업데이트는 집계 오차가 크면 글로벌 모델의 정확도가 낮아지기에 빔포밍을 통해 이를 최소화해야 한다. 그러나 단말 선택과 수신기 빔포밍을 동시에 최소화하는 문제는 non-convex(비볼록)하고 quadratic constraints(이차 제약)하기에 해결이 어렵다. 이를 해결하기 위해 sparse 및 low-rank로 문제를 최적화시킨 후 DC 기법을 적용해 효율적으로 풀 수 있도록 설계한다.
\[
\begin{aligned}
&\underset{S,\, \mathbf{m} \in \mathbb{C}^N}{\text{maximize}} && |S| \\
&\text{subject to} && \max_{i \in S} \left( \frac{\phi_i^2 \|\mathbf{m}\|^2}{| \mathbf{m}^H \mathbf{h}_i |^2} \right) \leq \gamma
\end{aligned}
\]
Sparsity는 l0-norm이 l1-norm과 Ky Fan k-norm의 차와 같다는 것을 통해 강화한다.
\[
\|\mathbf{x}\|_0 = \min \left\{ k : \|\mathbf{x}\|_1 - \|\mathbf{x}\|_k = 0,\; 0 \leq k \leq M \right\}
\]
Rank-one은 trace norm(모든 고유값의 합)과 spectral norm(가장 큰 고유값)의 차이가 0 임을 통해 확인한다.
\[
\text{rank}(M) = 1 \iff \text{Tr}(M) - \|M\|_2 = 0
\]
단말 선택을 효율적으로 하기 위해선 아래 그림과 같은 두 단계 절차가 요구된다.

먼저 첫번째 문제 PS1에선 각 단말이 MSE를 얼마나 잘 만족하는지 수치화하여 단말 선택 우선순위를 정한다. 이를 통해 rank-one 제약하에서 벡터 x의 sparce 구조를 얻는다.
\[
\begin{aligned}
\text{PS1:} \quad
& \min_{x,\, M} \; \|x\|_1 - \|x\|_k + \text{Tr}(M) - \|M\|_2 \\
& \text{subject to} \quad \text{Tr}(M) - \gamma_i h_i^H M h_i \leq x_i, \quad \forall i = 1, \dots, M, \\
& \qquad M \succeq 0, \quad \text{Tr}(M) \geq 1, \quad x \geq 0
\end{aligned}
\]
두번째 단계에선 앞에서 정한 우선순위를 바탕으로 실제 선택된 단말 집합이 MSE 제약을 만족하는지 검증한다. 앞 단계에서 우선순위를 정했기에 계산량을 줄일 수 있다.
\[
\begin{aligned}
\text{PS2:} \quad
& \min_{M} \; \text{Tr}(M) - \|M\|_2 \\
& \text{subject to} \quad \text{Tr}(M) - \gamma_i h_i^H M h_i \leq 0, \quad \forall i \in S[k], \\
& \qquad M \succeq 0, \quad \text{Tr}(M) \geq 1
\end{aligned}
\]
앞서 정의 했던 PS1, PS2에 DC 알고리즘의 수렴성을 보장하기 위해 α에 곱해진 이차항들을 더해 strongly convex 하게 만든다. 이후 각 문제 f를 g와 h의 차이로 변환해 두 개의 strongly convex 한 함수의 차이를 최소화하는 문제로 바꾼다. 이 변환을 통해 non-convex 최적화 문제에서 DC 알고리즘이 적용 가능한 형태로 문제가 바뀌었다.
\[
\min_{X \in \mathbb{C}^{m \times n}} f(X) = g(X) - h(X)
\]
\[
\begin{aligned}
f_1(x, M) &= x_1 - \|x\|_k + \mathrm{Tr}(M) - \|M\|_2 + I_{C_1}(x, M) \\
g_1(x, M) &= x_1 + \mathrm{Tr}(M) + I_{C_1}(x, M) + \tfrac{\alpha}{2} \left( \|x\|_F^2 + \|M\|_F^2 \right) \\
h_1(x, M) &= \|x\|_k + \|M\|_2 + \tfrac{\alpha}{2} \left( \|x\|_F^2 + \|M\|_F^2 \right)
\end{aligned}
\]
\[
\begin{aligned}
f_2(M) &= \mathrm{Tr}(M) - \|M\|_2 + I_{C_2}(M) \\
g_2(M) &= \mathrm{Tr}(M) + I_{C_2}(M) + \tfrac{\alpha}{2} \|M\|_F^2 \\
h_2(M) &= \|M\|_2 + \tfrac{\alpha}{2} \|M\|_F^2
\end{aligned}
\]
그 다음으로 수렴성을 보장하기 위해 기존의 문제와의 dual problem을 아래와 같이 정의한다. 각 iteration에서 오목한 부분을 선형화하여 convex로 근사하여 풀고 해를 갱신한다. 기존 문제와 dual 문제를 번갈아가며 풀면서 점점 더 좋은 해를 찾아간다. 이 과정을 거쳐 sparse/low-rank 한 해를 얻게 된다.
\[
\min_{Y \in \mathbb{C}^{m \times n}} \; h^{*}(Y) - g^{*}(Y)
\]
\[
Y[t] = \arg \inf_{Y \in \mathcal{Y}} \; h^{*}(Y) - \Big( g^{*}(Y[t-1]) + \langle Y - Y[t-1], X[t] \rangle \Big)
\]
\[
X[t+1] = \arg\min_{X \in \mathcal{X}} \; \Big( g(X) - h(X[t]) - \langle X - X[t], Y[t] \rangle \Big)
\]
이렇게 알고리즘을 반복 수행하면 해 M, x는 어떠한 임계점에 도달한다. 아래와 같이 iteration마다 목적 함수 f는 엄격히 감소함을 수학적으로 보였다. 이는 다시 말하면 알고리즘이 발산하지 않고 점점 더 좋은 해로 수렴하여 글로벌 수렴을 보장한다는 뜻이다.
\[
\begin{aligned}
\text{Avg}\,\|M[t] - M[t+1]\|_F^2 &\leq \frac{f[0]_1 - f^\star_1}{\alpha (t+1)} \\
\text{Avg}\,\|x[t] - x[t+1]\|_2^2 &\leq \frac{f[0]_1 - f^\star_1}{\alpha (t+1)}
\end{aligned}
\]
시뮬레이션은 M=20개의 모바일 디바이스와 N=6개의 안테나를 가진 BS에서 진행됐다. 채널은 i.i.d 복소 가우시안 분포를 따르는 무선 채널 모델을 사용하였으며 평균 SNR은 20dB로 설정되었다. 또한 각 기기들은 같은 수의 데이터 포인트를 갖는다고 가정했다. 아래 그림은 제안된 방법이 기존 SDR 방법보다 더 높은 feasability 확률을 가지며 global optimization과 거의 동일하게 근접함을 나타낸다. 가로축은 MSE threshold를 나타내며 이 값이 커질수록 제약 조건이 완화되어 더 많은 단말을 선택할 수 있어 feasability 확률이 증가한다.

아래 그림은 기지국 안테나 수가 늘어날수록 단말 선택의 제약을 만족할 확률이 증가하는 것을 보여준다. 즉 더 많은 안테나가 통신 및 학습 성능을 향상시킨다는 뜻이다.

다음 시뮬레이션 결과는 DC 알고리즘이 다른 방법들보다 더 많은 디바이스를 선택할 수 있음을 나타낸다. 더 많은 단말이 동시에 학습에 참여할 수 있으므로 이는 통신 효율성과 학습 정확도 향상으로 이어진다. MSE 조건이 까다로워질수록 기존 방법들은 선택 가능한 디바이스의 수가 주는 반면 DC 접근법은 상대적으로 더 많은 디바이스 수를 유지한다.

마지막으로 이 논문에선 제안된 DC 알고리즘을 실제 FL 환경에 적용해 성능을 평가했다. (a)에서 제안한 DC 알고리즘은 다른 방법보다 손실이 빠르게 줄어들며 benchmark에 근접한 성능을 보인다. (b)에서도 DC 알고리즘이 다른 기법보다 더 높은 정확도를 달성하고 benchmark와 유사한 수준까지 도달한 것을 확인할 수 있다.

한계
- 완벽한 채널 상태 정보 가정: 수신 빔포밍 과정에서 채널 상태 정보를 완벽히 알고 있다고 가정하나 실제 무선 환경에선 채널 추정 오류나 시간 변동성이 존재하기 때문에 성능 저하 가능성이 크다.
- 보안 및 공격 대응 미흡: 제안된 방법은 보안 위협에 대한 구체적인 방어 전략을 포함하지 않는다. 대규모 분산 학습에서 악의적 참여자가 성능을 크게 떨어뜨릴 수 있다.
- 비용 및 복잡성 문제: sparse/low-rank 최적화와 DC 알고리즘은 계산적으로 무겁고 실제 대규모 네트워크에서 실시간으로 적용하기 어렵다.
- 비현실적인 가정: 모든 디바이스가 동일한 데이터 크기를 가진다고 가정하거나 잡음 모델을 단순화하는 등 실제 환경과 괴리가 있다. 실제 모바일 환경에선 데이터 불균형, 디바이스 이질성이 심하다.
- 실험 범위의 제한: CIFAR-10 데이터셋과 SVM 분류기라는 비교적 단순한 설정에서만 검증을 진행하고 다양한 모델과 데이터 셋에서의 일반화 가능성은 충분히 입증되지 않았다.