EMA Target Encoder — Visual Experiments¶

Goal: See how a moving-average teacher filters student updates, how momentum controls memory, and why DINO needs more than EMA alone.

What we'll explore¶

#	Experiment	Question answered
1	Noise filtering	How does EMA turn a noisy student into a smoother target?
2	Teacher memory	How much history does each momentum value retain?
3	Stability vs lag	What happens when the representation changes suddenly?
4	DINO stabilization	Why are centering and sharpening needed around the EMA teacher?

---

Linked theory: research/05-ema-target-encoder.md

In [1]:

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)
plt.rcParams.update({
    'figure.dpi': 110,
    'font.size': 11,
    'axes.titlesize': 12,
    'axes.labelsize': 11,
    'figure.facecolor': 'white',
    'axes.facecolor': '#f8f9fa',
    'axes.grid': True,
    'grid.alpha': 0.35,
})

def ema(values, momentum):
    values = np.asarray(values, dtype=float)
    result = np.empty_like(values)
    result[0] = values[0]
    for index in range(1, len(values)):
        result[index] = momentum * result[index - 1] + (1 - momentum) * values[index]
    return result

def softmax(logits, temperature=1.0):
    scaled = np.asarray(logits) / temperature
    shifted = scaled - np.max(scaled, axis=-1, keepdims=True)
    probabilities = np.exp(shifted)
    return probabilities / probabilities.sum(axis=-1, keepdims=True)

import numpy as np import matplotlib.pyplot as plt np.random.seed(42) plt.rcParams.update({ 'figure.dpi': 110, 'font.size': 11, 'axes.titlesize': 12, 'axes.labelsize': 11, 'figure.facecolor': 'white', 'axes.facecolor': '#f8f9fa', 'axes.grid': True, 'grid.alpha': 0.35, }) def ema(values, momentum): values = np.asarray(values, dtype=float) result = np.empty_like(values) result[0] = values[0] for index in range(1, len(values)): result[index] = momentum * result[index - 1] + (1 - momentum) * values[index] return result def softmax(logits, temperature=1.0): scaled = np.asarray(logits) / temperature shifted = scaled - np.max(scaled, axis=-1, keepdims=True) probabilities = np.exp(shifted) return probabilities / probabilities.sum(axis=-1, keepdims=True)

---

Experiment 1: EMA filters noisy student updates¶

Question: How does an EMA teacher behave when the online parameter follows a useful trend plus minibatch noise?

We model one parameter coordinate. The underlying trajectory changes smoothly, while the observed student coordinate receives independent optimization noise at every step.

In [2]:

np.random.seed(1)
steps_exp1 = np.arange(300)
underlying_exp1 = 0.004 * steps_exp1 + 0.45 * np.sin(steps_exp1 / 34)
student_exp1 = underlying_exp1 + np.random.normal(0, 0.23, size=len(steps_exp1))
momenta_exp1 = [0.0, 0.90, 0.99]
teachers_exp1 = {momentum: ema(student_exp1, momentum) for momentum in momenta_exp1}
rmse_exp1 = {
    momentum: np.sqrt(np.mean((trajectory - underlying_exp1) ** 2))
    for momentum, trajectory in teachers_exp1.items()
}

np.random.seed(1) steps_exp1 = np.arange(300) underlying_exp1 = 0.004 * steps_exp1 + 0.45 * np.sin(steps_exp1 / 34) student_exp1 = underlying_exp1 + np.random.normal(0, 0.23, size=len(steps_exp1)) momenta_exp1 = [0.0, 0.90, 0.99] teachers_exp1 = {momentum: ema(student_exp1, momentum) for momentum in momenta_exp1} rmse_exp1 = { momentum: np.sqrt(np.mean((trajectory - underlying_exp1) ** 2)) for momentum, trajectory in teachers_exp1.items() }

In [3]:

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2))
fig.suptitle('EMA teacher as a low-pass filter over student parameters', fontsize=14, fontweight='bold')

ax = axes[0]
ax.plot(steps_exp1, student_exp1, color='#adb5bd', alpha=0.55, linewidth=1, label='Noisy student')
ax.plot(steps_exp1, underlying_exp1, color='#212529', linewidth=2.5, label='Underlying trajectory')
ax.plot(steps_exp1, teachers_exp1[0.90], color='#457b9d', linewidth=2, label='EMA m=0.90')
ax.plot(steps_exp1, teachers_exp1[0.99], color='#e76f51', linewidth=2, label='EMA m=0.99')
ax.set_xlabel('Optimization step')
ax.set_ylabel('Parameter coordinate')
ax.set_title('Teacher trajectories')
ax.legend()

