## 3 FedCKA
### Modifications to FedAvg
Only truly dissimilar updates are penalized, thus improving performance. Note Equation (3) includes an inner product in the numerator. While the l2-distance or cosine similarity are also inner products, CKA is more suitable for measuring similarities between matrices of higher dimension than the number of data points, as is the case for neural network representations [10]. Furthermore, while kernel CKA slightly outperforms linear CKA, we opt to use linear CKA considering the computational overhead. We also observe that the squared Frobenius norm and vectorized cosine similarity decreased performance only slightly. These methods outperformed most previous works. This verifies that while it is important to use an accurate similarity measure, it is more important to focus on regularizing naturally similar layers. ### Efficiency and Scalability

Efficient and scalable local training procedures are important engineering principles in Federated Learning. That is, for Federated Learning to be applied to real-world applications, we must assume that clients have limited computing resources. Thus, we analyze the local training time of all methods, as shown in Table 4. Note that FedAvg is the lower bound for training time, since all other methods add a regularization term. For a seven-layer CNN trained on CIFAR-10, the training time for all methods are fairly similar. FedCKA extends training by the largest amount, as the matrix multiplication operation to calculate the CKA similarity is proportionally expensive to the forward and back propagation of the small model. However, for ResNet-50 trained on Tiny ImageNet, we see that the training time of FedProx, SCAFFOLD, and MOON increased substantially. Only FedCKA has comparable training times to FedAvg. This is because FedProx and SCAFFOLD perform expensive operations on the weights of each layer, and MOON performs forward propagation on three models until the penultimate layer. All these operations scale substantially as the number of layers increase. While FedCKA also performs forward propagation on three models, the number of layers remains static, thus being most efficient with medium-sized models. We emphasize that regularization must remain scalable for Federated Learning to be applied to state-of-the-art models. Even on ResNet-50, which is no longer considered a large model, other Federated Learning regularization methods lack scalability. This causes \begin{table}
\begin{tabular}{l c c c c} \hline \hline
**Method** & **7 Layers** & **Time Extended** & **50 Layers** & **Time Extended** \\ \hline FedAvg & 54.82 & - & 638.79 & - \\ SCAFFOLD & 57.19 & 2.37 & 967.04 & 328.25 \\ FedProx & 57.20 & 2.38 & 862.12 & 223.33 \\ MOON & 97.58 & 42.76 & 1689.28 & 1050.49 \\
**FedCKA** & 104.17 & **49.35** & 750.97 & **112.18** \\ \hline \hline \end{tabular}
\end{table} 
Table 4: Average Training Duration Per Communication Round (in seconds). \begin{table}
\begin{tabular}{c c c} \hline \hline
**Similarity Metric** & **Accuracy** & **Training Duration (s)** \\ \hline None (FedAvg) & 64.37\% & 54.82 \\ Frobenius Norm & 65.54\% & 64.73 \\ Vectorized Cosine & 66.67\% & 65.75 \\ Kernel CKA & 67.93\% & 122.41 \\ Linear CKA & 67.86\% & 104.17 \\ \hline \hline \end{tabular}
\end{table} 
Table 3: Accuracy and training duration with FedCKA with different similarity metrics (CIFAR-10). difficulty in testing these methods with the current state-of-the-art models, such as ViT [26] possessing 1.843 billion parameters, or slightly older models, such as EfficientNet-B7 [27] possessing 813 layers. ## 5 Conclusions and Future Work

Improving the performance of Federated Learning on heterogeneous data is a widely researched topic. However, many previous studies suggested that regularizing every layer of neural networks during local training is the best method to increase performance. We propose FedCKA, an alternative approach built on the most up-to-date understanding of neural networks.