## 3 FedCKA
### Modifications to FedAvg
There are four main differences. First, SimCLR uses the representations of one model on different samples in a batch to calculate contrastive loss. FedCKA uses the representation of three models on the same samples in a batch to calculate \(\ell_{cka}\). \(a_{l_{i}}^{t}\), \(a_{l_{i}}^{t-1}\), and \(a_{g}^{t}\) are the representations of client \(i\)'s current local model, client \(i\)'s previous round local model, and the current global model, respectively. Second, SimCLR uses the temperature parameter \(\tau\) to increase performance on difficult samples. FedCKA excludes \(\tau\), as it was not seen to help performance. Third, SimCLR uses cosine similarity to measure the similarity between the representations of difference datasets. FedCKA uses CKA as its measure of similarity. Fourth, SimCLR calculates contrastive loss once per batch, using the representations of the projection head. Here, \(M\) represents the number of layers being regularized. FedCKA use calculates \(\ell_{cka}\)\(M\) times per batch, using the representations of the first \(M\) naturally similar layers, indexed by \(n\), and averages the loss based on the number of layers to regularize. \(M\) is set to two by default unless otherwise stated. As per [PERSON] et al. [10], CKA is shown in Equation (3). Here, the \(i^{th}\) eigenvalue of \(XX^{T}\) is \(\lambda_{X}^{i}\). While [PERSON] et al. [10] also presented a method to use kernels with CKA, we use the linear variant, as it is more computationally efficient, while having minimal impact on accuracy. \[ \begin{split} CKA(XX^{T},YY^{T})=\frac{\left\|Y^{T}X\right\|_{F}^ {2}}{\left\|X^{T}X\right\|_{F}\left\|Y^{T}Y\right\|_{F}}\\ =\frac{\sum_{i=1}^{p-1}\sum_{j=1}^{p2}\lambda_{X}^{i}\lambda_{Y}^ {i}\left\langle\mathbf{u}_{X^{\prime}}^{i}\mathbf{u}_{Y}^{j}\right\rangle^{2} }{\sqrt{\sum_{i=1}^{p1}(\lambda_{X}^{i})^{2}}\sqrt{\sum_{j=1}^{p2}(\lambda_{Y} ^{j})^{2}}}\end{split}  \tag{3}\]

## 4 Experimental Results and Discussion

### Experiment Setup

We compare FedCKA with the current state-of-the-art method, MOON [19], as well as FedAvg [3], FedProx [17], and SCAFFOLD [18]. We purposefully use a similar experimental setup to MOON, both because it is the most recent study and also reports the highest performance. In particular, CIFAR-10, CIFAR-100 [22], and Tiny ImageNet [23] datasets are used to test the performance of all methods. Figure 3: Training process of FedCKA

For CIFAR-10, we use a small convolutional neural network. Two \(5\times 5\) convolutional layers comprise the base encoder, with 16 and 32 channels and two \(2\times 2\) max-pooling layers following each convolutional layer. A projection head of four fully connected layers follow the encoder, with 120, 84, 84, and 256 neurons. The final layer is the output layer with the number of classes. Although FedCKA and other studies can perform without this projection head, we include it because MOON shows a high discrepancy in performance without it. For CIFAR-100 and Tiny ImageNet, we use ResNet-50 [21]. We also add the projection head before the output layer, as per MOON. We use the cross-entropy loss and SGD as our optimizer with a learning rate of 0.1, momentum of 0.9, and weight decay of 0.00001. Local epochs are set to 10. These are also the parameters used in MOON. Some small changes we made were with the batch size and communication rounds. We use a constant 128 for the batch size and trained for 100 communication rounds on CIFAR-10, 40 communication rounds on CIFAR-100, and 20 communication rounds on Tiny ImageNet. We used a lower number of communication rounds for the latter two datasets, because the ResNet-50 model overfit quite quickly. As with many previous studies, we use the Dirichlet distribution to simulate heterogeneous settings [9; 19; 24]. The \(\alpha\) parameter controls the strength of heterogeneity, with \(\alpha=0\) being the most heterogeneous, and \(\alpha=\infty\) being non-heterogeneous. We report results for \(\alpha\in[5.0,1.0,0.1]\), similarly to MOON.