ax = axes[1]
labels = ['m=0.00', 'm=0.90', 'm=0.99']
values = [rmse_exp1[momentum] for momentum in momenta_exp1]
ax.bar(labels, values, color=['#adb5bd', '#457b9d', '#e76f51'])
ax.set_xlabel('Teacher momentum')
ax.set_ylabel('RMSE to underlying trajectory')
ax.set_title('Filtering improves target fidelity')

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2)) fig.suptitle('EMA teacher as a low-pass filter over student parameters', fontsize=14, fontweight='bold') ax = axes[0] ax.plot(steps_exp1, student_exp1, color='#adb5bd', alpha=0.55, linewidth=1, label='Noisy student') ax.plot(steps_exp1, underlying_exp1, color='#212529', linewidth=2.5, label='Underlying trajectory') ax.plot(steps_exp1, teachers_exp1[0.90], color='#457b9d', linewidth=2, label='EMA m=0.90') ax.plot(steps_exp1, teachers_exp1[0.99], color='#e76f51', linewidth=2, label='EMA m=0.99') ax.set_xlabel('Optimization step') ax.set_ylabel('Parameter coordinate') ax.set_title('Teacher trajectories') ax.legend() ax = axes[1] labels = ['m=0.00', 'm=0.90', 'm=0.99'] values = [rmse_exp1[momentum] for momentum in momenta_exp1] ax.bar(labels, values, color=['#adb5bd', '#457b9d', '#e76f51']) ax.set_xlabel('Teacher momentum') ax.set_ylabel('RMSE to underlying trajectory') ax.set_title('Filtering improves target fidelity') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: Copying the student ($m=0$) preserves every noisy jump. Both EMA teachers are smoother, and $m=0.99$ removes the most high-frequency variation.
2. Why it looks this way: The teacher averages many historical students, so independent minibatch fluctuations cancel while persistent trends accumulate.
3. Key takeaway: EMA improves target stability by temporal ensembling, but stronger smoothing also implies more delay when the trend changes.

---

Experiment 2: Momentum determines teacher memory¶

Question: How much influence does a student checkpoint retain after many updates?

For constant momentum $m$, the contribution of a checkpoint that is $\Delta$ steps old is $(1-m)m^\Delta$. The cumulative curve shows how many recent checkpoints account for most of the teacher.

In [4]:

np.random.seed(2)
momenta_exp2 = [0.90, 0.99, 0.999]
ages_exp2 = np.arange(0, 5001)
weights_exp2 = {
    momentum: (1 - momentum) * momentum ** ages_exp2
    for momentum in momenta_exp2
}
cumulative_exp2 = {
    momentum: np.cumsum(weights)
    for momentum, weights in weights_exp2.items()
}
half_lives_exp2 = {
    momentum: np.log(0.5) / np.log(momentum)
    for momentum in momenta_exp2
}

np.random.seed(2) momenta_exp2 = [0.90, 0.99, 0.999] ages_exp2 = np.arange(0, 5001) weights_exp2 = { momentum: (1 - momentum) * momentum ** ages_exp2 for momentum in momenta_exp2 } cumulative_exp2 = { momentum: np.cumsum(weights) for momentum, weights in weights_exp2.items() } half_lives_exp2 = { momentum: np.log(0.5) / np.log(momentum) for momentum in momenta_exp2 }

In [5]:

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2))
fig.suptitle('EMA momentum defines an exponential memory over checkpoints', fontsize=14, fontweight='bold')
colors = ['#2a9d8f', '#457b9d', '#e76f51']

ax = axes[0]
for momentum, color in zip(momenta_exp2, colors):
    ax.semilogy(ages_exp2, weights_exp2[momentum], color=color, linewidth=2,
                label=f'm={momentum}, half-life={half_lives_exp2[momentum]:.0f}')
ax.set_xlim(0, 1500)
ax.set_ylim(1e-7, 2e-1)
ax.set_xlabel('Checkpoint age in steps')
ax.set_ylabel('Weight in current teacher (log scale)')
ax.set_title('Influence decays exponentially')
ax.legend()

ax = axes[1]
for momentum, color in zip(momenta_exp2, colors):
    ax.plot(ages_exp2, cumulative_exp2[momentum], color=color, linewidth=2, label=f'm={momentum}')
