Understanding Self-Supervised Learning Dynamics without Contrastive Pairs

In plain words

Self-supervised learning trains a network without labels by asking it to match two augmented views of the same image. Some recent methods, BYOL and SimSiam among them, do this with matching pairs only. There are no "negative" examples that push different images apart. The obvious worry is that the network could map every input to the same output and satisfy the matching goal trivially. These methods avoid that by adding a small learnable predictor, a stop-gradient on one branch, and, in BYOL, a slowly moving copy of the network. Why this works has been unclear.

The paper studies a deliberately simple linear version of this setup and works out its learning dynamics. In that model it explains what each ingredient does. The predictor has to learn at a suitable speed. The slowly moving copy acts like a curriculum that raises the goal in steps. Weight decay ties the strength of the online network to that of the predictor. The analysis is exact only in a simplified setting, but the authors check that its predictions match ablation studies on real image benchmarks.

The practical outcome is DirectPred. Instead of training the predictor by gradient descent, it estimates the correlation matrix of the network's own outputs and sets the predictor directly from that matrix's eigenvectors and eigenvalues. On ImageNet it comes close to the default two-layer BYOL predictor and beats a trained linear predictor. The basic form uses no gradient steps for the predictor at all.

TL;DR

Theory plus a method for non-contrastive SSL (BYOL, SimSiam). These methods avoid collapse without negative pairs because the predictor, the stop-gradient, the EMA target and weight decay jointly keep the online network growing alongside the predictor §2, §3. DirectPred replaces the trained linear predictor with a matrix built from a moving-average correlation estimate of the online features §4, eq. 18–19. Author-reported: it beats a gradient-trained linear predictor on STL-10 and ImageNet, and on ImageNet at 300 epochs it is comparable with the default two-layer BYOL predictor tab. 6, tab. 9. The official code is CC BY-NC licensed, so use it for reading only.

Problem

Non-contrastive SSL trains each branch to match the other using positive pairs only. The open question is why these methods do not collapse to trivial, constant representations §1. Earlier practice reported that the predictor's learning rate, the EMA rate and weight decay all matter, but without a theory to explain how §1. The predictor is also hard to get right by hand: plugging in its optimal least-squares solution performs poorly, and gradient training is needed to reach good accuracy tab. 1. The analysis assumes a linear network, isotropic data and augmentation noise, and a symmetric predictor §3.

Method

The paper has three parts: (a) a continuous-time model of BYOL and SimSiam dynamics in a bias-free two-layer linear network §2; (b) a reduction of that model to three scalar equations per eigenmode under simplifying assumptions, used to explain the effects of hyperparameters §3; (c) DirectPred, a predictor set without gradients from the input correlation matrix §4.

Setup §2. The online network is W \in \mathbb{R}^{n_2 \times n_1}, the predictor is W_p \in \mathbb{R}^{n_2 \times n_2} and the target is W_{ema}. Two augmented views x_1, x_2 of a sample give f_1 = W x_1 and f_{2,ema} = W_{ema} x_2. The objective is eq. 1:

J(W, W_p) = \tfrac{1}{2}\, \mathbb{E}\left[ \left\| W_p f_1 - \mathrm{StopGrad}(f_{2,ema}) \right\|_2^2 \right] eq. 1

Step 1: gradient-flow dynamics §2, Lemma 1. In the limit of large batches and small learning rates, the dynamics are eq. 2–4:

