# Research into Robust Federated Learning Methods Driven by Heterogeneity Awareness
## 3 Heterogeneity-Aware Federated Optimization Framework
### Overall Framework of Heterogeneity-Aware Federated Optimization
#### 3.1.2 Algorithm Implementation  
On the client side, a Global Feature Extractor is employed to perform the deep modeling of local data, generating unified embedding features with a consistent dimensionality across clients. These features provide a structural foundation for the subsequent global model aggregation. Algorithm 1 is an algorithmic implementation of the framework of heterogeneity-aware federated optimization.  
Next, the Latent Information Extractor statistically estimates the distributional characteristics of the local feature space by computing the sample mean and covariance matrix. Based on these statistics, a heterogeneity metric is introduced to measure the deviation of local data distributions from the global distribution. Specifically, the determinant of the covariance matrix is adopted as the heterogeneity indicator. This metric is incorporated into the training process as a regularization term, alongside center loss and weight regularization, to guide local model optimization.  
On the server side, after collecting model parameters and heterogeneity indicators from all the clients, a Heterogeneity-Weighted Aggregation Strategy is employed. This strategy adaptively adjusts each client's aggregation weight based on both the local sample size and the computed heterogeneity level. Clients with a high heterogeneity are assigned lower weights to reduce their negative impact on the global model. The concept of effective weight is introduced to ensure both an efficient model updating and the improved stability and generalization of the aggregated model.  
Figure 2: Overall framework of heterogeneity-aware federated optimization.  
**Algorithm 1.** Framework of Heterogeneity-Aware Federated Optimization  
**Inputs:**  
- Global model parameters: \(\theta_{0}\)  
- Client datasets: \(\{D_{1},D_{2},D_{3},\ldots D_{k}\}\)  
**Output:**  
- Optimized global model: \(\theta_{t}\)  
**Client-Side Operations** (Parallel Execution):  
1. Global Feature Extraction  
For each batch \(B\in D_{k}\):  
\(F_{k}=f_{\theta}(x_{k})\forall x_{k}\in B\)  
\(\rightarrow\) Output: Embedding matrix \(F_{k}\in R^{B\times D}\)  
2. Latent Information Extraction  
a. Compute feature statistics \(\mu_{k}\) and \(\sigma_{k}\) (Equation (1))  
b. Calculate heterogeneity metrics  
- Variance-based metric: \(H_{k}^{\text{var}}\) (Equation (2))  
- Manifold-based metric: \(H_{k}^{\text{mani}}\) (Equations (3) and (4))  
c. Unified heterogeneity score \(H_{k}\) (Equation (5))  
3. Multi-Loss Optimization  
a. Composite loss function \(L_{tatal}\) (Equation (8))  
b. Local model update:  
\(\theta_{k}^{t+1}=\theta_{k}^{t}-\eta\
abla_{\theta}L_{total}\)  
4. Transmit to Server  
Send tuple \((\theta_{k}^{t+1},H_{k},n_{k})\) where \(n_{k}=|D_{k}|\)  
**Server-Side Operations**  
1. Aggregation Initialization  
For round \(t=1\) to \(T\):  
- Receive client tuples: \((\theta_{k}^{t+1},H_{k},n_{k})\)  
2. Heterogeneity-Weighted Aggregation  
a. Compute effective weights \(\omega_{k}\) (Equation (9))  
b. Update global model:  
\(\theta^{t+1}=\frac{1}{\sum\
olimits_{k=1}^{K}\omega_{k}}\sum\limits_{k=1}^{K }\omega_{k}\theta_{k}\)  
3. Model Distribution  
Broadcast \(\theta_{t+1}\) to all clients  
**Termination**  
Return final global model \(\theta_{T}\) after T rounds