SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathbb {E} [ \| \widetilde {\nabla} _ {\Theta_ {g}} \ell (\Theta_ {g}, \Theta_ {k}) - \nabla_ {\Theta_ {g}} \ell (\Theta_ {g}, \Theta_ {k}) \| ^ {2} ] \leq \sigma_ {g} ^ {2}, \mathbb {E} [ \| \widetilde {\nabla} _ {\Theta_ {k}} \ell (\Theta_ {g}, \Theta_ {k}) - \nabla_ {\Theta_ {k}} \ell (\Theta_ {g}, \Theta_ {k}) \| ^ {2} ] \leq \sigma_ {k} ^ {2}.
$$
Assumption 3. (Partial gradient diversity). There exists a constant for all  $\theta_i^0$  and  $\Theta_g$ ,  $\theta_i$  and  $\Theta_k$ ,
$$
\sum_ {i = 1} ^ {m} \frac {n _ {i}}{n} \| \nabla_ {\Theta_ {g}} \ell (\Theta_ {g}, \theta_ {i}) - \nabla_ {\Theta_ {g}} \ell (\Theta_ {g}, \Theta_ {k}) \| ^ {2} \leq \delta^ {2}
$$
$$
\sum_ {i \in [ k ]} \frac {n _ {i}}{\sum_ {i \in [ k ]} n _ {i}} \| \nabla_ {\Theta_ {k}} \ell (\theta_ {i} ^ {0}, \Theta_ {k}) - \nabla_ {\Theta_ {k}} \ell (\Theta_ {g}, \Theta_ {k}) \| ^ {2} \leq \delta^ {2}.
$$
Assumption 4. (Convexity of cluster models). Fix  $\Theta_g$ , assume  $\ell(\Theta_k)$  is convex.
Theorem 1. (Convergence of Fed-CAM). Let Assumptions 1, 2, 3 and 4 hold, and learning rates chosen as  $\eta = \tau / (LE)$  for a  $\tau$  depending on the parameters  $L$ ,  $\sigma_g^2$ ,  $\sigma_k^2$ ,  $\delta^2$ ,  $s$ ,  $m$ ,  $T$ , provided clustering stable, we have (ignoring absolute constants),
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 1} ^ {T} \left(\frac {1}{L} \mathbb {E} \left[ \| \nabla_ {\Theta_ {g}} \ell \left(\Theta_ {g}, \Theta_ {k}\right) \| ^ {2} \right] + \frac {s}{m L} \frac {1}{m} \sum_ {i = 1} ^ {m} \mathbb {E} \left[ \| \nabla_ {\Theta_ {c (i)}} \ell \left(\Theta_ {g}, \Theta_ {c (i)}\right) \| ^ {2} \right]\right) (17) \\ \leq \frac {\left(\triangle_ {\ell} \sigma_ {s i m , 1} ^ {2}\right) ^ {1 / 2}}{T ^ {1 / 2}} + \frac {\left(\triangle_ {\ell} ^ {2} \sigma_ {s i m , 2} ^ {2}\right) ^ {1 / 3}}{T ^ {2 / 3}} + O \left(\frac {1}{T}\right), (18) \\ \end{array}
$$
where  $\triangle_{\ell} = \ell_0 - \ell^{\star}$ , and we define effective variance terms,
$$
\sigma_ {s i m, 1} ^ {2} = \frac {2}{L} \left(\delta^ {2} \left(1 - \frac {s}{m}\right) + \frac {\sigma_ {g} ^ {2}}{L} + \frac {\sigma_ {k} ^ {2} s}{m}\right) \tag {19}
$$
$$
\sigma_ {s i m, 2} ^ {2} = \frac {2}{L} \left(\delta^ {2} + \sigma_ {g} ^ {2} + \sigma_ {k} ^ {2}\right) \left(1 - \frac {1}{E}\right). \tag {20}
$$
Remark 1. It is straightforward to prove that the clustering of both IFCA-CAM and FeSEM-CAM converges, as evidenced by Ma et al. (2022). However, proving the stability of these clustering methods is more challenging due to the oscillation phenomenon often seen in K-means. The stability of clustering will be further demonstrated through experimental analysis in Section 5.2.
Remark 2. Besides the clustering structure, there is a distinct difference between FedSim [27] and Fed-CAM. In Fed-CAM, we need to aggregate both  $\Theta_{g}$  and  $\Theta_{1:K}$ , while in FedAlt [27], only  $\Theta_{g}$  requires aggregation. The  $\sigma_{sim,1}^{2}$  and  $\sigma_{sim,2}^{2}$  reflect the impact of sample number  $s$  and local steps  $E$ . Larger  $s$  or smaller  $E$ , better convergence rate. According to the results presented in [27], alternative gradient descent surpasses simultaneous gradient descent in terms of convergence rate. The asymptotic  $1/\sqrt{T}$  rate is achieved when each device is seen at least once on average, and the  $1/T$  term is dominated by the  $1/\sqrt{T}$  term, a situation that occurs when (ignoring absolute constants)
$$
T \geq \frac {\triangle_ {\ell}}{\sigma_ {s i m , 1} ^ {2}} \max \{(1 - \frac {1}{E}) \frac {m}{s}, 2 \}.
$$
Benchmark datasets and partitions. The proposed methods have been validated using several benchmark datasets. While detailed results for PathMNIST and TissueMNIST from the MedMNIST [28] are provided in the supplementary material, the other datasets used for validation include:
- Fashion-MNIST [29] includes 70,000 labeled fashion images  $(28 \times 28$  grayscale) in 10 classes, such as T-shirts, Trouser, and Bag, with others.
- CIFAR-10 [30] consists of 60,000 images  $(32 \times 32$  color) in 10 classes, including airplane, automobile, bird, and truck, among others. The divergence among classes in CIFAR-10 is relatively higher than in other datasets from the MNIST family.
Each dataset is split among 200 clients and we create the following non-IID scenarios:
- Client-wise non-IID by Dirichlet distribution ( $\alpha = 0.1$ ): This approach uses the Dirichlet distribution to control the randomness of non-IID data, as proposed by [31]. This is a standard method used in most personalized FL methods, which are usually client-wise non-IID.
- Cluster-wise non-IID by Dirichlet distribution  $(\alpha = (0.1, 10))$ : This strategy divides the dataset into  $K$  clusters with  $\alpha = 0.1$  to generate substantial variance in cluster-wise non-IID. Then, each cluster is divided into  $m / K$  clients with  $\alpha = 10$  to control the non-IID across clients.
- Client-wise non-III by n-class (2): This method randomly selects  $n$ -class out of all classes in the dataset for each client, as proposed by [1], and then samples instances from these datasets.
- Cluster-wise non-III by n-class (3, 2): This approach randomly assigns 3 classes to each cluster, ensuring a relatively balanced number of instances per class. It then assigns 2 classes to each client.
Baselines. We select baseline methods from four categories as follows:
- Single model-based FL: We choose FedAvg [1] and FedProx [3] with a coefficient of 230 and a regularization of 0.95 as the baselines.
- Ensemble FL: We train FedAvg and FedProx  $K$  times and then learn an ensemble model via soft voting to serve all clients, which are named FedAvg+ and FedProx+, respectively.
- Clustered FL: We choose FeSEM [8] and IFCA [6], which is similar to HypCluster [7].
- Clustered FL with additive modeling: We integrate CAM with IFCA and FeSEM, denoting them as IFCA-CAM and FeSEM-CAM, respectively.
Learning-related hyperparameters. We use the Convolutional Neural Network (CNN) [32] as the basic model architecture for each client, as detailed in the supplementary material. For optimization, we employ SGD with a learning rate of 0.001 and momentum of 0.9 to train the model, and the batch size is 32. We evaluate the performance using both micro accuracy  $(\%)$  and macro F1-score  $(\%)$  on the client-wise test datasets to better capture the non-IID nature per client.
FL system settings. We conduct 100 global communication rounds in the FL system, including 30 warm-up rounds if applicable. Each communication involves 10 local steps. For the clustering process of FeSEM-CAM, we measure distance on the flattened parameters of the fully-connected layers, and use K-Means as the clustering algorithm. The coefficient  $\lambda$  is chosen from 0.001, 0.01, 0.1 based on the best performance.
5.1 Main Results and Comparisons
Cluster-wise non-III scenarios make the assumption that there are underlying clustering structures among clients. Table 1 compares the methods using two benchmark datasets, namely Fashion-MNIST and CIFAR-10. Results using two biomedical datasets are presented in the appendix. The following are some notable observations and analyses:
- The application of the ensemble mechanism to FedAvg and FedProx yields minor improvements. This is because the server-side model in FedAvg or FedProx is already a relatively strong model, while ensemble mechanisms usually excel with weak models.
- The introduction of CAM significantly enhances the performance of IFCA, which typically struggles with clustering collapse in cluster-wise non-IID scenarios. Notably, CAM decomposes the shared components into a global model and personalized parts into cluster models. Thus, the clustering collapse is mitigated by isolating the dominant shared knowledge.
- FeSEM generally exhibits robust performance on cluster-wise non-IID without outliers. Implementing CAM in FeSEM further improves the Macro-F1 performance. The clustering process in FeSEM tends to overfit the label distribution (imbalanced classes) of clients to achieve higher accuracy. However, the application of CAM introduces a global model with a balanced label distribution by averaging all clients, thereby boosting the Macro-F1 performance while preserving the cluster-wise non-IID for high accuracy.
- With an increase in the number of clusters  $K$ , the CAM-based methods show substantial improvements in Macro-F1. The decomposition of shared knowledge and cluster-wise non-IID characteristics benefit from a reasonably larger  $K$ , which facilitates fine-grained, cluster-wise personalization.
Table 1: Test results (mean±std) in cluster-wise non-IID settings on Fashion-MNIST & CIFAR-10.
Datasets	Fashion-MNIST	CIFAR-10
Non-IID setting	Dirichlet α = (0.1, 10)	n-class (3, 2)	Dirichlet α = (0.1, 10)	n-class (3, 2)
#Cluster	Methods	Accuracy	Macro-F1	Accuracy	Macro-F1	Accuracy	Macro-F1	Accuracy	Macro-F1
1	FedAvg	86.08±0.70	57.24±2.26	86.33±0.44	46.09±1.08	24.38±3.30	11.69±3.15	21.33±3.83	9.00±0.58
FedProx	86.32±0.78	58.03±3.19	86.42±0.63	45.86±1.42	24.73±3.68	11.28±2.35	22.66±1.13	9.23±0.78
5	FedAvg+	87.61	59.48	86.95	65.61	25.97	12.16	24.35	9.06
FedProx+	87.94	59.83	86.52	65.73	26.05	12.53	24.83	9.31
IFCA	84.60±2.22	62.03±3.01	84.94±2.54	66.50±4.43	34.10±4.79	22.12±2.21	29.80±4.49	17.90±2.08
IFCA-CAM	93.33±0.95	79.64±4.09	95.38±0.49	77.56±1.14	58.13±3.82	28.09±3.68	54.56±3.58	27.27±1.06
FeSEM	94.64±1.54	82.90±2.38	94.20±1.96	77.07±6.05	59.06±3.24	32.33±7.25	58.76±3.35	35.75±2.54
FeSEM-CAM	95.13±1.78	85.10±3.17	95.69±1.05	78.82±1.17	64.35±2.33	38.33±1.77	65.58±1.21	38.63±1.17
10	FedAvg+	89.42	67.83	86.91	63.01	28.45	13.79	27.28	9.81
FedProx+	89.55	68.02	86.73	63.42	28.33	13.64	26.94	9.64
IFCA	82.10±5.40	62.62±8.22	86.58±4.97	66.22±5.69	34.84±5.82	22.76±3.99	34.06±2.60	18.70±1.31
IFCA-CAM	95.42±2.54	88.45±5.46	95.09±0.87	82.98±1.16	70.90±1.18	40.03±1.28	68.46±4.08	41.45±4.00
FeSEM	95.73±1.28	89.34±1.57	95.54±0.74	84.43±2.38	66.89±2.18	38.35±4.24	71.76±2.23	49.72±3.84
FeSEM-CAM	96.19±1.20	92.37±1.85	98.07±1.46	92.43±2.70	78.45±1.71	49.50±1.13	75.04±1.97	55.90±2.07
Client-wise non-IID Table 2 presents comparative results under client-wise non-IID scenarios using two benchmark datasets: Fashion-MNIST and CIFAR-10. Interestingly, IFCA maintains stable performance under client-wise non-IID conditions, primarily because it cannot form a single dominant cluster model - a primary cause of clustering collapse - in a highly heterogeneous environment. The application of CAM to IFCA and FeSEM shows a significant enhancement, particularly on the CIFAR-10 dataset. This improvement is likely due to FeSEM's typical restriction on knowledge sharing across clusters. In contrast, CAM utilizes a global model to capture more useful common
Figure 3: Cluster sizes during FeSEM vs. FeSEM+CAM in client/cluster-wise non-IID settings on CIFAR-10. Legend: cluster ID (cluster size) in the last round. CAM effectively mitigates clustering collapse/imbalance.
knowledge across clusters, thereby substantially enhancing the generalization capability of each cluster. Furthermore, CIFAR-10, being a relatively complex dataset with a diversity of images, underscores the importance of sharing common knowledge.
Table 2: Test results (mean±std) in client-wise non-IID settings on Fashion-MNIST & CIFAR-10.
Datasets	Fashion-MNIST	CIFAR-10
Non-IID setting	Dirichlet α = 0.1	n-class (2)	Dirichlet α = 0.1	n-class (2)
#Cluster	Methods	Accuracy	Macro-F1	Accuracy	Macro-F1	Accuracy	Macro-F1	Accuracy	Macro-F1
1	FedAvg	85.90±0.46	54.52±2.66	86.17±0.25	44.88±1.24	25.62±3.47	11.38±2.02	24.30±3.53	8.56±0.64
FedProx	86.03±0.58	54.69±3.32	86.47±0.23	44.89±1.38	25.72±3.29	11.14±1.49	24.19±2.45	8.69±0.74
5	FedAvg+	86.12	61.07	86.5	45.39	25.71	12.45	24.83	8.74
FedProx+	86.39	56.56	86.15	45.43	25.58	12.43	25.88	8.55
IFCA	90.13±6.81	68.47±5.23	91.54±5.04	72.30±5.32	47.21±10.28	22.67±1.48	46.54±12.8	17.78±1.29
IFCA-CAM	93.72±1.34	70.67±1.75	92.24±1.22	70.24±4.33	54.32±1.25	23.48±1.18	54.92±1.51	25.20±1.05
FeSEM	91.51±2.90	73.78±9.88	91.83±1.24	71.05±8.63	54.30±4.58	24.78±6.01	55.55±4.83	32.80±4.18
FeSEM-CAM	94.74±1.04	75.12±5.82	93.14±2.03	76.98±2.17	59.71±2.80	40.45±3.53	56.70±1.68	34.52±1.64
10	FedAvg+	86.81	60.43	86.91	47.12	27.83	13.65	27.71	9.65
FedProx+	86.24	56.2	86.78	42.83	25.86	12.84	26.16	9.94
IFCA	91.04±4.33	68.6±6.77	91.42±5.16	72.29±5.80	47.62±10.15	23.36±2.48	47.96±10.59	17.88±1.04
IFCA-CAM	95.70±1.19	79.17±1.91	92.57±2.63	76.31±4.39	72.54±2.7	42.86±4.36	61.01±2.41	31.63±2.17
FeSEM	93.3±2.0	80.47±11.05	93.75±1.53	79.39±6.57	67±1.57	31.69±8.52	63.64±6.51	42.97±6.08
FeSEM-CAM	95.25±1.93	81.5±2.24	95.15±1.48	86.16±3.19	80.11±1.82	59.19±4.67	69.88±1.7	49.5±1.42
5.2 Visualization: CAM combats clustering collapse
Figures 1 and 3 demonstrate the effectiveness of applying CAM to mitigate clustering collapse in IFCA and FeSEM under both cluster-wise and client-wise non-IID scenarios, using the CIFAR-10 dataset with  $K = 10$ . Each color represents a cluster, and the X-axis represents the iteration rounds.
In the case of IFCA, we observe a severe clustering collapse issue in cluster-wise non-IID scenarios. A single cluster can encompass all clients in the client-wise non-IID setting and up to  $50\%$  of clients in the cluster non-IID setting. Furthermore, the clustering remains unstable throughout the process. However, when CAM is applied in IFCA-CAM, it quickly identifies some clustering structures within a few rounds, and this structure closely approximates the ground truth.
As for FeSEM, while the phenomenon of clustering collapse is not as pronounced, a single cluster can still dominate up to  $25\%$  of all clients if there are no outliers. CAM can expedite the clustering convergence, sometimes achieving it in just one round. Moreover, under client-wise non-IID settings, the application of CAM results in lower variance and more uniform cluster size. In the case of cluster-wise non-IID settings, FeSEM-CAM can easily identify the ground truth.
We propose a novel structured FL model "clustered additive modeling (CAM)" and an efficient FL algorithmic framework Fed-CAM to address non-IID FL challenges with clustering structure. CAM is a general mode-agnostic tool that can improve various existing non-IID FL methods. It can capture more general non-IID structures with global knowledge sharing among clients than clustered FL and overcome several weaknesses such as clustering collapse, vulnerability to cluster imbalance/initialization, etc. Theoretically, Fed-CAM is capable of achieving an asymptotic convergence rate of  $O(1 / \sqrt{T})$ . Extensive experiments show that CAM brings substantial improvement to existing clustered FL methods, improves cluster balance, and effectively mitigates clustering collapse.
[24] Li Chou, Zichang Liu, Zhuang Wang, and Anshumali Shrivastava. Efficient and less centralized federated learning. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 772-787. Springer, 2021.
[25] Xueyang Tang, Song Guo, and Jingcai Guo. Personalized federated learning with contextualized generalization. In Lud De Raedt, editor, Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence, IJCAI-22, pages 2241–2247. International Joint Conferences on Artificial Intelligence Organization, 7 2022. Main Track.
[26] Alekh Agarwal, John Langford, and Chen-Yu Wei. Federated residual learning. arXiv preprint arXiv:2003.12880, 2020.
[27] Krishna Pillutla, Kshitiz Malik, Abdel-Rahman Mohamed, Mike Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. In International Conference on Machine Learning, pages 17716-17758. PMLR, 2022.
[28] Jiancheng Yang, Rui Shi, Donglai Wei, Zequan Liu, Lin Zhao, Bilian Ke, Hanspeter Pfister, and Bingbing Ni. Medmnist v2: A large-scale lightweight benchmark for 2d and 3d biomedical image classification. arXiv preprint arXiv:2110.14795, 2021.
[29] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017.
[30] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[31] Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06335, 2019.
[32] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436-444, 2015.