1 Introduction↩︎

Class imbalance remains a major challenge in federated learning (FL) despite prior works [1]. When the global label distribution aggregated over all clients is long-tailed, minority classes are underrepresented in training, which degrades their predictive accuracy. Prior works in centralized learning mitigate imbalance through loss reweighting, advanced sampling, or generative augmentation, but these techniques are difficult to deploy in FL because privacy constraints prevent data exchange and synthetic generators are either infeasible or produce unrealistic samples for sparse regimes or domain-specific tasks such as defect detection. For instance, when working with a chip-defect dataset, one of our primary evaluation use case in this paper, synthetic defect images fail to capture the true geometric details of actual defects. To address these challenges, we focus on tackling class imbalance at the optimization level. Unlike data-level augmentations, an optimization-based approach is more generally applicable and remains effective even when data-level tricks are restricted by privacy or the characteristics of datasets.

Another line of work, Per-Class Normalization (PCN), normalizes each class-specific gradient to unit length, preventing any class from dominating an update [2]. PCN effectively decreases the loss for all classes and, because it operates directly on gradients, integrates readily with other techniques. However, PCN introduces two critical limitations in multi-class classification: it is heavily affected by directional noise (since minority class gradients often misalign with the true descent direction), and the sum of many unit vectors produces a scaling mismatch that destabilizes convergence when the number of classes increases. PCN struggles even for a moderate number of classes.

To address these issues, our primary contribution is Federated Class-Grouped Normalized Momentum (FedCGNM). Instead of normalizing all \(C\) class-specific gradients, FedCGNM partitions classes into a small number of groups (e.g., majority vs.minority) and applies unit-norm normalization to a momentum vector per group. By reducing the number of normalized vectors from \(C\) to just a few, FedCGNM mitigates the scaling-mismatch problem while still giving every group equal magnitude. Furthermore, incorporating momentum attenuates the directional noise from minority classes. An additional benefit concerns client alignment, a primary cause of performance degradation in FL [3]. Heterogeneous clients often produce gradients whose magnitudes per class differ widely, amplifying misalignment. Because FedCGNM forces each group update to have the same unit norm, the aggregated directions across clients become more aligned.

Furthermore, during local training, each client applies resampling to reduce directional noise, so selecting the appropriate sampling-rates is crucial. This choice is nontrivial in FL settings. sampling-rates must be determined jointly across clients, which induces a combinatorial search space, and sampling schedules that fluctuate excessively can destabilize optimization. Our theoretical analysis indicates that stable convergence requires controlling the cumulative variation of sampling-rates across rounds. For these reasons, practical approaches adopt a fixed resampling policy, using a common rate across clients. To find sampling-rates efficiently, we introduce FedHOO as an auxiliary, plug-in algorithm designed specifically for small-scale federations. Based on the X-armed bandit framework, FedHOO exploits the inherent parallelism in FL. In every communication round, FedHOO requests each client to train with only two candidate rates, yet by linearly combining the returned updates the server can infer the validation performance metric for all \(2^K\) rate combinations, where \(K\) is the number of clients. The method therefore identifies effective sampling-rates early in training, while avoiding an exhaustive sweep of the hyperparameter space.

Across four public benchmarks, our methods consistently outperforms traditional reweighting and sampling baselines, with gains up to 29% over FedAvg combined with resampling. In a large-scale industrial chip defect dataset, our method achieves a 16% improvement over the best baseline. Our main contributions are as follows.

1. FedCGNM optimizer and variance-aware grouping rule. We introduce FedCGNM, the first client‐side optimizer that groups classes, applies unit-norm momentum per group to balance majority and minority influence while reducing noise and scale mismatch. We frame class partitioning as minimization of within‐group variance. This rule produces the optimal split found by exhaustive search, yielding a principled yet lightweight grouping strategy.

2. Convergence analysis. We incorporate the sampling-rate schedule into the convergence analysis, representing a novel effort combining three components: FL, group normalization, and dynamic sampling-rates. We demonstrate that the sampling-rate schedule impacts convergence.

3. FedHOO as an auxiliary sampling-rate tuner. While a grid search for sampling-rates is common, more refined strategies are advantageous when the number of clients is low. The proposed X-armed-bandit-based exploration scheme is the first optimization based algorithm to determine combinatorial local sampling-rates in FL. The method performs rapid, privacy-preserving search trading off exploration and exploitation.

4. Strong performance. We validate our framework across multiple public benchmarks and a real-world industrial semiconductor chip-defect dataset. Across all settings, it consistently shows improvement over strong baselines.

2 Related Works↩︎

Class imbalance poses a significant challenge in supervised learning, where limited data from minority classes leads to biased models and poor performance on those classes [4], [5]. Common training-level solutions include reweighting, which adjusts learning based on class frequency, optimization-level methods, and resampling, which alters the class distribution in training data.

2.0.0.1 Reweighting and Optimization Methods

Optimization-perspective methods generally adjust the learning process by modifying the loss function or the gradient update rule. Most reweighting methods adjust each sample’s contribution within the loss function to counteract class imbalance, such as weighted cross-entropy [6], focal loss [7], and class-balanced loss [8]. Beyond loss‐level re‐weighting, PCN [2] rescales each class‐specific gradient to unit norm, equalizing per‐class influence during optimization. [9] introduces a technique to adjust the weight of gradient dynamically in a class-incremental learning scenario. They consider reweighting in class-level which works poorly as the number of classes increases, and only consider balance of the gradient magnitude. Different from them, we tackle class imbalance by addressing gradient scaling and directional noise simultaneously.

Several methods have tailored re‐weighting to FL setting. CLIMB [10] addresses the class imbalance in FL by enforcing the similarity between local empirical losses, while Fed-GraB [11] introduces a self-adjusting gradient balancer that dynamically reweights gradients based on the class distribution inferred from classifier weights. FL Ratio Loss [12] builds on centralized Ratio Loss by estimating global class proportions via secure aggregation and adjusting local losses accordingly. FedNoRo [13] leverages knowledge distillation and distance‐aware aggregation to align client models, and incorporates a logical adjustment mechanism to address both data heterogeneity and class imbalance. Unlike these methods that focus on loss functions or models, we tackle class imbalance at the gradient level, pairing with a simple resampling technique.

2.0.0.2 Resampling and Data Synthesis

