ema.py

"""EMA target encoder helpers.

The target encoder ``θ̄`` is a deep copy of the (trainable) Decompositional
Encoder ``θ``. It receives no gradient and is updated in place after every
optimizer step:

    θ̄ ← α · θ̄ + (1 - α) · θ

It supplies the (stable) prediction targets for the temporal objective
``L_temp`` and is the encoder used at inference (see the paper, §3.3).
"""

from __future__ import annotations

import copy

import torch
from torch import nn


def make_ema_copy(module: nn.Module) -> nn.Module:
    """Deep-copy ``module`` for EMA use: ``requires_grad=False``, eval, fp32."""
    ema = copy.deepcopy(module)
    for p in ema.parameters():
        p.requires_grad = False
        p.data = p.data.float()
    ema.eval()
    return ema


@torch.no_grad()
def update_ema(ema_module: nn.Module, online_module: nn.Module, decay: float) -> None:
    """In place ``θ̄ ← decay·θ̄ + (1-decay)·θ`` over matching parameters.

    Online params may be bf16/fp16/fp32 (mixed-precision keeps fp32 masters);
    EMA params stay fp32 for numerical stability across many steps. Buffers
    (e.g. sinusoidal positional embeddings) are not trained and not updated.
    """
    for p_ema, p in zip(ema_module.parameters(), online_module.parameters()):
        p_ema.data.mul_(decay).add_(p.data.float(), alpha=1.0 - decay)


@torch.no_grad()
def sync_ema_from_online(ema_module: nn.Module, online_module: nn.Module) -> None:
    """Copy online → EMA (fp32). Used to warm-start the target encoder when a
    checkpoint lacks EMA weights."""
    for p_ema, p in zip(ema_module.parameters(), online_module.parameters()):
        p_ema.data.copy_(p.data.float())