EMA Update kernel

[superposition]

Description

Exponential moving average update for target networks (momentum encoders).

Operation

# Update target params: target = tau * target + (1 - tau) * source
ema_update(target_params, source_params, tau=0.999)

Implementation

• Fused kernel for all parameters
• In-place update (no allocation)
• Support for parameter groups with different tau

Tests

• Correct EMA formula applied
• In-place modification of target
• Works with tau=0 (copy) and tau=1 (no-op)
• Handles multiple parameter tensors
• No gradients through target (detached)

Used In

JEPA, BYOL, MoCo, EMA model averaging, Polyak averaging