Traditional resampling methods balance class balance by removing majority samples (under-sampling) or replicating minority ones (over-sampling) [14]. More recent techniques like SMOTE [15], GAMO [16], and I-GAN [17] generate synthetic minority data and show strong empirical performance. However, when minority samples are extremely scarce [18] or synthetic data risks being unrealistic or mislabeled [19], such methods become less viable. Thus, this work confines itself to conventional under- and over-sampling without synthetic generation since we focus on those situations.

Choosing how much to re-sample remains an open problem: systematic investigations in centralized deep learning reveal that the optimal under- or over-sampling-rate depends jointly on the dataset size and the severity of class skew [20]. Curriculum-based schemes, such as Dynamic Curriculum Learning [21], further highlight the need to adapt sampling ratios over the course of training rather than fixing them a priori. In the federated setting, [22] cast client-side resampling as a tunable policy, optimized to minimize the global loss while respecting privacy constraints. FAST [23] advances this idea by viewing each sampling ratio as an arm in a multi-armed-bandit framework, enabling dynamic exploration during training. However, it treats local sampling-rates independently, overlooking the combinatorial nature of FL. Our method adopts this adaptive philosophy, particularly suited to small-scale federations, while addressing the combinatorial optimization challenge.

3 Methodology↩︎

Consider a federated learning system with \(K\) clients for a \(C\)-class classification task. To cope with class imbalance, we incorporate a resampling strategy, which is a common strategy for tackling class imbalance. Let \(\mathcal{D}_k\) denote the original data distribution of client \(k\). We define \(\mathcal{D}_k(r_k)\) as the data distribution of client \(k\) after resampling with a sampling-rate \(r_k \geq 0\). If \(q_k^{(c)} \ne 0\) is the fraction of samples in class \(c\) for client \(k\), then under sampling-rate \(r_k\) client \(k\) resamples class \(c\) by a factor \(\left( \frac{\max_s{q_k^{(s)}}}{q_k^{(c)}} \right)^{r_k}\). Consequently, the fraction of class \(c\) in a mini-batch changes according to this resampled distribution.

We define the local objective function under resampling as \(f_k(x;r_k) = \mathbb{E}_{\xi \sim \mathcal{D}_k(r_k)} [\ell(x;\xi)]\), where \(\ell(x;\xi)\) is the sample loss for \(\xi\). Given sampling-rates \(r = (r_1,\dots, r_K)\), we define global objective as \(f(x;r)=\sum_{k=1}^K p_k f_k(x;r_k)\) where \(p_k\) is the weight of client \(k\). Furthermore, we assume a resampling-rate strategy determines the rates \(r^{(t)} \in \mathbb{R}^K\) before round \(t\) starts. If \(\mathcal{F}_t\) is the filtration containing all randomness up to round \(t-1\), we assume that \(r^{(t)}\) is \(\mathcal{F}_t\)-measurable.

3.1 Federated Class-Grouped-Normalized-Momentum↩︎

Per-Class Normalization (PCN) addresses gradient imbalance by normalizing each class-specific gradient and has shown promising performance in imbalanced binary classification. However, it faces a limitation in multi-class settings (see Appendix 9.1 for more details). The next-to-be-proposed FedCGNM is designed for multi-class problems and aims to overcome the limitations of PCN.

FedCGNM merges classes into a small number \(H \ll C\) of groups (typically majority and minority) and maintains a momentum per group. At communication round \(t\), the server broadcasts the global model \(x^{(t)}\) and local sampling-rates \(r_k^{(t)}\) based on a sampling strategy (see Section 3.4), and each client constructs a partition \(\{\mathcal{G}_{k,h}^{(t)}\}_{h=1}^{H}\) of the classes (see Section 3.2 discerning the grouping rule). If we denote \(f_k^{(t)}(x):=f_k(x;r_k^{(t)})\) as the local objective of the client \(k\) at round \(t\), then it decomposes as \(f_k^{(t)}=\sum_{h=1}^{H}f_{k,h}^{(t)}\), with \(f_{k,h}^{(t)}\) being the sum of the loss functions under the resampled distribution over the samples in \(\mathcal{G}_{k,h}^{(t)}\).

During local training, client \(k\) updates, for each group \(h\) and step \(i\), the momentum \[m_{k,h}^{(t,i)}=\beta m_{k,h}^{(t,i-1)}+(1-\beta)\,g_{k,h}^{(t,i)}, \quad m_{k,h}^{(t,0)}=0,\] where \(g_{k,h}^{(t,i)}(x) =\nabla f_{k,h}(x;r_k^{(t)};\xi_{k,h}^{(t,i)})\) is the stochastic gradient computed on samples in mini-batch \(\xi_k^{(t,i)}\) drawn solely from group \(\mathcal{G}_{k,h}^{(t)}\), and \(\beta\in[0,1)\) is the momentum factor. The per-step update direction is obtained by normalizing each momentum with a stabilizing constant \(\delta > 0\) and summing across the \(H\) groups as \[x_{k}^{(t,i)} = x_{k}^{(t,i-1)} - \eta\sum_{h=1}^{H} \frac{m_{k,h}^{(t,i)}}{\|m_{k,h}^{(t,i)}\| + \delta}\] with learning rate \(\eta>0\). The term \(\delta\) ensures numerical stability and provides a lower bound on the denominator, which is crucial for theoretical analysis and to avoid numerical instability during training, especially when the gradient norm is small [24]. After \(E\) iterations the client returns \(x_{k}^{(t,E)}\) to the server, which aggregates by \(x^{(t+1)}=\sum_{k}p_k x_{k}^{(t,E)}\). Operating on a handful of group momenta stabilizes the step norm, suppresses directional noise, and remains computationally efficient even when the number of classes is large.

3.2 Grouping of Classes↩︎

We formulate the problem of partitioning \(\{1,\ldots,C\}\) into \(H\) disjoint groups as one-dimensional variance reduction on the class proportions. For the sake of notation, we omit the dependency on client \(k\) and iteration \(t\) in the subsequent quantities. Let \(q_c\) be the (resampled) proportion of class \(c\), with \(\sum_{c=1}^C q_c=1\), and we assume \(q_1 \ge \cdots \ge q_C\). We treat \(\{q_c\}_{c=1}^C\) as points on the real line and select \(H-1\) thresholds to form contiguous groups in which proportions are as similar as possible.

