SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ziqing Fan $^{1,2}$ , Ruipeng Zhang $^{1,2}$ , Jiangchao Yao $^{1,2}$ , Bo Han $^{3}$ , Ya Zhang $^{1,2}$ , Yanfeng Wang $^{1,2,\text{图}}$
$^{1}$ Cooperative Medianet Innovation Center, Shanghai Jiao Tong University,
$^{2}$ Shanghai AI Laboratory,  $^{3}$ Hong Kong Baptist University
{zqfan_knight, zhangrp, sunarker}@sjtu.edu.cn
bhanml@comp.hkbu.edu.hk, {ya_zhang, wangyanfeng}@sjtu.edu.cn
Partially class-disjoint data (PCDD), a common yet under-explored data formation where each client contributes a part of classes (instead of all classes) of samples, severely challenges the performance of federated algorithms. Without full classes, the local objective will contradict the global objective, yielding the angle collapse problem for locally missing classes and the space waste problem for locally existing classes. As far as we know, none of the existing methods can intrinsically mitigate PCDD challenges to achieve holistic improvement in the bilateral views (both global view and local view) of federated learning. To address this dilemma, we are inspired by the strong generalization of simplex Equiangular Tight Frame (ETF) on the imbalanced data, and propose a novel approach called FedGELA where the classifier is globally fixed as a simplex ETF while locally adapted to the personal distributions. Globally, FedGELA provides fair and equal discrimination for all classes and avoids inaccurate updates of the classifier, while locally it utilizes the space of locally missing classes for locally existing classes. We conduct extensive experiments on a range of datasets to demonstrate that our FedGELA achieves promising performance (averaged improvement of  $3.9\%$  to FedAvg and  $1.5\%$  to best baselines) and provide both local and global convergence guarantees. Source code is available at: https://github.com/MediaBrain-SJTU/FedGELA.
Partially class-disjoint data (PCDD) [13, 18, 21] refers to an emerging situation in federated learning [14, 22, 43, 46, 50] where each client only possesses information on a subset of categories, but all clients in the federation provide the information on the whole categories. For instance, in landmark detection [39] for thousands of categories with data locally preserved, most contributors only have a subset of categories of landmark photos where they live or traveled before; and in the diagnosis of Thyroid diseases, due to regional diversity different hospitals may have shared and distinct Thyroid diseases [10]. It is usually difficult for each party to acquire the full classes of samples, as the participants may be lack of domain expertise or limited by demographic discrepancy. Therefore, how to efficiently handle the partially class-disjoint data is a critical (yet under-explored) problem in real-world federated learning applications for the pursuit of personal and generic interests.
Prevalent studies mainly focus on the general heterogeneity without specially considering the PCDD challenges: generic federated leaning (G-FL) algorithms adopt a uniform treatment of all classes and mitigate personal differences by imposing constraints on local training [17, 19], modifying logits [21, 47] adjusting the weights of submitted gradients [37] or generating synthetic data [54]; in contrast, personalized federated learning (P-FL) algorithms place relatively less emphasis on locally missing classes and selectively share either partial network parameters [1, 6] or class prototypes [33] to minimize the impact of personal characteristics, thereby separating the two topics. Those methods
Figure 1: Illustration of feature spaces and classifier vectors trained on the global dataset, two partially class-disjoint datasets (A and B), and restricted by federated algorithms. (a) is trained on the globally balanced dataset with full classes. (b) and (c) are trained on datasets A and B, respectively, which suffer from different patterns of classifier angle collapse problems. (d) is averaged in the server or constrained by some federated algorithms.
might directly or indirectly help mitigate the data shifts caused by PCDD, however, as far as we know, none of the existing works can mitigate the PCDD challenges to achieve holistic improvement in the bilateral views (global and local views) of federated learning. Please refer to Table 1 for a comprehensive comparison among a range of FL methods from different aspects.
Without full classes, the local objective will contradict the global objective, yielding the angle collapse for locally missing classes and the waste of space for locally existing classes. Ideally, as shown in Figure 1(a), global features and their corresponding classifier vectors shall maintain a proper structure to pursue the best separation of all classes. However, the angles of locally missing classes' classifier vectors will collapse, when trained on each client with partially class-disjoint data, as depicted in Figure 1(b), 1(c). FedRS [21] notices the degenerated updates of the classifier and pursues the same symmetrical structure in the local by restricting logits of missing classes. Other traditional G-FL algorithms indirectly restrict the classifier by constraining logits, features, or model weights, which may also make effects on PCDD. However, they cause another problem: space waste for personal tasks. As shown in Figure 1(d), restricting local structure will waste feature space and limit the training of the local model on existing classes. P-FL algorithms utilize the wasted space by selectively sharing part of models but exacerbate the angle collapse of classifier vectors. Recent FedRod [3] attempts to bridge the gap between P-FL and G-FL by introducing a two-head framework with logit adjustment in the G-head, but still cannot address the angle collapse caused by PCDD.
To tackle the PCDD dilemma from both P-FL and G-FL perspectives, we are inspired by a promising classifier structure, namely simplex equiangular tight frame (ETF) [9, 26, 41], which provides each class the same classification angle and generalizes well on imbalanced data. Motivated by its merits, we propose a novel approach, called FedGELA, in which the classifier is Globally fixed as a simplex ETF while Locally Adapted to personal tasks. In the global view, FedGELA merges class features and their corresponding classifier vectors, which converge to ETF. In the local view, it provides existing major classes with larger feature spaces and encourages to utilize the spaces wasted by locally missing classes. With such a bilateral curation, we can explicitly alleviate the impact caused by PCDD. In a nutshell, our contributions can be summarized as the following three points:
- We study a practical yet under-explored data formation in real-world applications of federated learning, termed as partially class-disjoint data (PCDD), and identify the angle collapse and space waste challenges that cannot be efficiently solved by existing prevalent methods (Sec. 3.2).
- We propose a novel method called FedGELA that classifier is globally fixed as a symmetrical structure ETF while locally adapted by personal distribution (Sec. 3.3), and theoretically show the local and global convergence analysis for PCDD with the experimental verification (Sec. 4.2).
- We conduct a range of experiments on multiple benchmark datasets under the PCDD case and a real-world dataset to demonstrate the bilateral advantages of FedGELA over the state-of-the-art methods from multiple views like the larger scale of clients and straggler situations (Sec. 5.2). We also provide further analysis like classification angles during training and ablation study. (Sec. 5.3).
2.1 Partially Class-Disjoint Data and Federated Learning algorithms
Partially class-disjoint data is one common formation among clients that can significantly impede the convergence, performance, and efficiency of algorithms in FL [18]. It belongs to the data heterogeneity
Table 1: Key differences between SOTA methods and our FedGELA categorized by targets (P-FL or G-FL), techniques (improve from the views of features, logits or model), and whether directly mitigate angle collapse of classifier vectors or save locally wasted feature spaces caused by PCDD.
Target	Research work	Feature View	Logit View	Model View	Mitigate Collapse	Save Space
G-FL	FedProx	-	-	✓	✓	-
MOON	✓	-	-	-	-
FedRS	-	✓	-	✓	-
FedGen	✓	-	-	✓	-
FedLC	-	✓	-	-	-
P-FL	FedRep	✓	-	✓	-	✓
FedProto	✓	-	✓	-	✓
FedBABU	✓	-	✓	-	✓
G&P-FL	FedRod	-	✓	✓	-	✓
FedGELA(ours)	✓	✓	✓	✓	✓
case, but does have a very unique characteristic different from the ordinary heterogeneity problem. That is, if only each client only has a subset of classes, it does not share the optimal Bayes classifier with the global model that considers all classes on the server side. Recently, FedRS [21] has recognized the PCDD dilemma and directly mitigate the angle collapse issue by constraining the logits of missing classes. FedProx [19] also can lessen the collapse by constraining local model weights to stay close to the global model. Other G-FL algorithms try to address data heterogeneity from a distinct perspective. MOON [17] and FedGen [54] utilizes contrastive learning and generative learning to restrict local representations. And FedLC [47] introduces logit calibration to adjust the logits of the local model to match those of the global model, which might indirectly alleviate the angle collapse in the local. However, they all try to restrict local structure as global, resulting in the waste space for personal tasks shown in Figure 1(d). P-FL algorithms try to utilize the wasted space by encouraging the angle collapse of the local classifier. FedRep [6] only shares feature extractors among clients and FedProto [33] only submits class prototypes to save communication costs and align the feature spaces. In FedBABU [25], the classifier is randomly initialized and fixed during federated training while fine-tuned for personalization during the evaluation. However, they all sacrifice the generic performance on all classes. FedRod [3] attempts to bridge this gap by introducing a framework with two heads and employing logit adjustment in the global head to estimate generic distribution but cannot address angle collapse. In Table 1, we categorize these methods by targets (P-FL or G-FL), skews (feature, logit, or model weight), and whether they directly mitigate the angle collapse of local classifier or saving personal spaces for personal spaces. It is evident that none of these methods, except ours, can handle the PCDD problem in both P-FL and G-FL. Furthermore, FedGELA is the only method that can directly achieve improvements from all views.
2.2 Simplex Equiangular Tight Frame
The simplex equiangular tight frame (ETF) is a phenomenon observed in neural collapse [26], which occurs in the terminal phase of a well-trained model on a balanced dataset. It is shown that the last-layer features of the model converge to within-class means, and all within-class means and their corresponding classifier vectors converge to a symmetrical structure. To analyze this phenomenon, some studies simplify deep neural networks as last-layer features and classifiers with proper constraints (layer-peeled model) [9, 12, 40, 53] and prove that ETF emerges under the cross-entropy loss. However, when the dataset is imbalanced, the symmetrical structure of ETF will collapse [9]. Some studies try to obtain the symmetrical feature and the classifier structure on the imbalanced datasets by fixing the classifier as ETF [41, 53]. Inspired by this, we propose a novel method called FedGELA that bilaterally curates the classifier to leverage ETF or its variants. See Appendix A for more details about ETF.
ETF under LPM. A typical L-layer DNN parameterized by  $\mathbf{W}$  can be divided into the feature backbone parameterized by  $\mathbf{W}^{-L}$  and the classifier parameterized by  $\mathbf{W}^{L}$ . From the view of layer-peeled model (LPM) [9, 12, 40, 53], training  $\mathbf{W}$  with constraints on the weights can be considered
as training the C-class classifier  $\mathbf{W}^L = \{\mathbf{W}_1^L,\dots,\mathbf{W}_C^L\}$  and features  $\mathbf{H} = \{h^{1},\dots,h^{n}\}$  of all  $n$  samples output by last layer of the backbone with constraints  $E_{W}$  and  $E_{H}$  on them respectively. On the balanced data, any solutions to this model form a simplex equiangular tight frame (ETF) that all last layer features  $h_c^{i,*}$  and corresponding classifier  $\mathbf{W}_c^{L,*}$  of all classes converge as:
$$
\frac {h _ {c} ^ {i , *}}{\sqrt {E _ {H}}} = \frac {\mathbf {W} _ {c} ^ {L , *}}{\sqrt {E _ {W}}} = m _ {c} ^ {*}, \tag {1}
$$
where  $m_c^*$  forms the ETF defined as  $\mathbf{M} = \sqrt{\frac{C}{C - 1}}\mathbf{U}\left(\mathbf{I}_C - \frac{1}{C}\mathbf{1}_C\mathbf{1}_C^T\right)$ . Here  $\mathbf{M} = [m_1^*,\dots ,m_C^* ]\in$ $\mathbb{R}^{d\times C}$ ,  $\mathbf{U}\in \mathbb{R}^{d\times C}$  allows a rotation and satisfies  $\mathbf{U}^T\mathbf{U} = \mathbf{I}_C$  and  $\mathbf{1}_C$  is an all-ones vector. ETF is an optimal classifier and feature structure in the balanced case of LPM.
FedAvg. On the view of LPM, given N clients and each with  $n_k$  samples, the vanilla federated learning via FedAvg consists of four steps [22]: 1) In round  $t$ , the server broadcasts the global model  $\mathbf{W}^t = \{\mathbf{H}^t,\mathbf{W}^{L,t}\}$  to clients that participate in the training (Note that here  $\mathbf{H}$  is actually the global backbone  $\mathbf{W}^{-L,t}$  instead of real features); 2) Each local client receives the model and trains it on the personal dataset. After  $E$  epochs, we acquire a new local model  $\mathbf{W}_k^t$ ; 3) The updated models are collected to the server as  $\{\mathbf{W}_1^t,\mathbf{W}_2^t,\dots ,\mathbf{W}_N^t\}$ ; 4) The server averages local models to acquire a new global model as  $\mathbf{W}^{t + 1} = \sum_{k = 1}^{N}p_{k}\mathbf{W}_{k}^{t}$ , where  $p_k = n_k / \sum_{k' = 1}^{N}n_{k'}$ . When the pre-defined maximal round  $T$  reaches, we will have the final optimized global model  $\mathbf{W}^T$ .
3.2 Contradiction and Motivation
Contradiction. In G-FL, the ideal global objective under LPM of federated learning is described as:
$$
\min  _ {\mathbf {H}, \mathbf {W} ^ {L}} \sum_ {k = 1} ^ {N} p _ {k} \frac {1}{n _ {k}} \sum_ {c \in C _ {k}} \sum_ {i = 1} ^ {n _ {k, c}} \mathcal {L} _ {C E} \left(h _ {k, c} ^ {i}, \mathbf {W} ^ {L}\right).
$$
Assuming global distribution is balanced among classes, no matter whether local datasets have full or partial classes, the global objective with constraints on weights can be simplified as:
$$
\min  _ {\mathbf {H}, \mathbf {W} ^ {L}} \frac {1}{n} \sum_ {c = 1} ^ {C} \sum_ {i = 1} ^ {n _ {c}} \mathcal {L} _ {C E} \left(h _ {c} ^ {i}, \mathbf {W} ^ {L}\right), \text {s . t .} \left\| \mathbf {W} _ {c} ^ {L} \right\| ^ {2} \leqslant E _ {W}, \left\| h _ {c} ^ {i} \right\| ^ {2} \leqslant E _ {H}. \tag {2}
$$
Similarly, the local objective of k-th client with a set of classes  $C_k$  can be described as:
$$
\min  _ {\mathbf {H} _ {k}, \mathbf {W} _ {k} ^ {L}} \frac {1}{n _ {k}} \sum_ {c \in C _ {k}} \sum_ {i = 1} ^ {n _ {k, c}} \mathcal {L} _ {C E} \left(h _ {k, c} ^ {i}, \mathbf {W} _ {k} ^ {L}\right), \text {s . t .} \left\| \mathbf {W} _ {k, c} ^ {L} \right\| ^ {2} \leqslant E _ {W}, \left\| h _ {k, c} ^ {i} \right\| ^ {2} \leqslant E _ {H}. \tag {3}
$$
When PCDD exists ( $C_k \neq C$ ), we can see the contradiction between local and global objectives, which respectively forms two structures, shown in Figure 3(a) and Figure 3(b). After aggregated in server or constrained by some FL methods, the structure in the local is restricted to meet the global structure, causing space waste for personal tasks shown in Figure 1(d).
Motivation. To verify the contradiction and related feature and classifier structures, we split CIFAR10 into 10 clients and perform FedAvg on it with Dirichlet Distribution (Dir  $(\beta = 0.1)$ ). As illustrated in Figure 2, the angle difference between existing classes and between missing classes becomes smaller and converges to a similar value in the global model. However, in the local training, angles between existing classes become larger while angles between missing classes become smaller, which indicates the contradiction. With this observation, to bridge the gap between Eq (3) and Eq (2) under PCDD, we need to construct the symmetrical and uniform classifier angles for all classes while encouraging local clients to expand existing classes' feature space. Therefore, we propose our method FedGELA that classifier can be Globally fixed as ETF but Locally Adapted based on the local distribution matrix to utilize the wasted space for the existing classes.
Figure 2: Averaged angles of classifier vectors between locally existing classes (existing angle) and between locally missing classes (missing angle) on CIFAR10 (Dir  $(\beta = 0.1)$ ) in local client and aggregated in global server (local epoch is 10). In global, "existing" angle and "missing" angle converge to similar values while in the local, "existing" angle expands but "missing" angle shrinks.
3.3 FedGELA
Global ETF. Given the global aim of achieving an unbiased classifier that treats all classes equally and provides them with the same discrimination and classifier angles, we curate the global model's classifier as a randomly initialized simplex ETF with scaling  $\sqrt{E_w}$  at the start of federated training:
$$
\mathbf {W} ^ {L} = \sqrt {E _ {w}} \mathbf {M}.
$$
Then the ETF is distributed to all clients to replace their local classifiers. In Theorem 1, we prove in federated training under some basic assumptions, by fixing the classifier as a randomly simplex ETF with scaling  $\sqrt{E_W}$  and constraints  $E_H$  on the last layer features, features output by last layer of backbone and their within class means will converge to the ETF similar to Eq (1), which meets the requirement of global tasks.
Local Adaptation. However, when PCDD exists in the local clients, naively combining ETF with FL does not meet the requirement of P-FL as analyzed in Eq (2) and Eq (3). To utilize the wasted space for locally missing classes, in the training stage, we curate the length of ETF received from the server based on the local distribution as below:
$$
\mathbf {W} _ {k} ^ {L} = \boldsymbol {\Phi} _ {k} \mathbf {W} ^ {L} = \boldsymbol {\Phi} _ {k} \sqrt {E _ {w}} \mathbf {M}, \tag {4}
$$
where  $\Phi_k$  is the distribution matrix of k-th client. Regarding the selection of  $\Phi_k$ , it should satisfy a basic rule for federated learning, wherein the aggregation of local classifiers aligns with the global classifier, thereby ensuring the validity of theoretical analyses from both global and local perspectives. Moreover, it is highly preferable for the selection process to avoid introducing any additional privacy leakage risks. To meet the requirement that averaged classifier should be standard ETF:  $\mathbf{W}^L = \sum_{k=1}^{N} p_k \mathbf{W}_k^L$  in the globally balanced case, its row vectors are all one's vector multiple statistical values of personal distribution:  $(\Phi_k^T)_c = \frac{n_{k,c}}{n_k \gamma} \mathbf{1}$  ( $\gamma$  is a constant, and  $n_{k,c}$  and  $n_k$  are the c-th class sample number and total sample number of the k-th client) respectively. We set  $\gamma$  to  $\frac{1}{|C|}$ . Finally, the local objective from Eq. (3) is adapted as:
$$
\min  _ {\mathbf {H} _ {k}} \frac {1}{n _ {k}} \sum_ {c = 1} ^ {C} \sum_ {i = 1} ^ {n _ {k, c}} - \log \frac {\exp \left(\Phi_ {k , c} \mathbf {W} _ {c} ^ {L ^ {T}} h _ {k , c} ^ {i}\right)}{\sum_ {c ^ {\prime} \in C _ {k}} \exp \left(\Phi_ {k , c ^ {\prime}} \mathbf {W} _ {c ^ {\prime}} ^ {L ^ {T}} h _ {k , c} ^ {i}\right)}, \tag {5}
$$
$$
\begin{array}{l l} \text {s . t .} & \| h ^ {i} \| ^ {2} \leqslant E _ {H}, \forall 1 \leqslant i \leqslant n _ {k}. \end{array}
$$
Total Framework. After introducing two key parts of FedGELA (Global ETF and Local Adaptation), we describe the total framework of FedGELA. As illustrated and highlighted in Algorithm 1 (refer to Appendix D for the workflow figure), at the initializing stage, the server randomly generates an ETF as the global classifier and sends it to all clients while local clients adjust it based on the personal distribution matrix as Eq (4). At the training stage, local clients receive global backbones and train with adapted ETF in parallel. After  $E$  epochs, all clients submit personal backbones to the server. In the server, personal backbones are received and aggregated to a generic backbone, which is broadcast to all clients participating in the next round. At the inference stage, on the client side, we obtain a generic backbone with standard ETF to handle the world data while on the client side, a personal backbone with adapted ETF to handle the personal data.
Algorithm 1 FedGELA
Input:  $(N,K,n_{k},c_{k},\mathbf{H}^{0},\mathbf{M},E_{W},E_{H},T,\eta ,E)$
Parallely for all clients:  $\mathbf{W}_k^L\gets \Phi_k\sqrt{E_W}\mathbf{M}.$
for  $t = 0,1,\dots ,T - 1$  do
$\triangleright$  on the server side
$\mathbf{H}^t\gets \sum_{k = 1}^{K}p_k^t H_k^{t - 1}.$
sample K clients from all N clients.
$\triangleright$  on the client side
do in parallel for  $\forall k\in K$  clients
receive  $\mathbf{H}^t$  from server,  $\mathbf{H}_k^t\gets \mathbf{H}^t$
for  $\tau = 0,1,\dots,E - 1$  do
sample a mini-batch  $b_{k}^{tE + \tau}$  in local data.
$$
H _ {k} ^ {t} \leftarrow H _ {k} ^ {t} - \eta \nabla F _ {k} \left(b _ {k} ^ {t}, \Phi_ {k} W _ {g} ^ {L}; \mathbf {H} _ {k} ^ {t}\right)
$$
submit  $\mathbf{H}_k^t$  to the server.
end in parallel
Output:  $(\mathbf{H}^T, W_g^L)$  and  $(\mathbf{H}_k^T, \Phi_k W_g^L)$ .
In this part, we first primarily introduce some notations and basic assumptions in Sec. 4.1 and then present the convergence guarantees of both local models and the global model under the PCDD with the proper empirical justification and discussion in Sec. 4.2. (Please refer to Appendix B for entire proofs and Appendix D for details on justification experiments.)
(a) Global convergence verification.
(b) Local convergence verification.
(c)Adapted structure.
Figure 3: Illustration of local and global convergence verification together with the effect of  $\Phi$ . (a) and (b) are the results of averaged angle between all class means and between locally existing class means in FedAvg, FedGE, and FedGELA on CIFAR10 under 50 clients and Dir ( $\beta = 0.2$ ). (c) is the illustration of how local adaptation utilizes the wasted space of missing classes for existing classes.
4.1 Notations
We use  $t$  and  $T$  to denote a curtain round and pre-defined maximum round after aggregation in federated training,  $tE$  to denote the state that just finishing local training before aggregation in round  $t$ , and  $tE + \tau$  to denote  $\tau$ -th local iteration in round  $t$  and  $0 \leq \tau \leq E - 1$ . The convergence follows some common assumptions in previous FL studies and helpful math results [15, 20, 29-31, 33, 36, 38, 45, 51] including smoothness, convexity on loss function  $F_{1}, F_{2}, \dots, F_{N}$  of all clients, bounded norm and variance of stochastic gradients on their gradient functions  $\nabla F_{1}, \nabla F_{2}, \dots, \nabla F_{N}$  and heterogeneity  $\Gamma_{1}$  reflected as the distance between local optimum  $\mathbf{W}_{k}^{*}$  and global optimum  $\mathbf{W}^{*}$ . Please refer to the concrete descriptions of those assumptions in Appendix B. Besides, in Appendix B, we additionally provide a convergence guarantee without a bounded norm of stochastic gradients, as some existing works [24, 32] point out the contradiction to the strongly convex.
4.2 Convergence analysis
Here we provide the global and local convergence guarantee of our FedGELA compared with FedAvg and FedGE (FedAvg with only the Globally Fixed ETF) in Theorem 1 and Theorem 2. To better explain the effectiveness of our FedGELA in local and global tasks, we record the averaged angle between all class means in global and existing class means in local as shown in Figure 3(a) and Figure 3(b). Please refer to Appendix B for details on the proof and justification of theorems.