Understanding how EMA works?

Understanding how EMA works with an simulation of a toy model with toy weights

Exponential Moving Average (EMA) in Weight Updates

EMA (Exponential Moving Average) is an incredibly useful concept that finds application in various scenarios:

• Weight Updates: EMA is used for updating model weights while retaining a historical record of previous weights. This enables the model to blend new information with past knowledge effectively.

• Self-Supervised Learning: EMA is commonly employed in Self-Supervised Learning setups. The weights obtained from Self-Supervised Learning are often utilized for downstream tasks like classification and segmentation.

Clarification on EMA’s Impact

Initially, there was a misconception that the EMA process directly impacts the ongoing training of model weights. However, this is not the case. In reality, the EMA process involves the creation of a duplicated set of weights. These duplicate weights are updated alongside the primary training process, and the updated weights are subsequently leveraged for validation purposes. As a result, the overall training procedure remains unaffected by the EMA process.

image.png

::: {#36edc225 .cell _cell_guid=‘b1076dfc-b9ad-4769-8c92-a6c4dae69d19’ _uuid=‘8f2839f25d086af736a60e9eeb907d3b93b6e0e5’ execution=‘{“iopub.execute_input”:“2023-08-14T15:14:27.136587Z”,“iopub.status.busy”:“2023-08-14T15:14:27.136194Z”,“iopub.status.idle”:“2023-08-14T15:14:30.874853Z”,“shell.execute_reply”:“2023-08-14T15:14:30.873447Z”}’ papermill=‘{“duration”:3.746175,“end_time”:“2023-08-14T15:14:30.877588”,“exception”:false,“start_time”:“2023-08-14T15:14:27.131413”,“status”:“completed”}’ tags=‘[]’ execution_count=1}

import torch.nn as nn
import torch
from copy import deepcopy

class Model(nn.Module):
    def __init__(self,):
        super().__init__()

        self.layer = nn.Linear(3,3,bias=False)

    def forward(self,x):
        return self.layer(x)
    
@torch.no_grad()
def init_weights_1(m):
    print(m)
    if type(m) == nn.Linear:
        m.weight.fill_(1.0)
        print(m.weight)

@torch.no_grad()
def init_weights_2(m):
    print(m)
    if type(m) == nn.Linear:
        m.weight.fill_(2.0)
        print(m.weight)

:::

teacher = Model().apply(init_weights_1) ## TO BE UPDATED
student = Model().apply(init_weights_2) ## From Training

Linear(in_features=3, out_features=3, bias=False)
Parameter containing:
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]], requires_grad=True)
Model(
  (layer): Linear(in_features=3, out_features=3, bias=False)
)
Linear(in_features=3, out_features=3, bias=False)
Parameter containing:
tensor([[2., 2., 2.],
        [2., 2., 2.],
        [2., 2., 2.]], requires_grad=True)
Model(
  (layer): Linear(in_features=3, out_features=3, bias=False)
)

@torch.no_grad()
def ema_model(teacher_model, student_model,decay=0.5):
    teacher_model.eval()

    for teacher_wt,student_wt in zip(teacher_model.state_dict().values(),student_model.state_dict().values()):
        teacher_wt.copy_(teacher_wt*decay + student_wt*(1-decay))

ema_model(teacher,student)

teacher.state_dict()

OrderedDict([('layer.weight',
              tensor([[1.5000, 1.5000, 1.5000],
                      [1.5000, 1.5000, 1.5000],
                      [1.5000, 1.5000, 1.5000]]))])

student.state_dict()

OrderedDict([('layer.weight',
              tensor([[2., 2., 2.],
                      [2., 2., 2.],
                      [2., 2., 2.]]))])