For a partition \(G=\{\mathcal{G}_h\}_{h=1}^H\), define the group mass \(S_h=\sum_{c\in\mathcal{G}_h} q_c\), group mean mass \(\mu_h=S_h/|\mathcal{G}_h|\), and the within-group distribution \(w_{c\mid h}=q_c/S_h\). In a mini-batch of size \(B\), let \(N_h\sim \mathrm{Binomial}(B,S_h)\) be the sample counts of group \(h\) and \(N_{h,c}\) be the sample counts of class \(c\) in group \(h\) based on the distribution. Define the empirical share vector \(\hat{\boldsymbol{w}}_h=(\hat{w}_{c\mid h})_{c\in \mathcal{G}_h}\) with \(\hat{w}_{c\mid h}=N_{h,c}/N_h\), and compare it with the uniform target \(u_h=(1/|\mathcal{G}_h|,\ldots,1/|\mathcal{G}_h|) \in \mathbb{R}^{|\mathcal{G}_h|}\). We define group imbalance vector \(\Delta_h=\tfrac{N_h}{B}(\hat{\boldsymbol{w}}_h-u_h)\). Taking expectation, we obtain \[\mathbb{E}\|\Delta_h\|^2 =\Big(S_h^2+\tfrac{S_h(1-S_h)}{B}\Big)\|\hat{\boldsymbol{w}}_h-u_h\|^2 +\tfrac{S_h}{B}\Big(1-\sum_{c\in\mathcal{G}_h} w_{c\mid h}^2\Big).\] The optimization problem for grouping is \(\min_{\mathcal{G}_h} \mathbb{E} \| \Delta_h \|^2\). We solve this problem heuristically by finding the best threshold that yields groups whose class distributions have minimal within-group variance.

The dominant term of \(\mathbb{E} \| \Delta_h \|^2\) is \(S_h^2\|\boldsymbol{w}_h-u_h\|^2\) with \(O(1/B)\) corrections. We next link this imbalance to the variance of raw proportions. Using \(\|\boldsymbol{w}_h-u_h\|^2 =\sum_{c\in\mathcal{G}_h}\Big(\frac{q_c}{S_h}-\frac{1}{|\mathcal{G}_h|}\Big)^2=\frac{|\mathcal{G}_h|}{S_h^2}\,\sigma_h^2\) where \(\sigma_h^2=\frac{1}{|\mathcal{G}_h|}\sum_{c\in\mathcal{G}_h}(q_c-\mu_h)^2\), the dominant term becomes \(S_h^2\|\boldsymbol{w}_h-u_h\|^2=|\mathcal{G}_h|\,\sigma_h^2\). Summing over groups and normalizing per class yields \(\tfrac{1}{C}\sum_{h=1}^H |\mathcal{G}_h|\sigma_h^2=\sum_{h=1}^H \omega_h\sigma_h^2\) with \(\omega_h=|\mathcal{G}_h|/C\), which is the within-group variance objective on the sorted \(\{q_c\}_{c=1}^C\). Consequently, selecting \(H-1\) thresholds by this strategy produces the partition that minimizes the class-balanced within-group dispersion and, by the argument above, asymptotically minimizes expected per-batch imbalance.

Figure 2 reports the test accuracy obtained when we exhaustively vary the split threshold \(\tau_g\), which assigns the \(\tau_g\) rarest classes to the minority group and the remaining \(C-\tau_g\) classes to the majority group. Accuracy rises to a clear maximum near \(\tau_g=69\) and declines when either too few or too many classes are treated as minority. The red line marks the threshold chosen by our variance-based grouping rule, which, in this example, coincides with the empirical optimum. Additional experiments in Appendix 9.2 exhibit the same pattern, confirming that minimizing within-group variance serves as a reliable proxy for an exhaustive threshold search.

In summary, variance-aware grouping via minimizing the within-group variance on empirical class distribution provides a data-driven, lightweight mechanism that, when coupled with normalized momentum, substantially attenuates gradient noise while ensuring balanced directional contributions from majority and minority classes.

3.3 Theoretical Convergence Analysis↩︎

We analyze the convergence of FedCGNM by explicitly incorporating a resampling-rate schedule \(\{r^{(t)}\}\). By modeling the global objective \(f^{(t)}(x) = f(x; r^{(t)})\) as a function of these time-varying rates, we can rigorously quantify how the choice of a sampling strategy impacts training dynamics. For the theoretical convergence analysis, we make the following assumptions.

Assumption 5 asserts that, at each client–class update, the stochastic gradient cannot be both perfectly aligned and equal in length to its exponentially averaged momentum. One of the two gaps is virtually guaranteed in practice, and the assumption is strictly weaker than the simultaneous angle-and-norm bounds adopted in earlier analyses of normalized momentum methods. Note that we include the factor \(\alpha_i:=1-\beta^i\) in the norm-gap branch of Assumption 5 because \(\|m_{k,h}^{(t,i)}\|\) is naturally on the scale of \(\|\alpha_i \nabla f_{k,h}(x^{(t)})\|\) rather than \(\|\nabla f_{k,h}(x^{(t)})\|\). The following theorem shows that FedCGNM attains the standard \(\mathcal{O}(T^{-1/2})\) stationary-point rate under smoothness and bounded-variance assumptions. A complete proof is given in Appendix 6.

The second term on the right-hand side of ?? represents the path variation of the sampling-rates. For the algorithm to converge to a stationary point at the standard rate of \(\mathcal{O}(T^{-1/2})\), the sampling strategy must ensure that \(V_T /\sqrt{T} \to 0\).

3.4 Resampling Strategy↩︎

The path variation term in Theorem 1 implies a fundamental trade-off where aggressive resampling can improve minority class learning, but an indefinite oscillation prevents the global objective from stabilizing. Therefore, a practical resampling strategy must eventually stabilize.

While dynamic sampling-rate search has been explored in centralized curriculum learning [21] and in per-client sampling-rate search for FL [23], it presents unique challenges in FL due to client heterogeneity and large search-space. If a search algorithm fails to identify sampling rates quickly enough to reduce oscillation during training, convergence may deteriorate. We note, however, that training stability does not strictly guarantee superior test performance. While a full exploration of the optimal balance between these two competing objectives is beyond the scope of this work, our convergence analysis establishes a critical requirement of \(V_T\) for practitioners.