\dot W_p = \alpha_p \left( -W_p W (X + X') + W_{ema} X \right) W^\top - \eta W_p eq. 2

\dot W = W_p^\top \left( -W_p W (X + X') + W_{ema} X \right) - \eta W eq. 3

\dot W_{ema} = \beta \left( -W_{ema} + W \right) eq. 4

Here X = \mathbb{E}[\bar x \bar x^\top], where \bar x is the average augmented view of a sample, and X' is the augmentation covariance averaged over samples §2. The learning-rate ratios are \alpha_p (predictor) and \beta (target), \eta is weight decay, and in discrete time the EMA rate is \alpha\beta = 1 - \gamma_{ema} and the weight decay is \bar\eta = \alpha\eta §2. SimSiam is the special case W_{ema} = W, with the stop-gradient kept and the EMA equation ignored §2.

Step 2: balancing invariant §2, Theorem 1. Whatever the target dynamics, the online and predictor weights satisfy eq. 5:

W(t) W(t)^\top = \alpha_p^{-1} W_p(t)^\top W_p(t) + e^{-2\eta t} C eq. 5

with C fixed by the initialisation §2, Thm. 1. Weight decay removes the constant term, so the online and predictor weights end up balanced §2.

Step 3: stop-gradient is necessary §2, Theorem 2. With W_{ema} = W and the stop-gradient removed, the online weights obey \frac{d}{dt}\mathrm{vec}(W) = -H(t)\,\mathrm{vec}(W) with a positive semi-definite H eq. 6. If its smallest eigenvalue stays bounded below by a positive constant, W(t) \to 0, which is representational collapse §2:

\frac{d}{dt}\mathrm{vec}(W) = -H(t)\,\mathrm{vec}(W) eq. 6

where H = X' \otimes (W_p^\top W_p + I) + X \otimes \tilde W_p^\top \tilde W_p + \eta I and \tilde W_p = W_p - I §2, Thm. 2.

Step 4: simplifying assumptions §3. Three assumptions reduce the dynamics. (i) Proportional EMA: W_{ema}(t) = \tau(t) W(t) §3, Assumption 1. (ii) Isotropic data and augmentation: data has zero mean and identity covariance, and the augmentation has mean x and covariance \sigma^2 I, so X = I and X' = \sigma^2 I §3, Assumption 2. (iii) Symmetric predictor: W_p is initialised symmetric and its flow is symmetrised §3, Assumption 3. Let F = W X W^\top (the correlation of predictor inputs). Under these, eq. 7 holds:

\dot W_p = -\tfrac{\alpha_p}{2}(1+\sigma^2)\{W_p, F\} + \alpha_p \tau F - \eta W_p eq. 7

with \{A, B\} = AB + BA §3.1.

Step 5: eigenspace alignment §3.1, Theorem 3. The commutator [F, W_p] decays as \frac{d}{dt}[F, W_p] = -[F, W_p] K - K [F, W_p] eq. 8, where eq. 9:

K(t) = (1+\sigma^2)\left[ \tfrac{\alpha_p}{2} F(t) + W_p^2(t) - \tfrac{\tau}{1+\sigma^2} W_p(t) \right] + \tfrac{3}{2}\eta I eq. 9

If \inf_t \lambda_{\min}(K(t)) = \lambda_0 > 0, then \|[F(t), W_p(t)]\|_F \le e^{-2\lambda_0 t}\|[F(0), W_p(0)]\|_F \to 0 eq. 10 §3.1, Thm. 3. So the eigenspaces of W_p and F align over time §3.1. The paper checks this on BYOL with ResNet-18 on STL-10 fig. 2.

Step 6: decoupled per-mode dynamics §3.1. Once the eigenspaces agree, with W_p = U \Lambda_{W_p} U^\top and F = U \Lambda_F U^\top, each mode j with eigenvalues p_j (of W_p) and s_j (of F) follows eq. 11–13:

\dot p_j = \alpha_p s_j \left[ \tau - (1+\sigma^2) p_j \right] - \eta p_j eq. 11

\dot s_j = 2 p_j s_j \left[ \tau - (1+\sigma^2) p_j \right] - 2\eta s_j eq. 12

s_j \dot\tau = \beta (1 - \tau) s_j - \tfrac{1}{2}\tau \dot s_j eq. 13

This is a three-dimensional system for BYOL and a two-dimensional one for SimSiam (where \tau = 1) §3.1.

Step 7: exact invariant §3.2. The quantity eq. 14 is conserved:

s_j(t) = \alpha_p^{-1} p_j^2(t) + e^{-2\eta t} c_j, \qquad c_j = \alpha_p^{-1} p_j^2(0) - s_j(0) eq. 14

Without weight decay, c_j is never forgotten. With weight decay, the state approaches the parabola s_j = \alpha_p^{-1} p_j^2 §3.2.

Step 8: dynamics on the parabola §3.2. Setting c_j = 0 and using the target's slow integral eq. 15:

\tau(t) = p_j^{-1}(t)\, \beta e^{-\beta t} \int_0^t p_j(t') e^{\beta t'}\, dt' eq. 15

with \tau(0) = 0, the predictor follows eq. 16:

\dot p_j = p_j^2 \left[ \tau(t) - (1+\sigma^2) p_j \right] - \eta p_j eq. 16

For 0 < \eta \le \tau^2 / (4(1+\sigma^2)) there are three fixed points, p^*_{j\pm} = \frac{\tau \pm \sqrt{\tau^2 - 4\eta(1+\sigma^2)}}{2(1+\sigma^2)} and p^*_{j0} = 0. The collapsed point 0 and the upper point p^*_{j+} are stable, and p^*_{j-} is unstable, which separates their basins §3.2. For \eta > \tau^2/(4(1+\sigma^2)) only the collapsed point is stable §3.2. Unlabelled equations for these fixed points and the stability criterion are in §3.2.

Step 9: positive-definite condition §3.2. Let \Delta_j = p_j[\tau - (1+\sigma^2)p_j] - \eta. Eigenspace alignment stays stable when eq. 17:

\Delta_j < \tfrac{1}{2}\left[ \alpha_p (1+\sigma^2) s_j + \eta \right] eq. 17

Together with \Delta_j > 0 (needed to leave the collapsed point) this explains why a faster predictor and more weight decay help escape collapse, while very fast predictors stop the online network from growing §3.2, Obs. 1–3.

Step 10: DirectPred §4. Estimate the correlation \hat F of the predictor inputs, eigendecompose it as \hat F = \hat U \hat\Lambda_F \hat U^\top with eigenvalues s_j, and set eq. 18:

p_j = \sqrt{s_j} + \epsilon \max_j s_j, \qquad W_p = \hat U\, \mathrm{diag}[p_j]\, \hat U^\top eq. 18

The estimate is a moving average over batches, without centring the features eq. 19 §4:

\hat F = \rho \hat F + (1 - \rho)\, \mathbb{E}_B[f f^\top] eq. 19

The paper states that the magnitude of s_j does not matter because of ℓ2 normalisation in the experiments §4. A hybrid variant runs the eigendecomposition only every freq minibatches and uses regular gradient updates in between §4.

Pseudocode: DirectPred predictor update. This follows eq. 18–19 and §4. Names are ours; the paper gives no code for this block.

# DirectPred predictor update  [§4] [eq. 18-19]
# Input:  features (B, d) from the online network, one batch
# State:  F_hat (d, d), a running correlation estimate
# Output: predictor weight W_p (d, d), symmetric

def update_correlation(F_hat, features, rho):
    # Moving average of the batch second moment, no centring  [§4] [eq. 19]
    batch_moment = features.T @ features / batch_size  # (d, d)
    return rho * F_hat + (1 - rho) * batch_moment

def direct_predictor(F_hat, eps):
    # Eigendecomposition of the symmetric estimate  [§4] [eq. 18]
    scores, basis = eigh(F_hat)  # scores = s_j, basis = U_hat
    # Square root of each eigenvalue plus eps times the largest one  [eq. 18]
    gains = sqrt(scores) + eps * max(scores)  # p_j
    return basis @ diag(gains) @ basis.T  # (d, d)

Pseudocode: where the predictor sits in one training step. This is the linear-model version from §2 with the loss normalised as in the experiments; it is a sketch, not the authors' code.

# One step of BYOL with a DirectPred predictor  [§2] [§4] [app. C]
# Input:  two augmented views, step counter, update period freq
# Side effects: updates the online and target networks, F_hat and W_p

def training_step(view_1, view_2, step, freq):
    # Online features, before the predictor  [§2]
    feat_1 = online(view_1)
    feat_2 = online(view_2)
    # Correlation estimate from both views  [§4] [eq. 19]
    F_hat = update_correlation(F_hat, concat(feat_1, feat_2), rho)
    # Re-set the predictor every freq steps  [§4] [eq. 18]
    if step % freq == 0:
        W_p = direct_predictor(F_hat, eps)
    # Target branch has no gradient; EMA update with gamma = 0.996  [app. C]
    target_2 = stop_gradient(target(view_2))
    # Loss uses l2-normalised vectors in the experiments  [§4] [code]
    loss = l2_regression(W_p @ feat_1, target_2)
    return loss

Hyperparameters

Results

Author-reported. STL-10 and CIFAR-10 use ResNet-18 with the linear-probe top-1 metric app. C. ImageNet uses ResNet-50 with top-1 and top-5 §4.

Other results: plugging the optimal least-squares predictor into BYOL gives 32–41% top-1 on STL-10 at 100 epochs, far below gradient training (about 74–75%) tab. 1. Without EMA, a symmetric predictor collapses to 36.62 ± 1.85 top-1 on STL-10 with a one-layer predictor, against 72.85 ± 0.16 for the regular predictor tab. 3. DirectPred is run with EMA in all reported experiments tab. 6.

Paper vs code

Official code: the paper's footnote links facebookresearch/luckmatters, directory ssl/, at commit 625b5f7de69dc321f24376f5c6fe3446af3a364b. The STL-10 and CIFAR-10 DirectPred path is in ssl/real-dataset. The repository's root LICENSE is Attribution-NonCommercial 4.0, so the code is cited for reading only and nothing is copied.

• Eigenvalue map. Paper: p_j = \sqrt{s_j} + \epsilon \max_j s_j eq. 18. Code: eigenvalues are divided by the largest one byol_trainer.py, raised to the power 1/dyn_convert (default 2 in byol_config.yaml), then dyn_eps is added to the normalised value byol_trainer.py and clamped below at 1e-4 byol_trainer.py. Dividing the paper's map by \sqrt{\max_j s_j} gives \sqrt{s_j/\max s} + \epsilon\sqrt{\max s}, which equals the code only when the largest eigenvalue is 1. The code's clamp at 1e-4 is not in the paper.
• ε value. The README's DirectPred command uses dyn_eps=0.01 and dyn_lambda=0.3 README.md. The paper's best STL-10 cell is ε = 0.1 with ρ = 0.3 (77.86 ± 0.16) tab. 6. The text cites 77.38% for "additional regularisation through ε" §4, a value that appears in no table. Which setting produced the headline numbers cannot be settled from the paper and code alone.
• Correlation estimate. Paper: moving average with ρ eq. 19. Code: the accumulator computes \lambda c + (1-\lambda) d with \lambda = dyn_lambda byol_trainer.py, and the batch matrix averages the outer products of both views byol_trainer.py byol_trainer.py. The two agree. The code starts the average from the first batch; the paper does not state an initial value.
• Input normalisation. Paper: ℓ2 normalisation is used so that the magnitude of s_j does not matter §4. Code: the correlation is taken on the raw online outputs before the predictor byol_trainer.py with no visible ℓ2 normalisation of those features. Magnitude invariance comes from dividing by the largest eigenvalue byol_trainer.py. The loss is ℓ2-normalised by default byol_trainer.py.
• Centring. Paper: no centring §4. Code: dyn_zero_mean defaults to false byol_config.yaml. Consistent.
• Update period. Paper: freq = how often the eigendecomposition runs §4. Code: the update happens when niter % predictor_freq == 0 byol_trainer.py. Consistent.
• Config defaults. The default config uses predictor_reg: symmetric, dyn_lambda: 0.8, dyn_eps: 0.0 and balance_type: clamp byol_config.yaml byol_config.yaml byol_config.yaml byol_config.yaml. DirectPred needs the overrides given in the README.
• Version. The README says the DirectPred command was validated at commit cb23d10c3018df6bf275ad537f23675c8a627253 README.md. The clone used here is the repository's HEAD at 625b5f7, recorded above. Whether the two differ in these files was not checked.
• ImageNet. Paper: the ImageNet runs use the authors' own BYOL implementation §4. The repository has no ImageNet pipeline. The ImageNet numbers cannot be checked against code.
• Hyperparameters that do agree. Learning rate 0.03, momentum 0.9, weight decay 0.0004, batch 128, EMA 0.996 app. C match byol_config.yaml byol_config.yaml byol_config.yaml. The predictor lr equals the online lr byol_config.yaml, which matches α_p = 1.

Pitfalls

• Two ε conventions. Eq. 18 scales ε by the unnormalised maximum eigenvalue, but the code normalises first and then adds ε. Values of ε do not transfer between the two without rescaling. Check which one a reimplementation follows before comparing numbers.
• Which ε? The best reported setting is ε = 0.1 tab. 6; the README uses 0.01. Start from the paper's table value, and expect a small gap to the README's setting.
• The "+5%" claim. The introduction says DirectPred beats a gradient-trained linear predictor by 5% on STL-10 and ImageNet at 60 epochs §1. The STL-10 table gives 77.86 against 74.51, a gap of about 3.4 points tab. 6 tab. 3. The 5% gap holds on ImageNet at 60 epochs tab. 9.
• Symmetric predictors without EMA. A symmetric predictor without EMA fails badly unless the predictor learning rate is above 1 or weight decay is applied only to the predictor tab. 3 tab. 5 fig. 6. DirectPred's reported results all use EMA tab. 6, so do not drop EMA without checking the results.
• Positive constant c_j. The variant p_j = sqrt(max(s_j − c_j, 0)) with a positive c_j collapses. With freq=1 the result is 46.57 ± 18.43 at c_j = 0.1 tab. 7. Keep c_j ≤ 0 or use the default.
• Cost. Every update needs a d×d eigendecomposition. With d = 256 on ImageNet §4 and freq=1, this runs every minibatch. The paper notes freq>1 reduces the cost and can improve accuracy §4, tab. 8.
• Predictor updates between eigendecompositions. With freq>1 the predictor receives regular gradient updates in between §4. The freq=1 variant has no such step.
• Table 8 duplicate values. The SGD baseline at 300 and 500 epochs on STL-10 is the same, 75.25 ± 0.74 tab. 8. The paper gives no explanation; treat it as possibly a copy error.
• Integral typo. The integral in eq. 15 has dt rather than dt' inside the integral in the source. The intended meaning is clear from context.
• Theory assumptions. Results in §3 assume isotropic data and augmentation, EMA proportional to the online weights, and a symmetric predictor. The paper says the non-symmetric case is left to future work §3.1. DirectPred's predictor is symmetric by construction eq. 18.
• Normalised vs unnormalised loss. The objective in eq. 1 is unnormalised, but the experiments use ℓ2 normalisation §4. Reproduce the normalised version, which the code uses byol_trainer.py.
• Unstated ImageNet details. Learning rate, weight decay, EMA value and LARS settings for ImageNet are not stated §4. The paper says they follow BYOL, but BYOL's values are not given in this paper.

Lineage

• Builds on: BYOL (arXiv:2006.07733) and SimSiam (arXiv:2011.10566), whose dynamics the paper analyses. Neither has a card here; nothing from their content is used.
• Compared against: the default two-layer BYOL predictor tab. 9, a gradient-trained linear predictor tab. 6 tab. 9, and SGD-trained predictors on STL-10 and CIFAR-10 tab. 8. Contrastive methods are discussed as background in §1 only.
• Superseded by: not stated.