Augmented Self-bootstrap Distillation (ASD)

• The paper introduces ASD, a logit-level self-bootstrap distillation method that enhances feature representations and classifier calibration in federated learning on non-IID, long-tailed data.
• ASD employs dual augmentation where weakly augmented samples serve as teachers to distill knowledge to strongly augmented student samples without relying on external networks.
• Combined with distribution-aware logit adjustment, ASD mitigates class imbalance and drives feature prototypes toward near-optimal neural-collapse geometry, significantly improving accuracy.

Augmented Self-bootstrap Distillation (ASD) is a logit-level client-side self-distillation mechanism central to the FedYoYo framework for federated learning (FL) under heterogeneous and long-tailed data distributions. ASD enables each client to utilize its own labeled data to bootstrap improved feature representations by distilling information from weakly augmented samples (serving as teachers) to strongly augmented samples (as students). The mechanism operates solely within a single model, leverages no external data or auxiliary networks, and is tightly integrated with distribution-aware logit adjustment (DLA) to correct for pervasive class imbalance. ASD’s design addresses the fundamental limitations of prior neural-collapse-inspired and self-distillation methods in federated settings—specifically, the inability to drive feature spaces toward optimal collapse and to correct for classifier bias under non-IID and long-tailed data regimes (Yan et al., 10 Mar 2025).

1. Motivation and Conceptual Overview

Federated learning with real-world data is impeded by severe heterogeneity, as local client distributions are typically non-IID and may present significant long-tailed imbalance. These challenges manifest in poor local feature representations, classifier bias toward majority classes, and substantial drift of local models away from centralized optima upon aggregation. Previous approaches fall into two main categories: (1) neural-collapse-inspired strategies (e.g., FedETF, FedLoGe) impose synthetic classifier structures but fail to adequately regularize the feature backbone in highly heterogeneous conditions, and (2) self-distillation or self-supervised pipelines (e.g., BYOL, SimCLR), which require auxiliary models or unlabeled data and are frequently uncalibrated for skewed class distributions.

ASD is introduced as a lightweight, logit-space knowledge distillation protocol operating on local data. Each labeled sample is stochastically augmented twice: a weak augmentation (e.g., random crop/flip/rotation) and a strong augmentation (e.g., via AutoAugment or RandAugment). The local model produces logits for both versions, enabling distillation from the weak (teacher) to the strong (student) view. ASD strictly considers only correctly classified “teacher” samples, alleviating the amplification of noise from erroneous pseudo-labels. The approach obviates the need for auxiliary networks and integrates directly with a distribution-aware softmax to actively mitigate class imbalance via DLA.

2. Formalism and Notation

Consider client $k$ with local labeled data set $\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$. For each $x_i$, weak and strong augmentations yield $x_i^w$ and $x_i^s$. The local model, $f_k(\cdot)$, outputs logits for each class $y=1 \ldots C$. A fused class prior $\pi_{\rm mix}$ is computed as

$\pi_{\rm mix}=(1-\gamma)\,\pi_g + \gamma\,\pi_k,$

where $\pi_k$ is the local effective prior (estimated via Pearson-area feature correlations) and $\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$0 is the FedAvg-aggregated global prior; $\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$1 determines the local/global tradeoff.

Using temperature $\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$2 (default 1.5), the distribution-aware softmax for input $\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$3 is

$\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$4

The ASD loss is defined via Kullback-Leibler divergence between the “teacher” (weak) and “student” (strong) views: $\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$5 where $\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$6, $\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$7, and only correctly classified $\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$8 are retained (others masked). The loss is combined with cross-entropy classification under DLA, resulting in a total local loss of

$\mathcal D_k=\{(x_i,y_i)\}_{i=1}^{n_k}$9

with $x_i$0 a tunable coefficient.

3. Local Client Algorithmic Workflow

A typical local training round under ASD consists of the following procedure (pseudocode structure given in (Yan et al., 10 Mar 2025)):

1. Input: Global model parameters $x_i$1, local prior $x_i$2, global prior $x_i$3.
2. Data Augmentation: For each minibatch, generate weak ($x_i$4) and strong ($x_i$5) augmentations.
3. Forward Propagation: Compute logits for both augmentations using $x_i$6.
4. Fusion of Priors: Calculate $x_i$7 from $x_i$8 and $x_i$9 with fusion rate $x_i^w$0.
5. Distribution-aware Softmax: Convert logits to class probabilities for both augmentations.
6. Teacher Correctness Masking: Retain only samples where the weak-augmented prediction matches the true label.
7. Loss Computation:
   • Compute ASD loss ($x_i^w$1) as KL divergence between weak and strong distributions for masked samples.
   • Calculate DLA classification loss ($x_i^w$2) via two-view cross-entropy under the same softmax.
   • Aggregate as $x_i^w$3.