For small federations where \(K\) is small, we introduce Federated Hierarchical-Optimistic-Optimization (FedHOO) as a strategy designed to search effective rates rapidly. This approach leverages the observation that aggregation in FL is linear with respect to client updates, so the server can “mix and match” client updates computed under different local hyperparameter choices. We cast the joint sampling-rate search as an X-armed bandit over a continuous space, and the server follows a HOO-style optimistic rule to balance exploration and exploitation when choosing which rate intervals to probe each round. Here, we outline only the key idea of FedHOO and the full algorithmic details are provided in Appendix 7. The following methodology is applicable to general hyperparameter tuning in FL, but we present it from the perspective of sampling-rates.

To search for joint sampling-rates in FL, the server maintains a plausible interval of sampling-rates for each client \(k\). In standard FL, the server can evaluate only one joint hyperparameter configuration per round. In contrast, FedHOO enables a much broader exploration by instructing each client to train at two probe rates (e.g., one lower and one upper probe) inside the current interval. Each client performs two local training sessions starting from the same global model using the assigned probe rates and returns the two corresponding client updates, i.e., the model changes produced by local training. Since the server aggregates these updates via a weighted sum, it can construct and validate all \(2^K\) models corresponding to every combination of the local rates, select the best performing model for the next starting global model, and shrink each client’s interval toward the chosen probe. Iterating this procedure yields rapid rate identification and stabilization.

3.4.0.1 Federated validation.

FedHOO does not require centralized validation data. After synthesizing the candidate global models, the server asks clients to evaluate them on their private validation splits. Clients return only scalar validation metrics, such as validation loss or accuracy, and the server aggregates these metrics to select the best candidate model. Thus, FedHOO adds validation overhead but does not require transferring raw validation data to the server. If a server-side hold-out set is available, it can be used as an alternative, but it is not required.

Given that FedHOO requires validating \(2^K\) candidates per round, it is designed for small-scale regimes (e.g., \(K \leq 5 \;\text{or}\; 6\)). For larger federations where the exponential validation cost becomes prohibitive, we revert to a simpler strategy such as a uniform global rate. Further, in practice, we run FedHOO only in the early rounds to quickly find stable client-specific rates, then fix them for the remaining rounds. Empirically, Appendix 10.2 shows that FedHOO reaches stable resampling-rates within a small number of rounds, whereas a standard exploration strategy that evaluates only one configuration per round fails to stabilize within the same training rounds.

4 Experiments↩︎

4.1 Experiment Settings↩︎

Our experiments evaluate FedCGNM and FedHOO on five classification benchmarks: two long-tailed image collections (CIFAR-10-LT and CIFAR-100-LT [25]), two tabular datasets (Adult Income [26] and UNSW-NB15 [27]), and a semiconductor chip-defect dataset. We control the degree of class imbalance via the imbalance rate \(\xi\) where larger \(\xi\) corresponds to more extreme imbalance. Details of dataset and imbalance settings are provided in Appendix 8.

The semiconductor Chip-Defect-Detection (CDD) corpus consists of approximately 780,000 high-resolution (224 × 224 after pre-processing) images captured on select years after 2020 from five factories. The ratio of defect images is 1.7%, with seven defect categories observed across hundreds of product types. Additional heterogeneity distributions are presented in Figure 4 and Figure 7. We split seventy percent of the dataset to training, fifteen percent to validation, and the remaining fifteen percent to test. We report the company’s weighted accuracy that balances defect recall and non-defect precision. Some exact counts are withheld in accordance with the non-disclosure agreement. For full details of the implementation, including the selection of hyperparameters, see Appendix 8.

We compare FedCGNM against seven alternatives: (1) FedAvg [28] with standard SGD, (2) FedAvg using class-weighted cross-entropy [6], (3) FedAvg with Ratio Loss [12], (4) FedProx [29], (5) CLIMB [10], (6) FedGraB [11], and (7) FedCGN, which adapts Per-Class Normalization to a grouped setting and employs CGN for local optimization. We compare optimization-based imbalance methods because our approach also targets the optimization dynamics, and both our method and these baselines can be readily combined with complementary remedies such as resampling or knowledge distillation. All methods share identical backbones and hyper-parameter schedules, but see Appendix for details of models, initial learning rate, batch size, weight decay and learning-rate scheduler. We use \(H=2\) for all grouping implementations.

3pt

3pt

Two federation regimes are considered. In the small-scale scenario includes five clients with full participation. In the moderate-scale scenario, twenty clients are available but only half are sampled each round. Public datasets are partitioned IID/Non-IID across clients, where we generate disjoint non-IID client data using a Dirichlet distribution, Dir(\(\psi\)), with a concentration parameter \(\psi\) set to 0.5, as described in [30]. The chip corpus is divided according to natural factory boundaries, resulting in highly non-IID client data.

4.2 Primary Real-World Evaluation: Chip-Defect Detection (CDD)↩︎

The semiconductor CDD task provides the realistic evaluation in our study since this dataset reflects actual production condition, with only 1.7% of samples exhibiting defects. Figure 4 shows that the standard FedAvg baseline stalls at around 73, while other FL baselines provide only marginal improvements. FedCGN pushes performance into 85, and FedCGNM adds another improvement. Using FedHOO as resampling strategy consistently improves training. This demonstrates that variance-aware grouping, per-group momentum, and efficient rate exploration are not only effective on public benchmarks but also critical for real industrial deployments.

4.3 Evaluation on Public Datasets↩︎

Table 1 summarizes F1 under two federation regimes. Across both federation regimes, FedCGNM outperforms all baselines by several points on average in all setting. In the small-client setting (\(K=5\)), coupling FedCGNM with FedHOO yields the best accuracy in all scenarios. Appendix 10.2 further isolates the effect of FedHOO by comparing it against other local exploration strategies . FedHOO consistently improves multiple FL baselines, supporting that its benefit comes from joint rate selection. Even without FedHOO, FedCGNM alone consistently beats competing optimizers. The only exception occurs under mild skew on Adult Income, where Ratio Loss briefly matches FedCGNM, suggesting that simple loss reweighting can suffice when imbalance is mild.

As the number of clients increases from five to twenty, the accuracy decreases for all methods due to the greater heterogeneity. However, FedCGNM exhibits noticeably milder degradation compared to FedCGN, which is the best baseline: on average, FedCGN loses about 14–37% performance, while FedCGNM drops by only 7–32%. This indicates that group momentum not only boosts performance in small federations but also makes optimization more robust to client scaling.

4.4 Sensitivity Analysis↩︎

4.4.0.1 Sensitivity to Imbalance Severity

