[NeurIPS 2021] FjORD - (1)

논문 제목: FjORD: Fair and Accurate Federated Learning under heterogeneous targets with Ordered Dropout

출처: https://arxiv.org/abs/2102.13451

 

 이번 paper는 system heterogeniety에 관한 내용을 다루고 있습니다. 만약 특정 device의 computing power가 부족해서 학습이 매우 느리게 진행되거나 혹은 아예 model이 메모리에 올라가지 못 한다면, 이러한 device는 곧바로 straggler가 될 것입니다. 해당 논문은 device 별로 model의 size를 조절하여 이 문제를 해결하고자 하였는데, 과연 어떠한 idea를 활용하였는지 지금부터 살펴보겠습니다.

 

 1. 연구 배경

 

 우리가 잘 알고 있듯이, FL은 크게 cross-silo setting과 cross-device setting으로 나뉩니다. 전자의 경우 보통 사용되는 device들의 computing power가 충분한 편이어서 크게 문제되지 않으나, 후자의 경우 스마트워치와 같이 computing power가 낮은 기기부터 GPU가 달려 있는 컴퓨터까지 아우르는 다양한 종류의 기기가 학습에 참여하는 관계로 학습 속도가 client마다 상이하다는 문제점이 있습니다. 극단적으로는 아예 메모리에 model과 data가 올라가지 못하는 경우도 존재할 수 있습니다.

 

 이러한 setting에서 우리가 기존에 사용하였던 방법은 크게 두 가지가 있었습니다. 하나는 사양이 낮은 기기에 맞춰 model size를 조절하는 방법, 다른 하나는 사양이 낮은 기기들을 straggler로 취급하고 일정 수준 이상의 computing power를 갖춘 device에 대해서만 학습을 진행하는 방법입니다. 사실 두 가지 다 좋은 방법은 아닙니다. 전자의 경우 자연스레 학습 속도가 늦춰질 것이고, 후자의 경우 버려진 device가 가진 data에 대한 학습이 전혀 이루어지지 않게 되기 때문입니다. 저자들은 이러한 상황을 해결해야 할 필요성을 느껴 해당 연구를 진행하였습니다.

 

2. Ordered Dropout(OD)

 

 해당 논문이 제안하는 알고리즘인 FjORD를 설명하기에 앞서, 우리는 Ordered Dropout이라는 tool을 이해할 필요가 있습니다. (물론, 이 역시 해당 논문에서 제안한 것입니다.) 기존의 Random Dropout(RD)과 OD의 차이점을 오른쪽 그림이 잘 보여주고 있는데, RD는 random한 node를 dropout하지만, OD는 오른쪽의 node부터 dropout한다는 것(즉, 왼쪽의 node들을 살려둔다는 것)을 알 수 있습니다. (여기에서, $p$는 살려두는 node의 비율입니다. 일반적으로 $p$는 dropout하는 node의 비율을 의미하는데, 해당 논문은 notation을 반대로 사용하였다는 점을 꼭 염두에 두시기 바랍니다.)

 

 OD를 이야기하기 위해서, 우선 $\mathcal{P} := \left\{s_1, s_2, \cdots, s_{|\mathcal{P}|} \right\}$를 정의합니다. OD는 이러한 $\mathcal{P}$ 안에서 하나의 원소 $p \in \mathcal{P}$를 뽑아서 dropout rate로 사용하는데, 이때 $\mathcal{P}$는 임의의 $i = 1, 2, \cdots, |\mathcal{P}|$에 대해서 다음을 만족합니다.

 

$(1)$ $0 < s_1 < s_2 < s_{|\mathcal{P}|} \leq 1$

$(2)$ $\textbf{P} (p = s_i) > 0$