ax.axhline(0.95, color='#6c757d', linestyle='--', label='95% of total weight')
ax.set_xlim(0, 3000)
ax.set_ylim(0, 1.03)
ax.set_xlabel('Number of recent checkpoints included')
ax.set_ylabel('Cumulative teacher weight')
ax.set_title('Higher momentum averages a much longer history')
ax.legend()

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2)) fig.suptitle('EMA momentum defines an exponential memory over checkpoints', fontsize=14, fontweight='bold') colors = ['#2a9d8f', '#457b9d', '#e76f51'] ax = axes[0] for momentum, color in zip(momenta_exp2, colors): ax.semilogy(ages_exp2, weights_exp2[momentum], color=color, linewidth=2, label=f'm={momentum}, half-life={half_lives_exp2[momentum]:.0f}') ax.set_xlim(0, 1500) ax.set_ylim(1e-7, 2e-1) ax.set_xlabel('Checkpoint age in steps') ax.set_ylabel('Weight in current teacher (log scale)') ax.set_title('Influence decays exponentially') ax.legend() ax = axes[1] for momentum, color in zip(momenta_exp2, colors): ax.plot(ages_exp2, cumulative_exp2[momentum], color=color, linewidth=2, label=f'm={momentum}') ax.axhline(0.95, color='#6c757d', linestyle='--', label='95% of total weight') ax.set_xlim(0, 3000) ax.set_ylim(0, 1.03) ax.set_xlabel('Number of recent checkpoints included') ax.set_ylabel('Cumulative teacher weight') ax.set_title('Higher momentum averages a much longer history') ax.legend() plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: At $m=0.90$, old checkpoints vanish within tens of steps. At $m=0.999$, hundreds of steps still have visible influence and roughly three thousand recent checkpoints are needed for 95% of the mass.
2. Why it looks this way: Each additional step multiplies historical weight by $m$. A seemingly small change from 0.99 to 0.999 increases the effective memory by about ten times.
3. Key takeaway: Momentum is a time-scale parameter, so it must be interpreted relative to the number of optimizer steps.

---

Experiment 3: Stability and adaptation pull in opposite directions¶

Question: How quickly can a smooth teacher follow a sudden representation shift?

The synthetic student changes regime halfway through training. We measure target noise before the shift and the number of steps needed to reach 90% of the new level.

In [6]:

np.random.seed(3)
steps_exp3 = np.arange(400)
change_step = 180
true_level_exp3 = np.where(steps_exp3 < change_step, 0.0, 1.0)
student_exp3 = true_level_exp3 + np.random.normal(0, 0.18, size=len(steps_exp3))
momenta_exp3 = [0.50, 0.90, 0.99]
teachers_exp3 = {momentum: ema(student_exp3, momentum) for momentum in momenta_exp3}
noise_exp3 = {}
delay_exp3 = {}

for momentum, trajectory in teachers_exp3.items():
    noise_exp3[momentum] = np.std(trajectory[40:change_step] - true_level_exp3[40:change_step])
    after_change = trajectory[change_step:]
    reached = np.flatnonzero(after_change >= 0.9)
    delay_exp3[momentum] = reached[0] if len(reached) else len(after_change)

np.random.seed(3) steps_exp3 = np.arange(400) change_step = 180 true_level_exp3 = np.where(steps_exp3 < change_step, 0.0, 1.0) student_exp3 = true_level_exp3 + np.random.normal(0, 0.18, size=len(steps_exp3)) momenta_exp3 = [0.50, 0.90, 0.99] teachers_exp3 = {momentum: ema(student_exp3, momentum) for momentum in momenta_exp3} noise_exp3 = {} delay_exp3 = {} for momentum, trajectory in teachers_exp3.items(): noise_exp3[momentum] = np.std(trajectory[40:change_step] - true_level_exp3[40:change_step]) after_change = trajectory[change_step:] reached = np.flatnonzero(after_change >= 0.9) delay_exp3[momentum] = reached[0] if len(reached) else len(after_change)

In [7]:

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2))
fig.suptitle('A smoother teacher pays with slower adaptation', fontsize=14, fontweight='bold')
colors = ['#2a9d8f', '#457b9d', '#e76f51']

ax = axes[0]
ax.plot(steps_exp3, student_exp3, color='#adb5bd', alpha=0.45, linewidth=1, label='Noisy student')
ax.plot(steps_exp3, true_level_exp3, color='#212529', linewidth=2.4, label='True regime')
for momentum, color in zip(momenta_exp3, colors):
    ax.plot(steps_exp3, teachers_exp3[momentum], color=color, linewidth=2, label=f'EMA m={momentum}')