On the multi-class image benchmarks (CIFAR-10 and CIFAR-100), all methods show declining F1 scores as the class distribution becomes more imbalanced, but FedCGNM’s drop is noticeably milder. By contrast, on the Adult Income task, FedCGNM offers only a slight improvement, reflecting the relative ease of a binary prediction problem.

4.4.0.2 Effect of Grouping and Number of Groups

We examine the impact of the number of groups using training and validation losses for \(H=2,3,4\) in Figure [fig:loss95curve95by95H] and Table 5. For additional results, see Appendix 9.3). With two groups, validation loss remains lowest and most stable, whereas three or four groups lead to earlier increases in validation loss despite continued training loss reduction, indicating accelerated overfitting. On CIFAR-100, \(H=3\) occasionally yields slightly better accuracy, but the gap is marginal, so two groups offer the best trade-off between balanced class influence and generalization. Beyond the number of groups, our grouping strategy based on minimizing within-group variance provides clear benefits compared to naive half-splits: on CIFAR-10-LT20 with \(K=5\), our rule achieves an F1 score of 0.8565 versus 0.8382, and on CIFAR-100-LT20 with \(K=5\), it improves from 0.4339 to 0.4983. These gains confirm that variance-aware grouping better balances class contributions and yields superior generalization across datasets.

4.4.0.3 Sensitivity to Group Momentum Factor

We perform an ablation study on the group momentum factor \(\beta\) across multiple datasets and client settings. The results in Table ¿tbl:tab:comparison95by95beta? show that very small (\(\beta=0\), which is FedCGN) or very large (\(\beta=0.9\)) values generally underperform, while moderate values yield the strongest results. For example, on CIFAR-10-LT20 with \(K=5\), accuracy improves from 83.35 at \(\beta=0\) to a peak of 85.68 at \(\beta=0.5\), and on CIFAR-100-LT20 with \(K=20\), performance rises from 38.27 to 40.51 at the same setting. Similar trends appear consistently across datasets, indicating that moderate momentum factors strike the right balance between stability and adaptivity, thereby offering the best overall performance.

Additional experiments (Appendix 10) confirm that FedHOO accelerates sampling-rate search compared to standard HOO, and FedCGNM maintains its advantage under large-client federations. These results demonstrate the robustness and generality of our approach across diverse and challenging federated learning conditions.

5 Conclusion↩︎

We introduced FedCGNM, a client‐side optimizer that balances majority and minority class influence by grouping labels and applying group-wise normalized momentum, and FedHOO, a federated exploration strategy that exploits FL parallelism to efficiently tune client sampling rates without exhaustive search (in small-\(K\) regimes). We develop a convergence analysis that accounts for dynamic sampling-rate schedules. Experiments on multiple public benchmarks and a large-scale industrial chip-defect dataset demonstrate consistent improvements over competing baselines. Together, these components form a practical framework for mitigating global class imbalance in privacy‐preserving federated settings.

Acknowledgements↩︎

The authors express sincere gratitude to Intel Corporation for their generous financial support, which was instrumental in completion of this research. Their commitment to advancing innovation has significantly contributed to the success of this project.

Impact Statement↩︎

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

6 Proof↩︎

6.1 Supplementary Lemmas↩︎

6.2 Proof of Theorem 1↩︎

7 Details of FedHOO↩︎

This section formalizes FedHOO as a continuous hyperparameter search problem over client-wise resampling-rates. We first cast rate selection as an X-armed bandit on the domain \(\mathcal{X}=[r_{\min},r_{\max}]^{K}\), and briefly review how classical hierarchical optimistic optimization (HOO) navigates this space via a tree-based partition. We then introduce FedHOO, which retains HOO’s hierarchical structure but exploits the linear aggregation in federated learning to evaluate \(2^{K}\) joint rate configurations per round from only two local runs per client, enabling substantially faster exploration and stabilization in small-client regimes.

7.1 FedHOO Algorithm↩︎

A possible way to find sampling-rates is to view the problem through the lens of the X-armed bandit (XAB) over the continuous domain \(\mathcal{X} = [ r_{\min}, r_{\max}]^K\). An arm \(r = (r_1,\ldots,r_K) \in \mathcal{X}\) specifies client rates, and pulling \(r\) consists of executing one FL round where client \(k\) trains with \(r_k\), aggregating the local models, and evaluating the aggregated model. The reward \(f(r)\) is the performance of the resulting aggregated model, and the learner must balance exploration and exploitation to identify a near-optimal \(r\in\mathcal{X}\). Hierarchical Optimistic Optimization (HOO, [31]) addresses XAB by creating a hierarchical, tree-based partition of the search space \(\mathcal{X}\). The algorithm iteratively navigates this tree, selecting the most promising subregions to explore based on an optimistic estimate of their potential reward.

However, standard XAB methods (e.g., exhaustive search or HOO) converge too slowly under FL’s limited rounds and early sensitivity to hyperparameters. We introduce FedHOO, which exploits the parallelism of FL and is suitable for a small number of clients. FedHOO retains HOO’s tree of boxes and enables much faster exploration by letting clients perform local training with two rates and exploring all combinations of local rates for validations.

The space \(\mathcal{X}=[r_{\min},r_{\max}]^{K}\) is organized as a \(2^K\)-ary tree. The root covers the entire search space, and each child \(\nu\) halves its parent’s range, and the midpoint \(c(\nu)\) of node \(\nu\) ’s interval acts as its representative point. Each node \(\nu\) in depth \(h\) stores a box \(I(\nu) = \{[L_k(\nu), U_k(\nu)]\}_{k=1}^K\), an evaluation count \(N(\nu)\), a running reward estimate \(V(\nu)\), and an optimistic score \(B(\nu)\).

At round \(t\), the server selects the leaf \(\nu\) with largest \(B(\nu)\). Let \(\mathcal{I}(\nu)=\prod_{k}[L_{k},U_{k}]\) be its box and \(h\) be its depth in the tree. For each client \(k\), the server sends two rates \(r_{k}^{L}=(3L_{k}+U_k)/4\), \(r_{k}^{U}=(L_k+3U_{k})/4\), and the current model \(x^{(t)}\). Client \(k\) trains two local models (one per rate) and returns the deltas. By linearly combining these deltas, the server aggregates the models corresponding to all \(2^K\) lower/upper choices, thereby evaluating the rewards of the \(2^K\) child nodes in one round.

