Federated Feature Augmentation and Alignment

Federated learning is a distributed paradigm that allows multiple parties to collaboratively train deep learning models without direct exchange of raw data. Nevertheless, the inherent non-independent and identically distributed (non-i.i.d.) nature of data distribution among clients results in significant degradation of the acquired model. The primary goal of this study is to develop a robust federated learning algorithm to address <i>feature shift</i> in clients’ samples, potentially arising from a range of factors such as acquisition discrepancies in medical imaging. To reach this goal, we first propose federated feature augmentation (<small>FedFA</small><inline-formula><tex-math notation="LaTeX">$^{l}$</tex-math></inline-formula>), a novel feature augmentation technique tailored for federated learning. <small>FedFA</small><inline-formula><tex-math notation="LaTeX">$^{l}$</tex-math></inline-formula> is based on a crucial insight that each client's data distribution can be characterized by first-/second-order statistics (<i>a.k.a.</i>, mean and standard deviation) of latent features; and it is feasible to manipulate these local statistics <i>globally</i>, i.e., based on information in the entire federation, to let clients have a better sense of the global distribution across clients. Grounded on this insight, we propose to augment each local feature statistic based on a normal distribution, wherein the mean corresponds to the original statistic, and the variance defines the augmentation scope. Central to <small>FedFA</small><inline-formula><tex-math notation="LaTeX">$^{l}$</tex-math></inline-formula> is the determination of a meaningful Gaussian variance, which is accomplished by taking into account not only biased data of each individual client, but also underlying feature statistics represented by all participating clients. Beyond consideration of <i>low-order</i> statistics in <small>FedFA</small><inline-formula><tex-math notation="LaTeX">$^{l}$</tex-math></inline-formula>, we propose a federated feature alignment component (<small>FedFA</small><inline-formula><tex-math notation="LaTeX">$^{h}$</tex-math></inline-formula>) that exploits <i>higher-order</i> feature statistics to gain a more detailed understanding of local feature distribution and enables explicit alignment of augmented features in different clients to promote more consistent feature learning. Combining <small>FedFA</small><inline-formula><tex-math notation="LaTeX">$^{l}$</tex-math></inline-formula> and <small>FedFA</small><inline-formula><tex-math notation="LaTeX">$^{h}$</tex-math></inline-formula> yields our full approach <small><b>FedFA<inline-formula><tex-math notation="LaTeX">$+$</tex-math><alternatives><mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"><mml:mo>+</mml:mo></mml:math><inline-graphic xlink:href="yuan-ieq8-3457751.gif" xmlns:xlink="http://www.w3.org/1999/xlink"/></alternatives></inline-formula></b></small>. <small>FedFA<inline-formula><tex-math notation="LaTeX">$+$</tex-math><alternatives><mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"><mml:mo>+</mml:mo></mml:math><inline-graphic xlink:href="yuan-ieq9-3457751.gif" xmlns:xlink="http://www.w3.org/1999/xlink"/></alternatives></inline-formula></small> is non-parametric, incurs negligible additional communication costs, and can be seamlessly incorporated into popular CNN and Transformer architectures. We offer rigorous theoretical analysis, as well as extensive empirical justifications to demonstrate the effectiveness of the algorithm.