ax.axvline(change_step, color='#6c757d', linestyle='--', label='Regime change')
ax.set_xlabel('Optimization step')
ax.set_ylabel('Parameter coordinate')
ax.set_title('Response to a sudden shift')
ax.legend(fontsize=8)

ax = axes[1]
for momentum, color in zip(momenta_exp3, colors):
    ax.scatter(noise_exp3[momentum], delay_exp3[momentum], s=110, color=color, label=f'm={momentum}')
    ax.annotate(f'{momentum}', (noise_exp3[momentum], delay_exp3[momentum]), xytext=(6, 5), textcoords='offset points')
ax.set_xlabel('Pre-shift teacher noise (lower is smoother)')
ax.set_ylabel('Steps to reach 90% of new level')
ax.set_title('Stability-lag trade-off')
ax.legend()

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2)) fig.suptitle('A smoother teacher pays with slower adaptation', fontsize=14, fontweight='bold') colors = ['#2a9d8f', '#457b9d', '#e76f51'] ax = axes[0] ax.plot(steps_exp3, student_exp3, color='#adb5bd', alpha=0.45, linewidth=1, label='Noisy student') ax.plot(steps_exp3, true_level_exp3, color='#212529', linewidth=2.4, label='True regime') for momentum, color in zip(momenta_exp3, colors): ax.plot(steps_exp3, teachers_exp3[momentum], color=color, linewidth=2, label=f'EMA m={momentum}') ax.axvline(change_step, color='#6c757d', linestyle='--', label='Regime change') ax.set_xlabel('Optimization step') ax.set_ylabel('Parameter coordinate') ax.set_title('Response to a sudden shift') ax.legend(fontsize=8) ax = axes[1] for momentum, color in zip(momenta_exp3, colors): ax.scatter(noise_exp3[momentum], delay_exp3[momentum], s=110, color=color, label=f'm={momentum}') ax.annotate(f'{momentum}', (noise_exp3[momentum], delay_exp3[momentum]), xytext=(6, 5), textcoords='offset points') ax.set_xlabel('Pre-shift teacher noise (lower is smoother)') ax.set_ylabel('Steps to reach 90% of new level') ax.set_title('Stability-lag trade-off') ax.legend() plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: High momentum produces the cleanest pre-shift target, but it takes the longest to approach the new regime. Low momentum adapts quickly and transmits more noise.
2. Why it looks this way: The same long memory that cancels fluctuations also preserves obsolete student checkpoints after a real change.
3. Key takeaway: EMA cannot maximize stability and responsiveness simultaneously; schedules that increase momentum over training are a practical compromise.

---

Experiment 4: DINO needs centering and sharpening around EMA¶

Question: Why does a stable teacher still need output-level collapse controls?

We create teacher logits with a global bias toward output dimension 0 plus a sample-specific semantic signal. Centering removes the global bias; sharpening lowers each sample's entropy so its target remains informative.

In [8]:

np.random.seed(4)
batch_size_exp4 = 600
dimensions_exp4 = 6
semantic_class_exp4 = np.arange(batch_size_exp4) % dimensions_exp4
teacher_logits_exp4 = np.random.normal(0, 0.35, size=(batch_size_exp4, dimensions_exp4))
teacher_logits_exp4[np.arange(batch_size_exp4), semantic_class_exp4] += 1.8
teacher_logits_exp4[:, 0] += 2.4
center_exp4 = teacher_logits_exp4.mean(axis=0, keepdims=True)

uncentered_probs_exp4 = softmax(teacher_logits_exp4, temperature=0.20)
centered_soft_exp4 = softmax(teacher_logits_exp4 - center_exp4, temperature=1.00)
centered_sharp_exp4 = softmax(teacher_logits_exp4 - center_exp4, temperature=0.20)
usage_exp4 = np.vstack([
    uncentered_probs_exp4.mean(axis=0),
    centered_soft_exp4.mean(axis=0),
    centered_sharp_exp4.mean(axis=0),
])
sample_index_exp4 = 4
sample_distributions_exp4 = np.vstack([
    centered_soft_exp4[sample_index_exp4],
    centered_sharp_exp4[sample_index_exp4],
])

