# Research into Robust Federated Learning Methods Driven by Heterogeneity Awareness
## 3 Heterogeneity-Aware Federated Optimization Framework
### Multi-Loss Design  
To enhance models' generalization and convergence stability under non-IID data environments, a multi-loss function was designed during the client training stage. This design further responds to the client distribution disparities revealed by the aforementioned heterogeneity quantification mechanism. By integrating the heterogeneity scores into the training objective, the model can perceive and adaptively adjust to various heterogeneity factors, thereby achieving coordination between heterogeneity modeling and optimiza  
Figure 3: Feature distribution and heterogeneity score of four representative clients after PCA projection. tion strategies. The overall loss function consists of four components, classification loss, heterogeneity loss, regularization term, and feature center constraint, as detailed below:  
* Classification Loss (Cross-Entropy Loss). As the fundamental component in supervised learning, the classification loss evaluates the model's prediction capability on the input data. It adopts the cross-entropy form to ensure the model learns discriminative features corresponding to local task labels. * Heterogeneity Loss. Based on the multi-scale fusion heterogeneity quantification mechanism proposed in Section 3.2, a corresponding heterogeneity loss term is introduced at the client-side training stage. This loss transforms the heterogeneity measurement into an optimizable training signal, guiding the model to suppress unstructured feature diffusions and structural drift. Specifically, this loss comprises two parts: the first is a statistical metric based on feature variance, utilizing per-dimension variance and the center deviation of intermediate embedding features to reflect the spatial dispersion of client features; the second is a nonlinear manifold-based metric, where t-SNE is applied to map features into a low-dimensional manifold space, quantifying their structural complexity and distributional divergence. * Feature Center Constraint. Considering the semantic feature drift among heterogeneous clients, a feature center constraint is introduced to enhance the consistency of feature representations. Specifically, in each training round, the mean vector of the embedded features from the current batch is computed, and the Euclidean distance between each sample's feature and this mean is measured. This term encourages a more compact distribution of features within the semantic space, reduces the disruptive effect of feature dispersion on training stability, and improves the model's discriminative capability after aggregation. It is defined as \[L_{center}=\frac{1}{B}\sum_{i=1}^{B}\lVert f_{i}-\mu_{k}\rVert_{2}^{2}\] (6)
* Regularization Term (L2 Regularization). To prevent the model from overfitting to local data distributions, an L2 regularization term is imposed on model parameters as a structural constraint. By penalizing the L2 norm of all network weight parameters, this term reduces the model's complexity and enhances it generalization capability during global aggregation. This is particularly crucial in non-IID settings, as it mitigates local performance shifts caused by overfitting. Formally, \[R_{L2}=\sum_{i}\lVert\theta_{i}\rVert^{2}\] (7)  
Finally, the total loss function optimized during client training is formulated as  
\[L_{total}=L_{\text{sup}}+\alpha\cdot H_{k}+\lambda\cdot R_{L2}+\beta\cdot L_{ center} \tag{8}\]  
where \(\alpha\), \(\lambda\), and \(\beta\) are the tuning coefficients for the heterogeneity loss, regularization term, and feature center constraint, respectively. This multi-loss design not only improves the model's representation ability on local data but also enables effective modeling and adaptation to inter-client distribution differences under the guidance of the heterogeneity quantification mechanism. By introducing heterogeneity-aware losses based on feature statistics, the model dynamically perceives and responds to non-IID characteristics, thereby guiding the local model convergence toward the global distribution during training.