The optimistic score of \(\nu\) is \[B(\nu)\;=\;V(\nu)\;+ \; \tau\; \text{diam}(\nu)^h \; +\;\sqrt{\tfrac{\alpha\ln (t+1)}{N(\nu)}}\;,\] where \(\mathrm{diam}(\nu)\) is the diameter of \(\mathcal{I}(\nu)\) and \(\tau, \alpha>0\) are hyperparameters. The second term of \(B(\nu)\) is borrowed from UCB bandits [32] exploration. In FedHOO, server also keep track of reward \(V(\nu)\). Finally, the server expands \(\nu\) into its \(2^K\) child nodes corresponding to all combinations of two rates per client, and sets \(x^{(t+1)}\) to the one with the highest reward among the \(2^{K}\) candidates just evaluated.

Doubling each client’s local training time unlocks an exponential exploratory gain. Enumerating those combinations explicitly would be prohibitive, but FedHOO obtains the same information at linear cost by utilizing parallelism of FL, making it vastly more efficient than existing search methods.

The strategy is especially advantageous in small federations, a configuration frequently encountered in industrial deployments. We therefore apply FedHOO when the number of clients is low. When \(K\) is large, we revert to a uniform global sampling-rate because significant exploration cannot be completed within a reasonable training budget and a single rate promotes update alignment, a consideration that becomes increasingly critical as \(K\) grows. We present the pseudocode of FedHOO in Algorithm 6. We write \(\mathrm{diam}(\nu) = \max_k \big(U_k(\nu)-L_k(\nu)\big)\) corresponding to the \(\ell_\infty\) width of the interval and use \(\odot\) as element-wise product.

7.2 Intuition with an Example↩︎

To illustrate how FedHOO works, consider a simple setting with \(K=2\) clients, and the search space is \([0,1]^2\). The root node of the search tree corresponds to the full box \([0,1]\times[0,1]\). If we split this root into four quadrants, the child nodes are \[[0,0.5]\times[0,0.5], \quad [0,0.5]\times[0.5,1.0], \quad [0.5,1.0]\times[0,0.5], \quad [0.5,1.0]\times[0.5,1.0].\] Exploring further, the child \([0,0.5]\times[0,0.5]\) can itself be split into four sub-boxes such as \([0,0.25]\times[0,0.25]\) and \([0,0.25]\times[0.25,0.5]\), and so on. This recursive partitioning continues as the algorithm zooms in on promising regions.

In the standard HOO algorithm, only one node can be explored in each round. For example, at round 1, HOO would select just a single quadrant, say \([0,0.5]\times[0.5,1.0]\), and update its statistics based on a single evaluation. Over many rounds, this gradually builds information, but the search may proceed slowly because each evaluation provides feedback for only one region of the tree.

By contrast, FedHOO leverages the federated setting to explore many regions at once using parallelism of FL. Because each client performs two local runs (with two different probe rates), the server can synthesize outcomes corresponding to all corner combinations of the current interval. In the two-client example, this means that in a single round, FedHOO obtains rewards for all four quadrants simultaneously. For instance, at round 1, FedHOO evaluates the four children at depth \(h=1\): \[[0,0.5]\times[0,0.5], \quad [0,0.5]\times[0.5,1.0], \quad [0.5,1.0]\times[0,0.5], \quad [0.5,1.0]\times[0.5,1.0].\] This parallel evaluation dramatically accelerates the search. Rather than spending four rounds to cover these quadrants (as in HOO), FedHOO requires only one. As the depth increases, the same principle applies: four sub-intervals at depth \(h=2\) can be evaluated together by reusing the two local runs per client.

7.3 Summary and Limitations↩︎

The main advantage of FedHOO is that the federated setting allows the server to combine a small number of local runs into exponentially many synthetic evaluations. In the two-client example, two runs per client yield four corner evaluations in each round, and more generally \(2^K\) corners can be evaluated from only two runs per client. This exponential coverage greatly accelerates the search compared to classical HOO, which can only evaluate a single node per round.

FedHOO is not intended as a universal tuner for very large federations. Its main purpose is to exploit the small-client regime, which appears in industrial deployments such as our CDD benchmark. In this regime, the exponential candidate set is still manageable, and the ability to evaluate joint rate combinations is more valuable than independent local tuning. For large-scale FL, we recommend either using a uniform global rate or applying FedHOO at a coarser level, such as over client clusters rather than individual clients.

Another limitation of this approach is the overhead of validating \(2^K\) candidate models at each round, which may become expensive for large \(K\). In addition, extending the procedure to partial participation is not straightforward, since missing client updates can prevent consistent synthesis of all corners. These issues suggest that while FedHOO is powerful for moderate \(K\) or cluster-level dimensions, further work is needed to make it scalable to very large federations.

8 Detail of the Experiment Setting↩︎

In this appendix, we provide full details of our experimental setup, including datasets, model architectures, training hyperparameters, and federated protocols.

8.1 Datasets and Models↩︎

We evaluate on five benchmarks. CIFAR-10-LT and CIFAR-100-LT are long-tailed variants of CIFAR-10/100 with imbalance rates 20 and 100, constructed via the protocol of [33]. Concretely, for imbalance rate \(\xi\), the number of training samples in class \(c\in\{0,\dots,C-1\}\) (sorted by frequency) is \[N_c \;=\; N_{\max}\,\xi^{-c/(C-1)},\] where \(N_{\max}\) is the majority-class size and \(C\) is the number of classes. If a dataset has balanced classes, we randomly sample from each class \(c\) to have \(N_c\) samples to make an imbalanced dataset. Test sets remain unaltered to test the performance of algorithms to produce balanced performance.

Adult Income [26] is a binary classification task with original imbalance ratio 3.17 and additional settings of 10 and 20 obtained by subsampling the minority class. UNSW-NB15 [27] provides 175,341 network-flow records spanning ten classes with a natural imbalance of roughly 434. For Adult Income and UNSW datasets, we preserve the same class imbalance in the test splits as in their corresponding training sets to mirror real-world conditions. The proprietary Chip-Defect-Detection (CDD) corpus comprises of approximately 780k optical micrographs (224×224) from five semiconductor fabs, in which only 1.7% of samples contain defects. CDD dataset is split by factory into 70% train, 15% validation, and 15% test. To expose the class imbalance inherent in our dataset, Figure 7 reports the number of defect and non‑defect samples recorded by each factory.