np.random.seed(4) batch_size_exp4 = 600 dimensions_exp4 = 6 semantic_class_exp4 = np.arange(batch_size_exp4) % dimensions_exp4 teacher_logits_exp4 = np.random.normal(0, 0.35, size=(batch_size_exp4, dimensions_exp4)) teacher_logits_exp4[np.arange(batch_size_exp4), semantic_class_exp4] += 1.8 teacher_logits_exp4[:, 0] += 2.4 center_exp4 = teacher_logits_exp4.mean(axis=0, keepdims=True) uncentered_probs_exp4 = softmax(teacher_logits_exp4, temperature=0.20) centered_soft_exp4 = softmax(teacher_logits_exp4 - center_exp4, temperature=1.00) centered_sharp_exp4 = softmax(teacher_logits_exp4 - center_exp4, temperature=0.20) usage_exp4 = np.vstack([ uncentered_probs_exp4.mean(axis=0), centered_soft_exp4.mean(axis=0), centered_sharp_exp4.mean(axis=0), ]) sample_index_exp4 = 4 sample_distributions_exp4 = np.vstack([ centered_soft_exp4[sample_index_exp4], centered_sharp_exp4[sample_index_exp4], ])

In [9]:

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2))
fig.suptitle('EMA stabilizes weights; centering and sharpening stabilize targets', fontsize=14, fontweight='bold')
dimensions = np.arange(dimensions_exp4)
width = 0.25

ax = axes[0]
ax.bar(dimensions - width, usage_exp4[0], width, color='#e76f51', label='No center, tau=0.20')
ax.bar(dimensions, usage_exp4[1], width, color='#adb5bd', label='Centered, tau=1.00')
ax.bar(dimensions + width, usage_exp4[2], width, color='#457b9d', label='Centered, tau=0.20')
ax.axhline(1 / dimensions_exp4, color='#212529', linestyle='--', label='Uniform average usage')
ax.set_xlabel('Teacher output dimension')
ax.set_ylabel('Mean probability across batch')
ax.set_title('Centering prevents one dimension dominating')
ax.set_xticks(dimensions)
ax.legend(fontsize=8)

ax = axes[1]
ax.bar(dimensions - width / 2, sample_distributions_exp4[0], width, color='#adb5bd', label='Centered, tau=1.00')
ax.bar(dimensions + width / 2, sample_distributions_exp4[1], width, color='#2a9d8f', label='Centered, tau=0.20')
ax.set_xlabel('Teacher output dimension')
ax.set_ylabel('Probability for one sample')
ax.set_title('Sharpening makes each target informative')
ax.set_xticks(dimensions)
ax.legend()

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2)) fig.suptitle('EMA stabilizes weights; centering and sharpening stabilize targets', fontsize=14, fontweight='bold') dimensions = np.arange(dimensions_exp4) width = 0.25 ax = axes[0] ax.bar(dimensions - width, usage_exp4[0], width, color='#e76f51', label='No center, tau=0.20') ax.bar(dimensions, usage_exp4[1], width, color='#adb5bd', label='Centered, tau=1.00') ax.bar(dimensions + width, usage_exp4[2], width, color='#457b9d', label='Centered, tau=0.20') ax.axhline(1 / dimensions_exp4, color='#212529', linestyle='--', label='Uniform average usage') ax.set_xlabel('Teacher output dimension') ax.set_ylabel('Mean probability across batch') ax.set_title('Centering prevents one dimension dominating') ax.set_xticks(dimensions) ax.legend(fontsize=8) ax = axes[1] ax.bar(dimensions - width / 2, sample_distributions_exp4[0], width, color='#adb5bd', label='Centered, tau=1.00') ax.bar(dimensions + width / 2, sample_distributions_exp4[1], width, color='#2a9d8f', label='Centered, tau=0.20') ax.set_xlabel('Teacher output dimension') ax.set_ylabel('Probability for one sample') ax.set_title('Sharpening makes each target informative') ax.set_xticks(dimensions) ax.legend() plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: Without centering, output dimension 0 absorbs most probability across the batch. Centering restores balanced average usage, while a lower teacher temperature turns a diffuse per-sample target into a confident one.
2. Why it looks this way: EMA smooths parameter updates but does not constrain output occupancy or entropy. Centering and sharpening act directly on those two collapse modes.
3. Key takeaway: DINO is stable because EMA, stop-gradient, centering, and sharpening cooperate; an EMA teacher alone is insufficient.

---

Summary / Key Takeaways¶

Experiment	Key insight
1. Noise filtering	EMA creates a smoother temporal ensemble of student parameters.
2. Teacher memory	Moving from 0.99 to 0.999 increases memory by roughly an order of magnitude.
3. Stability vs lag	Smooth targets adapt slowly after genuine representation shifts.
4. DINO stabilization	Weight averaging does not replace output-level anti-collapse controls.

Connection to the broader theory¶

The EMA target encoder supplies a slowly moving coordinate system for latent prediction. JEPA uses that coordinate system as the target space in which a context representation must predict masked content.