$(3)$ $s_i  \sim D_{\mathcal{P}}$ (단, $D_{\mathcal{P}}$는 $\mathcal{P}$의 distribution)

 

 여기에서 $\mathcal{P}$의 분포로 사용될 수 있는 것은 다양하게 존재하지만, 저자들은 discrete uniform distribution을 사용하였습니다. 즉, $\mathcal{P} = \left\{\frac {i} {k} \mid i = 1, 2, \cdots, k \right\}$입니다. (discrete한 분포를 사용하는 이유는 node의 수가 discrete하기 때문입니다.) 또한, layer 별로 $p_l \sim D_{\mathcal{P}}^l$를 다르게 지정할 수도 있지만, 저자들은 이에 대한 이야기는 논외로 하기로 하고, 한 model에 대한 dropout rate는 동일하게 $p \in \mathcal{P}$로 지정하였습니다.

 

 그렇다면, 한 model이 가진 모든 layer는 동일하게 $p$만큼의 node(그리고 이러한 node들을 연결해주는 neuron)만이 남은 형태가 되는데, 이를 $p$-pruned subnetwork라고 부르고, $F_p$로 정의합니다. 그리고 이에 해당하는 weight는 $w_p$로 정의합니다. (여기에서, $F$와 $w$는 각각 원 model과 그 model의 weight입니다.) 정의에 따라서, $F_p$는 각 layer $l$의 크기가 $K_l$일 때, 왼쪽에서부터 $\{ 0, 1, \cdots, \lceil p \cdot K_l \rceil - 1 \}$번 node까지만 남아있고, $\{ \lceil p \cdot K_l \rceil, \cdots, K_l - 1 \}$번 node는 pruned된 상황입니다. 그런데, 여기에서 이야기하는 pruning은 일반적으로 통용되는 pruning과는 다소 차이가 있습니다. 보통의 pruning은 여러 번의 실험을 통해서 어떠한 node를 남겨두고, 어떠한 node를 제거할 것인지 정해야 하므로 시간이 다소 걸리는데, 해당 논문에서 이야기하는 $F_p$의 경우 그냥 왼쪽에서부터 차례대로 남겨두면 되기 때문에 추가적인 fine-tuning이 필요 없습니다. 저자들은 이러한 경량화 과정에서의 cost 절약도 하나의 이점이라고 언급하고 있습니다.

 

3. Knowledge Distillation OD(OD w/ KD)

 

 저자들은 이렇게 절약한 cost를 다른 연산 과정에 사용하려고 시도합니다. $\mathcal{P}$에서 임의의 두 원소 $p_1 < p_2$를 선택할 경우, 왼쪽에서부터 차례대로 node를 살려두는 OD의 특성 상 $F_{p_1} \subsetneq F_{p_2}$를 만족하게 됩니다. (이를 OD의 nested structure라고 이야기합니다.) 즉, 더 큰 $p$를 사용한 model은, 더 작은 $p$를 사용한 model을 완전히 포함합니다. 저자들은 이 구조에서 착안하여 $F_{\max{\mathcal{P}}}$를 teacher로, $F_p$(단, $p < \max \mathcal{P}$)를 student로 두고 KD를 진행하는 메커니즘을 고안하였는데, 이를 Knowledge Distillation OD로 명명합니다. (구분이 필요할 경우, 앞서 이야기한 OD를 plain OD라고 부릅니다.) 이 경우, $F_p$의 loss function은 다음과 같이 정의됩니다. (여기에서, $\text{SM}_p$는 $F_p$의 softmax output, $\text{CE}$는 cross-entropy function, $\text{KL}$은 Kullback-Leibler divergence이고, $\alpha$는 오른쪽 두 term 간의 relative weight, $T$는 distillation temperature입니다.)

$$\mathcal{L}_d (\text{SM}_p, \text{SM}_{\max \mathcal{P}}, y_{label}) = (1 - \alpha) \text{CE} (\max(\text{SM}_p), y_{label}) + \alpha \text{KL} (\text{SM}_p, \text{SM}_{\max \mathcal{P}}, T)$$

 

 다음 포스트에서 OD의 특성을 계속해서 살펴보도록 하겠습니다.