We train a ResNet18 [34] model with group normalization from scratch on all public image benchmarks. For tabular tasks we employ a four layer fully connected network with hidden dimensions of 32, 16, 8, and 2 and ReLU activations on the Adult Income data and we adopt the CNN–LSTM architecture of [35], consisting of stacked one-dimensional convolutional filters feeding into an LSTM to capture temporal correlations in network-flow features for UNSW-NB15. For CDD, we use ResNet-34. Table 3 summarizes each architecture.

8.2 Implementation and Training Details↩︎

All algorithms—FedAvg (SGD), FedProx, weighted cross-entropy, Ratio Loss, FedGraB, CLIMB, FedCGN, and FedCGNM—share the same training search strategy. We tune the initial learning rate over the set \(\{0.4,0.2,0.1,0.05,0.01\}\) by grid search and then decay it with a cosine annealing curve that reaches \(10^{-4}\) at the final round. Weight decay is selected from \(\{10^{-4},10^{-3}\}\) and the batch size is selected from \(\{32, 64, 128\}\). For FedCGNM, we set stabilizing constant \(\delta=0.1\) across all datasets since the algorithm is robust under the choice of \(\delta\) values. We also test the group momentum coefficient \(\beta\) over \(\{0.5,0.6,0.7,0.8,0.9\}\). Each communication round runs three local epochs when five clients participate and five local epochs when twenty clients participate on the image benchmarks; Adult Income and UNSW NB15 use a single local epoch per round because their training sets are comparatively small. All jobs are executed with PyTorch \(2.4.1\) and CUDA \(11.4\) on four NVIDIA TITAN Xp GPUs.

sampling-rates \(r_k\in[0.4,0.8]\) are tuned via FedHOO when \(K=5\), initializing a hierarchical tree over \([0.4,0.8]^K\) with optimism constant \(\alpha=1.0\). In practice, since validation accuracy increases over rounds, we maintain each node’s reward as an exponential moving average of validation accuracies for fair comparison. In each round, clients evaluate two candidate rates \((r_k^L,r_k^U)\), return update deltas, and the server extrapolates rewards for all \(2^K\) combinations at linear cost. We run FedHOO only in the initial training phase and terminate it once the selected sampling-rates remain within the same per-client search intervals for five consecutive rounds. We then fix the resulting rates and use them for the remaining rounds to avoid long-term oscillation and validation overhead. For \(K=20\), a uniform global rate is used. In implementation, we perform three rounds of training without expanding tree to tackle cold start problem, since the initial training rounds are sensitive to sampling-rate choices, and cap the FedHOO search tree at depth 5. This provides a fine enough resolution of sampling-rates while keeping the exponential branching factor computationally manageable. Empirically, we found that deeper trees do not yield meaningful accuracy gains, as the noise inherent in federated training outweighs the benefit of additional granularity, whereas depth 5 provides a stable balance between exploration and efficiency.

In the small-scale regime (\(K=5\)), all clients participate each round. In the moderate-scale regime (\(K=20\)), we sample 50 % of clients per round. Public datasets are split IID/Non-IID across clients; the CDD corpus is partitioned by factory based on real meta data to simulate non-IID conditions.

9 Further Discussion on Grouping Mechanism↩︎

This appendix provides additional discussion and ablation results that clarify the role of our grouping mechanism in handling class imbalance. The central idea is to balance two competing effects in multi-class optimization under skewed data. Treating each class independently can be overly sensitive to noisy, small-sample gradients and may introduce unstable update scaling, whereas grouping multiple classes can reduce variance by aggregating signals but may increase bias when heterogeneous classes are merged. The following subsections examine these trade-offs from complementary angles. We first show the instability of per-class normalization in practice. We then validate our variance-based threshold selection and study how the number of groups affects the bias–variance and overfitting behavior.

9.1 Limitations of Per-Class Normalization (PCN)↩︎

Per-Class Normalization (PCN) rescales each class-specific gradient to unit norm. While this equalizes per-class magnitudes, it can induce unstable optimization dynamics and degrade generalization in multi-class regimes. First, aggregating \(C\) unit-norm class gradients yields an update whose norm can vary widely from near \(0\) to as large as \(C\), depending on the alignment among class directions. This introduces substantial scaling variability across iterations, which can impede stable convergence. Second, normalization does not reduce directional noise. When class-wise gradients fluctuate or are estimated from limited samples, PCN can amplify noisy directions and may overfit minority classes. Third, when \(C\) is large, many classes are often absent from a mini-batch, and maintaining class-wise gradients becomes computationally and memory intensive.

Figure 8 provides empirical evidence of these issues in a centralized learning setting on CIFAR-10 with imbalance rate \(20\), using the same backbone (ResNet18) for both methods. Compared with SGD, PCN exhibits pronounced training instability with large loss spikes, while validation accuracy fluctuates and intermittently collapses, consistent with high directional noise and unstable effective step scaling.

9.2 Validation of the Variance‐Based Grouping Threshold↩︎

To confirm that our grouping rule reliably identifies the optimal split, we sweep the threshold \(t\in\{1,\dots,9\}\), i.e. the number of rarest classes assigned to the minority group, and record test accuracy on CIFAR-10-LT under imbalance rates 20 and 100. As shown in Figure 9, the threshold \(t=6\) selected by minimizing within‐group variance (vertical dashed red line) coincides with the highest performance in both settings, demonstrating that our data‐driven rule matches the empirical optimum without exhaustive search.

9.3 The Effect of the Number of Groups↩︎

To confirm that two groups offer the best bias–variance trade-off beyond the setting in Figure [fig:loss95curve95by95H], we repeat the analysis on other settings. Figure 10 shows the training and validation losses when the number of groups \(H\) is set to \(\{2,3,4\}\) for CIFAR-10 and CIFAR-100.

In both benchmarks the pattern mirrors our earlier finding: validation loss stays lowest and most stable for \(H=2\), whereas \(H=3, 4\) starts to drift upward sooner, which is a signal of accelerated over-fitting. The training loss, by contrast, continues to fall for all values of \(H\), which widens the train–validation gap when more than two groups are used. These results reinforce the conclusion that splitting classes into exactly two groups strikes the right balance between reducing gradient variance and controlling over-fitting for CIFAR-10 and CIFAR-100.

10 Additional Experiment Results↩︎

In this appendix, we provide supplementary analyses to validate our design choices and assess the proposed algorithms under more challenging conditions.

10.1 Effectiveness of FedHOO↩︎