8. Model Update: Backpropagate and update parameters.
9. Prior Update: Refine local class prior $x_i^w$4 via updated Pearson-area statistics. 10. Upload: Return updated $x_i^w$5 and $x_i^w$6 to server.

Default hyperparameters are: $x_i^w$7, $x_i^w$8, and $x_i^w$9; performance is stable across these ranges.

4. Integration with Distribution-aware Logit Adjustment (DLA)

Distribution-aware Logit Adjustment ensures calibration across class frequencies, directly addressing the adverse effects of long-tailed and non-IID data in federated settings. DLA employs the same fused prior $x_i^s$0 for adjusting softmax scores, balancing global and local class proportion knowledge. Both ASD loss and DLA classification loss use this adjusted softmax, ensuring that knowledge distillation remains invariant to class frequency and that minority classes are not overwhelmed during training.

For each client, the local prior $x_i^s$1 is estimated in feature space and periodically averaged across the federation to obtain $x_i^s$2, with $x_i^s$3 blending both. The combination of ASD’s invariance-enforcing KL loss and DLA’s balanced classification anchors learning to a regime where neither class bias nor poor feature utilization dominates, enabling effective aggregation and reduced drift.

5. Theoretical Insights

ASD instantiates a self-bootstrap process in the logit space, imposing consistency between weak and strong sample augmentations. This constraint encourages learning representations that inhabit flatter loss basins and resist overfitting to idiosyncratic client distributions. Notably, performing distillation in the logit space (rather than penultimate-layer features) embeds calibration properties into the learned representations themselves, aligning class prototypes across the heterogeneous federation.

Empirical and theoretical analyses demonstrate that ASD, in conjunction with DLA, results in global feature prototypes that approach the angles associated with an Equiangular Tight Frame (ETF) as predicted by neural-collapse theory, even though no explicit ETF structure is enforced. This suggests that ASD-based self-distillation is effective at harmonizing representation geometry under challenging data regimes.

6. Empirical Performance and Ablation Findings

Comprehensive evaluations on benchmarks such as CIFAR-100-LT and CIFAR-10-LT across non-IID and long-tailed settings establish ASD’s efficacy:

• On CIFAR-100-LT with imbalance factor (IF) 100 and Dirichlet $x_i^s$4=0.5, incrementally introducing augmentations and ASD yields overall accuracy jumps from 33.34% (baseline) to 41.99% (+8.65%) with ASD, and to 46.13% (+12.79%) when combined with DLA.
• ASD boosts “Many-shot” class performance (+8.42% with RandAugment), while DLA is critical to elevate “Few-shot” accuracy (from 13.93% to 28.63%).
• The (weak → strong) distillation direction outperforms alternatives (46.13% vs. 42.92% for strong → strong).
• On vanilla non-IID CIFAR-10 (α=0.01), FedYoYo (with ASD) reaches 81.70% vs. 72.31% for FedETF (+9.39%).
• Under global long-tailed, non-IID settings, FedYoYo exceeds prior FL methods (e.g. 81.45% vs. 70.63% for Fed-Grab, +10.82%).
• Feature-space analysis reveals that class prototypes under FedYoYo approach the ETF-predicted 96.4° angle, in contrast to substantial gaps observed in prior methods.
• The global-to-local discrepancy is sharply reduced as measured after aggregation and local re-adaptation.

These results demonstrate that ASD, especially when used with DLA, mitigates the representation and calibration failures endemic to federated learning under mixed heterogeneity, enabling performances near or beyond centralized baselines in long-tailed regimes (Yan et al., 10 Mar 2025).

7. Contextual Significance and Implications

Augmented Self-bootstrap Distillation represents a principled synthesis of modern self-distillation and calibration approaches adapted to the federated learning paradigm. By relying exclusively on local operations and logit-space distillation from weak to strong augmentations, ASD sidesteps the need for auxiliary teacher networks and the computational or data communication burdens they entail. Its demonstrated ability to induce near-neural-collapse geometry and dramatically narrow the generalization gap across federated and centralized models positions ASD as a central technique for future FL systems operating in real-world, heterogeneous, and unbalanced environments. A plausible implication is that similar one-model, logit-space self-distillation mechanisms, when paired with dynamic class-prior calibration, may generalize to other distributed learning architectures with analogous challenges.