Projector-only EMA in Self-Supervised Learning

• The paper introduces projector-only EMA, which applies EMA updates only to the projector block, significantly reducing computational overhead compared to full-encoder EMA.
• Projector-only EMA is defined as restricting momentum updates to the projector, targeting layers with high gradient volatility to stabilize training dynamics.
• Empirical results reveal that projector-only EMA nearly matches full-encoder EMA in top-1 accuracy while offering substantial savings in compute and memory resources.

Projector-only Exponential Moving Average (EMA) is a variant of the momentum encoder paradigm employed in self-supervised learning (SSL), which restricts the application of the EMA update to only the projector module at the top of the encoder stack. This contrasts with the conventional approach of applying EMA to the entire encoder, as seen in frameworks such as MoCo, BYOL, and DINO. Projector-only EMA retains almost all of the performance and stability benefits associated with full-encoder EMA, while significantly reducing compute and memory overhead by avoiding redundant backbone computations. This method leverages detailed observations regarding gradient stability within network layers and offers an optimized trade-off between computational efficiency and representation quality (Pham et al., 2022).

1. Motivation and Background

In standard SSL settings (e.g., MoCo, BYOL), a "teacher" network is constructed by applying an exponential moving average (EMA, or momentum) update to a copy of the student network, generating more stable and consistent targets. This often doubles both memory and computation, since two deep encoders (student/backbone and momentum teacher) are stored and evaluated per view. The motivation for projector-only EMA arises from empirical observations that most of the training instability—measured through loss gradients and parameter "spikes"—concentrates in the uppermost network layers (block4 and projector), which experience both larger and more volatile gradients than earlier layers [(Pham et al., 2022): Fig. 6]. EMA's stabilizing effect is therefore disproportionately beneficial when focused on those final layers.

2. Methodology: EMA Restriction to Projector

Conventional full-encoder EMA performs the following parameter update at each training step $t$:

$$\xi^{(t)} = \alpha\,\xi^{(t-1)} + (1-\alpha)\,\theta^{(t)}$$

where $\theta$ denotes the online encoder parameters, $\xi$ the momentum teacher parameters, and $\alpha \in [0,1]$ the momentum coefficient (typical values: 0.99–0.9995).

Projector-only EMA restricts this update to the projector block, so that only the projector’s parameters receive EMA smoothing, while the backbone is shared between student and teacher ($\theta_b = \xi_b$, i.e., no EMA smoothing on the backbone). The update rule for the projector block is:

$$\xi_p^{(t)} = \alpha\,\xi_p^{(t-1)} + (1-\alpha)\,\theta_p^{(t)}$$

Here, $\theta_p$ and $\xi_p$ denote the projector parameters of the online and momentum branches, respectively [(Pham et al., 2022): Eq. in Sec. 2].

Ablation studies where EMA is selectively applied to Conv1, intermediate blocks, or projector reveal that only the latter yields substantial accuracy gains (from 43.5% with no EMA to 59.6% with projector EMA; nearly matching the 62.0% of full-encoder EMA on CIFAR-100, ResNet-18, 200 epochs) [(Pham et al., 2022): Fig. 5a].

3. Empirical Analysis: Stability and Gradient Dynamics

Gradient-norm analysis during SSL training (e.g., in BYOL) shows that the $\infty$-norm of the gradients with respect to the output of each encoder block is highest and most volatile in the final block and projector. Absence of momentum leads to the occurrence of weight "spikes" in these top layers, which correlate with sudden drops in linear classification accuracy [(Pham et al., 2022): Figs. 3, 4]. EMA's "memory" of prior weights dampens these instabilities, improving the smoothness of learning dynamics. Notably, backbone layers (blocks 1–3) exhibit comparatively small, stable gradients and do not benefit substantially from EMA, as their parameters are less susceptible to destabilizing fluctuations.

4. Computational Efficiency Gains

Full-encoder EMA doubles the computational and memory requirements per augmented view by requiring forward passes and storage for both student and momentum teacher backbones. By restricting EMA to the projector (typically a small MLP), projector-only EMA reduces the extra forward pass to two lightweight projectors, while sharing the expensive backbone:

• Memory: Only the small projector’s parameters and activations are duplicated for the teacher forward; backbone activations/parameters are not.
• Compute: Extra computation of a full backbone forward per step is eliminated.

For example, in DINO training with ResNet-50 on ImageNet-1k (200 epochs, 2×A6000 GPUs), per-epoch time is reduced from 56.5 min (full-encoder EMA) to 38.5 min (projector EMA), with virtually unchanged top-1/top-5 accuracy (67.9%/88.3% vs. 67.8%/88.2%, respectively) [(Pham et al., 2022): Table 5].

5. Experimental Results and Comparisons

Representative results across various SSL frameworks and datasets highlight the efficacy of projector-only EMA. On CIFAR-100 (ResNet-18, 200 epochs):

• No EMA baseline (e.g., BYOL $$\xi^{(t)} = \alpha\,\xi^{(t-1)} + (1-\alpha)\,\theta^{(t)}$$0): top-1 accuracy = 43.48%
• Full-encoder EMA (BYOL $$\xi^{(t)} = \alpha\,\xi^{(t-1)} + (1-\alpha)\,\theta^{(t)}$$1): top-1 = 62.01%
• Projector-only EMA: top-1 = 59.61% (–2.4% vs. full EMA)

For ImageNet-1k with DINO:

Setting	Top-1 Acc.	Top-5 Acc.	Compute/Memory Cost
Full-encoder EMA	67.9%	88.3%	2× backbones
Projector-only EMA	67.8%	88.2%	1× backbone + 2× proj

The marginal (≤0.1%) drop in accuracy for large-scale settings suggests that the trade-off will be favorable in most practical circumstances.

6. Hyperparameter Sensitivity and Robustness

Projector-only EMA demonstrates reduced sensitivity to the momentum coefficient $$\xi^{(t)} = \alpha\,\xi^{(t-1)} + (1-\alpha)\,\theta^{(t)}$$2. When $$\xi^{(t)} = \alpha\,\xi^{(t-1)} + (1-\alpha)\,\theta^{(t)}$$3 is set extremely high (e.g., 0.9995 for >500 epoch runs), full-encoder EMA methods suffer a performance collapse (–6% to –19%), but projector-only EMA experiences only a –2% to –4% drop [(Pham et al., 2022): Table 6, Fig. 7]. This robustness is attributed to the residual stability in the shared backbone and the fast adaptation retained by restricting EMA to the jumpy aggregator at the top of the stack.

Applying BatchNorm to the output of block4 and/or the projector can further stabilize training and close any residual accuracy gap [(Pham et al., 2022): Appendix Sec. C, Table 10].

7. Practical Adoption and Recommendations

To use projector-only EMA in any student–teacher (momentum-based) SSL codebase, one replaces the standard full-encoder EMA update with the projector-restricted update. The backbone parameters remain identical for student and teacher (i.e., set $$\xi^{(t)} = \alpha\,\xi^{(t-1)} + (1-\alpha)\,\theta^{(t)}$$4 for the backbone). Recommended settings include $$\xi^{(t)} = \alpha\,\xi^{(t-1)} + (1-\alpha)\,\theta^{(t)}$$5 for 200-epoch runs, increasing up to 0.9995 for longer training schedules, with projector-only EMA being substantially more resilient to high $$\xi^{(t)} = \alpha\,\xi^{(t-1)} + (1-\alpha)\,\theta^{(t)}$$6 values.

A plausible implication is that by focusing computational resources on stabilizing only the volatile upper layers, projector-only EMA provides an efficient, high-throughput alternative for scenarios where compute/memory constraints are limiting or scaling to large datasets is required.

Summary Table: Comparison of EMA Strategies

EMA Application	Top-1 Accuracy (CIFAR-100)	Computational Cost	Stability Benefit
None	43.48%	1× backbone + 1× projector	Unstable/projector
Full-encoder	62.01%	2× backbone + 2× projector	High
Projector-only	59.61%	1× backbone + 2× projector	Nearly as high

Projector-only EMA, by targeting regularization and smoothing to the layers exhibiting maximum instability, achieves a favorable balance between maintaining SSL performance and dramatically reducing training resource requirements (Pham et al., 2022).