Table 4 compares the company’s metric obtained on the proprietary CDD benchmark when each federated optimizer runs with a fixed uniform resampling-rate versus when per-client rates are tuned by our FedHOO strategy. We include this result to demonstrate that FedHOO remains effective even when integrated with alternative training paradigms. FedHOO consistently boosts performance, confirming that our method is beneficial even when combined with optimizers other than FedCGNM.

To further isolate the benefit of FedHOO as a sampling-rate tuner, we compare it against two local rate-selection alternatives in the small-client regime. The first is a FAST-type [23] local exploration strategy, where each client independently explores its own rate using a bandit-style reward based on local loss decrease. The second is a locally optimal strategy, where each client independently selects its best local rate according to local validation and the resulting client-wise rates are combined for FL training. These baselines capture natural alternatives that tune resampling-rates locally rather than jointly.

Tables 5 and 6 show that FedHOO consistently outperforms both alternatives across FL optimizers. This supports the central design motivation of FedHOO: in FL, the best client-specific rates cannot generally be recovered by optimizing each client independently, because the final global model depends on interactions induced by federated aggregation. FedHOO explicitly searches this joint combinatorial space, while local strategies ignore cross-client interactions.

10.2 Resampling-rates in FedHOO↩︎

To better understand and illustrate the benefit of FedHOO, we plot and compare the trajectories of sampling-rate selection under FedHOO and standard HOO when each search method is run throughout training. In FedHOO, the selected sampling-rates correspond to the candidates yielding the best-performing model in each round, which is then adopted as the subsequent training initialization. As outlined in the implementation details, we restrict the search tree to a maximum depth of five.

Figure 11 illustrates how FedHOO adaptively selects per-client resampling-rates across different datasets and imbalance severities, and Figure 12 shows how HOO selects per-client resampling-rates on CIFAR-10-LT20 dataset. Each curve corresponds to one client, with the \(y\)-axis showing its selected sampling-rate, and the \(x\)-axis showing the communication rounds.

Under standard HOO, the algorithm explores only one branch of the search tree per round. As a result, the sampling-rate assigned to each client fluctuates heavily throughout training, and fail to converge in reasonable training time. The per-client trajectories remain noisy even after hundreds of rounds, reflecting the limited feedback that HOO gathers in each step.

In contrast, FedHOO leverages federated parallelism: by evaluating two candidate rates per client and linearly combining their updates, it effectively observes all \(2^K\) combinations at once. This parallel exploration dramatically accelerates the search. The sampling-rates quickly stabilize after the initial rounds, with each client settling into a distinct but consistent rate. The stabilized patterns observed across datasets in Figure 11 confirm that FedHOO both reduces variance and identifies effective configurations much earlier in training.

Overall, these plots highlight the key difference: while HOO suffers from slow and noisy adaptation due to sequential exploration, FedHOO achieves rapid and stable convergence by exploiting the structure of federated training.

10.3 Alignment Effect of FedCGNM↩︎

Beyond mitigating directional noise, FedCGNM can also improve client alignment, which is a primary cause of performance degradation in federated learning [3]. Under heterogeneous and imbalanced data, clients often produce updates whose class-wise magnitudes differ substantially, which can amplify misalignment after aggregation. FedCGNM addresses this issue by normalizing each group update to a common unit scale, making client contributions more comparable and improving the alignment of aggregated directions.

To empirically assess this effect, we track an alignment metric over communication rounds and compare it against a stochastic-gradient baseline. At each round \(t\), we compute the cosine similarity between a client update direction and a server’s update direction (the corresponding aggregated update). We report (i) the similarity induced by FedCGNM updates and (ii) the similarity induced by stochastic-gradient updates under the same training pipeline. Figure 13 shows that FedCGNM typically achieves higher similarity across rounds, and the per-round difference (bottom plot) is predominantly positive, supporting that FedCGNM improves update alignment across clients.

10.4 FedCGNM in Large-Federation Regimes↩︎

To further assess robustness, we evaluate FedCGNM under the challenging condition beyond the main experiments: large numbers of clients with partial participation.

10.4.0.1 Large-client regime with partial participation.

We evaluate FedCGNM when the number of clients is large (\(K=100\)) and only small fraction participate (10%) per round, which is a scenario that amplifies variance and typically harms optimization. We also use \(E=1\) for this scenario. Table 7 reports results on CIFAR-10-LT and CIFAR-100-LT with imbalance rates \(20\) and \(100\). While all methods degrade in this setting, FedCGNM consistently achieves the best performance, surpassing both reweighting-based and per-group gradient-normalization baselines.

Together, these results demonstrate that FedCGNM retains its advantage in more challenging federated environments: it scales to non-IID data and large-client regimes where baseline methods suffer the most.

10.5 Robustness under Various Data Heterogeneity↩︎

To further evaluate robustness under different degrees of client heterogeneity, we conduct additional experiments on CIFAR-100-LT20 with \(K=5\) clients. Table 8 reports the results under two heterogeneous data partitions. FedCGNM consistently outperforms the baselines, and FedHOO further improves performance in both settings. These results suggest that the proposed class-grouped normalized momentum remains effective under different degree of heterogeneity, while FedHOO provides additional gains by selecting client-specific resampling-rates.

10.6 Per-Class Accuracy Distribution↩︎

From Figure 14, FedCGNM (blue) shifts the per-class accuracy distribution markedly to the right compared to FedAvg (orange), substantially reducing the number of near-zero classes. Whereas FedAvg produces a heavy tail of classes below 0.1 accuracy, FedCGNM elevates most classes into a moderate accuracy range. FedCGN also achieves a more balanced distribution, confirming the effectiveness of the grouping strategy in balancing class-wise performance.

10.7 Empirical Diagnostics for Assumption 5↩︎

Assumption 5 excludes the degenerate case where the group-wise momentum is both nearly perfectly aligned with the corresponding gradient and nearly identical in magnitude to the scaled gradient. We empirically examine this condition by logging the two quantities appearing in the assumption: the gradient–momentum cosine similarity and the magnitude gap.

Figure 15 shows the results on AdultIncome and CIFAR-100-LT. On AdultIncome, the cosine similarity is often high, but the magnitude gap remains consistently positive and non-negligible. Thus, the magnitude-gap branch of the assumption is supported. On CIFAR-100-LT, the cosine similarity remains far below perfect alignment throughout training, and the magnitude gap is also positive. Thus, the low-alignment branch is strongly supported. These complementary behaviors suggest that Assumption 5 is mild in practice, indicating that the gradient and momentum are not simultaneously perfectly aligned and equal in magnitude along the observed training trajectories.

References↩︎