Model EMA는
원본 model의 weight를 평균내어 저장하는
shadow model 을 사용함.
Model EMA(Exponential Moving Average, 지수 이동 평균)
- 학습 중 계속 변화하는 원본 model의 weight(가중치)를
- 지수 가중 방식으로 평균내어 저장하는 기법임.
일반적인 Model EMA에서는 다음 두 model을 함께 유지함.
- 원본 model: gradient와 optimizer를 통해 실제로 학습되는 model
- EMA model: 원본 model의 weight를 받아 평균값으로 갱신되는 model
핵심은 다음과 같음.
Gradient는 원본 model에서만 계산되며,
EMA model은
원본 model의 weight를 평균내어 갱신만 수행함.
원본 model의 학습 과정에서 생성된 weight들을 부드럽게 평균내어 저장하는 shadow model(그림자 모델) 에 해당함.
1. EMA의 배경
EMA는 딥러닝에서 처음 등장한 개념이 아님.
- 시계열 분석과 신호 처리 분야에서
- 새로운 관측값에는 상대적으로 큰 비중을 부여하고,
- 오래된 관측값의 영향은 점차 감소시키기 위해 사용되어 온 이동 평균 방식임.
기존의 parameter에 더 큰 비중을 부여하는 EMA와는 세부적인 계산 방식에서 차이가 있음.
- 딥러닝에서 model weight의 EMA를 명확한 구조로 사용한 대표적인 연구는
- Tarvainen과 Valpola가 2017년에 제안한 Mean Teacher임.
Mean Teacher는 다음 두 model로 구성됨.
- Student model: gradient descent를 통해 학습되는 model
- Teacher model: Student model weight의 EMA로 갱신되는 model
Teacher라는 이름을 사용하지만, Teacher model이 별도의 optimizer를 통해 학습되는 것은 아님.
Teacher model은 Student model의 weight를 지수 이동 평균하여 구성한 평균 model임.
2. 원본 model의 weight가 흔들리는 이유
신경망은 일반적으로 전체 학습 데이터를 한 번에 사용하지 않고 mini-batch 단위로 학습함.
- 각 mini-batch에 포함되는 sample이 달라지므로, 매 학습 step에서 계산되는 gradient에도 차이가 발생함.
- 원본 model의 weight를 $\boldsymbol{\theta}$라고 하면 SGD의 weight update는 다음과 같이 표현됨.
$$
\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \eta \mathbf{g}_t
$$
각 기호의 의미는 다음과 같음.
- $\boldsymbol{\theta}_t$: $t$번째 step에서의 원본 model weight
- $\mathbf{g}_t$: 현재 mini-batch에서 계산된 gradient
- $\eta$: learning rate
Mini-batch가 달라질 때마다 $\mathbf{g}_t$도 달라짐.
따라서 model weight는
- loss가 감소하는 방향으로 이동하지만,
- 동시에 mini-batch sampling으로 인한 noise에도 영향을 받음.
예를 들어 하나의 weight가 다음과 같이 변화할 수 있음.
초기값: 1.00
Step 1: 1.20
Step 2: 0.95
Step 3: 1.15
Step 4: 1.05
문제는
- 마지막 step의 값인
1.05가 반드시 가장 우수한 generalization performance(일반화 성능)를 보이는 값이라고 보장할 수 없음. - 마지막 mini-batch의 영향이 우연히 크게 반영된 값일 수도 있기 때문임.
EMA는 마지막 weight 하나만 사용하는 대신, 최근까지의 weight를 부드럽게 평균내어 이러한 변동을 완화하는 방법임.
3. 원본 model과 EMA model의 역할
학습 시작 시 원본 model을 복사하여 EMA model을 생성함.
model # 실제로 학습되는 원본 model
ema_model # 원본 model의 weight를 평균내는 EMA model3-1. 원본 model
원본 model은 일반적인 neural network 학습 과정을 수행함.
- 입력 data를 이용한 forward 수행
- 예측값과 정답을 이용한 loss 계산
backward()를 통한 gradient 계산optimizer.step()을 통한 weight update
즉, 원본 model의 parameter는 loss로부터 전달된 gradient에 의해 직접 변경됨.
3-2. EMA model
EMA model은 일반적으로 다음 과정을 수행하지 않음.
- 별도의 loss 계산 없음
- 별도의 backward 수행 없음
- 별도의 gradient 계산 없음
- 별도의 optimizer 사용 없음
EMA model의 갱신 과정은 다음과 같음.
- 원본 model의 최신 weight 확인
- 기존 EMA weight와 원본 model의 최신 weight 결합
- 계산된 값으로 EMA weight 갱신
원본 model이 mini-batch마다 빠르게 변하는 동안,
EMA model은 원본 model의 변화를 더 천천히 추적함.
이러한 특성 때문에 EMA model을 shadow model이라고 부르기도 함.
4. EMA update 공식
원본 model의 weight를 $\boldsymbol{\theta}$, EMA model의 weight를 $\bar{\boldsymbol{\theta}}$라고 하면 EMA update는 다음과 같음.
$$
\bar{\boldsymbol{\theta}}_{t+1} = \beta \bar{\boldsymbol{\theta}}_t + (1-\beta)\boldsymbol{\theta}_{t+1}
$$
각 기호의 의미는 다음과 같음.
- $\bar{\boldsymbol{\theta}}_t$: 기존 EMA model weight
- $\boldsymbol{\theta}_{t+1}$: optimizer update가 완료된 원본 model의 최신 weight
- $\beta$: EMA decay
예를 들어
$\beta=0.9$인 경우 다음과 같이 계산됨.
새 EMA weight
= 기존 EMA weight × 0.9
+ 현재 model weight × 0.1$\beta=0.999$인 경우 다음과 같음.
새 EMA weight
= 기존 EMA weight × 0.999
+ 현재 model weight × 0.001$\beta$가 클수록 기존 EMA weight의 비중이 커짐.
따라서 EMA model이 원본 model의 최신 변화를 더 천천히 반영함.
- 작은 decay: 원본 model의 최신 변화가 빠르게 반영됨
- 큰 decay: weight 변화가 강하게 평활화됨
5. Gradient가 계산되는 위치
학습 data는 EMA model이 아니라 원본 model에 입력됨.
prediction = model(x)
loss = criterion(prediction, y)이후 loss.backward()를 실행하면 원본 model의 parameter에 gradient가 계산됨.
loss.backward()계산된 gradient는 다음과 같이 원본 model의 parameter에 저장됨.
model.weight.grad
model.bias.gradOptimizer는 이 gradient를 사용하여 원본 model의 weight를 갱신함.
optimizer.step()EMA model은 원본 model의 update가 끝난 이후에 갱신됨.
update_ema(ema_model, model)
전체 학습 순서는 다음과 같음.
- 입력 data를 원본 model에 전달하여 forward 수행
- 원본 model의 예측값과 정답을 이용하여 loss 계산
loss.backward()를 실행하여 원본 model의 gradient 계산optimizer.step()을 실행하여 원본 model의 weight 갱신- 갱신된 원본 model의 최신 weight를 이용하여 EMA model 갱신
따라서 gradient가 EMA model에 직접 전달되는 것은 아님.
- Gradient는 먼저 원본 model의 weight를 변경함.
- EMA model은 변경된 원본 model의 weight를 이용하여 간접적으로 갱신됨.
6. 수치 예시
하나의 weight만 존재하는 model을 가정함.
원본 model weight = 1.00
EMA model weight = 1.00
learning rate = 0.10
EMA decay = 0.90
원본 model에서 계산된 gradient가 -2.0이라고 가정함.
1단계: 원본 model update
SGD update는 다음과 같음.
$$
\theta_{\text{new}} = \theta_{\text{old}}-\eta g
$$
따라서 원본 model의 새로운 weight는 다음과 같음.
새 model weight
= 1.00 - 0.10 × (-2.0)
= 1.20원본 model의 weight가 1.00에서 1.20으로 변경됨.
2단계: EMA model update
EMA model에는 gradient를 적용하지 않음.
Optimizer update가 끝난 원본 model의 weight 1.20을 이용하여 다음과 같이 갱신함.
새 EMA weight
= 0.90 × 1.00 + 0.10 × 1.20
= 0.90 + 0.12
= 1.02최종 결과는 다음과 같음.
원본 model weight = 1.20
EMA model weight = 1.02- 원본 model은 gradient에 의해
0.20만큼 변경됨. - 반면 EMA model은 원본 model의 최신 weight 중 일부만 반영하므로
0.02만큼 변경됨.
따라서
EMA model은
원본 model보다 천천히 이동하며,
학습 step별 weight 변동을 완화함.
7. EMA model은 별도로 학습되는 model이 아님
일반적인 model 학습은 forward, loss 계산, backward, gradient 계산, optimizer update의 순서로 수행됨.
EMA model은 이러한 학습 과정을 수행하지 않음.
일반적인 Model EMA에서는 다음 코드가 존재하지 않음.
ema_loss.backward()
ema_optimizer.step()
EMA model에는 다음 update만 수행됨.
update_ema(ema_model, model)
두 model의 역할은 다음과 같이 구분됨.
원본 model
- 학습 data를 입력받음
- loss를 계산함
- gradient를 계산함
- optimizer를 이용하여 weight를 갱신함
EMA model
- 별도의 학습 data를 이용한 loss 계산 없음
- 별도의 gradient 계산 없음
- 별도의 optimizer 사용 없음
- 원본 model의 weight를 평균내어 저장함
따라서 EMA model을 두 번째 학습 model이라고 표현하는 것은 정확하지 않음.
원본 model의 parameter trajectory(파라미터 궤적)를 평균내는 shadow model로 이해하는 것이 적절함.
8. Adam의 EMA와 Model EMA의 차이
Adam optimizer도 EMA를 사용함.
다만 Adam의 EMA와 Model EMA는 평균을 계산하는 대상이 다름.
Adam은 optimizer 내부에서 다음 두 통계를 관리함.
- gradient의 1차 moment(모멘트) EMA
- gradient 제곱의 2차 moment EMA
이를 단순화하면 다음과 같이 표현됨.
$$
\mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1-\beta_1)\mathbf{g}_t \\
\mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1-\beta_2)\mathbf{g}_t^2
$$
각 기호의 의미는 다음과 같음.
- $\mathbf{m}_t$: gradient의 1차 moment 추정값
- $\mathbf{v}_t$: gradient 제곱의 2차 moment 추정값
- $\mathbf{g}_t$: 현재 step의 gradient
Adam은 이러한 gradient 통계를 사용하여 원본 model의 weight를 갱신함.
반면 Model EMA는 optimizer update가 끝난 원본 model의 weight 자체를 평균냄.
두 EMA의 차이는 다음과 같음.
- Adam 내부 EMA: gradient와 gradient 제곱을 평균내어 원본 model의 weight update에 사용함
- Model EMA: 원본 model의 weight를 평균내어 EMA model의 weight를 갱신함
따라서 Adam과 Model EMA는 두 기법을 동시에 사용할 수 있음.
optimizer = torch.optim.Adam(
model.parameters(),
lr=1e-3,
)
loss.backward()
optimizer.step()
update_ema(ema_model, model)optimizer.step()에서는 Adam의 gradient EMA가 사용됨.update_ema()에서는 원본 model weight의 EMA가 계산됨.
9. PyTorch 직접 구현
다음은 가장 단순한 Model EMA 구현임.
import copy
import torch
import torch.nn as nn
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 1)
def forward(self, x):
return self.linear(x)
model = SimpleModel()
# 원본 model을 복사하여 EMA model 생성
ema_model = copy.deepcopy(model)
# EMA model의 gradient 계산 비활성화
ema_model.requires_grad_(False)
EMA model은 optimizer에 포함하지 않음.
optimizer = torch.optim.Adam(
model.parameters(),
lr=1e-3,
)
EMA update 함수는 다음과 같이 구현할 수 있음.
@torch.no_grad()
def update_ema(
ema_model: nn.Module,
model: nn.Module,
decay: float = 0.999,
) -> None:
for ema_param, model_param in zip(
ema_model.parameters(),
model.parameters(),
):
ema_param.mul_(decay)
ema_param.add_(
model_param,
alpha=1.0 - decay,
)
다음 코드는 EMA 공식을 구현한 것으로 다음과 같음:
# ema_param.mul_(decay)
# ema_param.add_(model_param, alpha=1.0 - decay)
ema_param = (
decay * ema_param
+ (1.0 - decay) * model_param
)@torch.no_grad()를 사용하는 이유는 EMA update가 loss와 gradient를 계산하기 위한 연산이 아니기 때문임.- EMA update는 parameter 값을 직접 수정하는 작업이므로 autograd graph를 생성할 필요가 없음.
10. 전체 학습 과정
다음은 원본 model의 학습과 EMA model의 update를 함께 수행하는 예시임.
import copy
import torch
import torch.nn as nn
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 1)
def forward(self, x):
return self.linear(x)
@torch.no_grad()
def update_ema(
ema_model: nn.Module,
model: nn.Module,
decay: float = 0.999,
) -> None:
for ema_param, model_param in zip(
ema_model.parameters(),
model.parameters(),
):
ema_param.mul_(decay)
ema_param.add_(
model_param,
alpha=1.0 - decay,
)
device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)
# 실제로 학습되는 원본 model
model = SimpleModel().to(device)
# 원본 model의 weight를 평균내는 EMA model
ema_model = copy.deepcopy(model).to(device)
ema_model.requires_grad_(False)
optimizer = torch.optim.Adam(
model.parameters(),
lr=1e-3,
)
criterion = nn.MSELoss()
for step in range(1000):
x = torch.randn(32, 10, device=device)
y = torch.randn(32, 1, device=device)
model.train()
optimizer.zero_grad()
prediction = model(x)
loss = criterion(prediction, y)
# 원본 model의 gradient 계산
loss.backward()
# 원본 model의 weight update
optimizer.step()
# EMA model의 weight update
update_ema(
ema_model=ema_model,
model=model,
decay=0.999,
)
가장 중요한 부분은 다음과 같음:
loss.backward() # 원본 model의 gradient 계산
optimizer.step() # 원본 model의 weight update
update_ema(ema_model, model) # EMA model의 weight update- EMA update는 일반적으로
optimizer.step()이후에 수행함. - 현재 학습 step에서 optimizer가 생성한 최신 원본 model의 weight를 EMA에 반영하기 위한 순서임.
- Gradient accumulation을 사용하는 경우에도 실제
optimizer.step()이 실행된 이후에 EMA model을 갱신하는 것이 일반적임. - Mini-batch마다 gradient를 누적하는 중간 단계에서는 원본 model의 weight가 아직 갱신되지 않았기 때문임.
11. 검증과 추론에서의 EMA model 사용
학습에는 원본 model을 사용하고, validation(검증) 또는 inference(추론)에서는 EMA model을 사용할 수 있음.
ema_model.eval()
with torch.no_grad():
prediction = ema_model(x_val)- 학습: 원본 model을 사용하여 gradient를 계산하고 optimizer로 weight를 갱신함
- EMA 갱신: 원본 model의 최신 weight를 이용하여 EMA model을 갱신함
- 검증 및 추론: EMA model을 평가 model로 사용할 수 있음
EMA model의 weight는 여러 학습 step의 원본 model weight를 평균낸 결과임.
- 따라서 특정 mini-batch에 의해 발생한 일시적인 weight 변동의 영향을 덜 받을 수 있음.
- 다만 EMA model이 항상 원본 model보다 높은 성능을 보장하는 것은 아님.
- 실제 학습에서는 원본 model과 EMA model의 validation performance를 모두 측정한 뒤 더 우수한 model을 최종 checkpoint로 선택하는 것이 적절함.
주의: BatchNorm buffer 처리
앞의 단순한 구현은 model.parameters()에 포함된 parameter만 EMA 대상으로 사용함.
그러나 BatchNorm에는 gradient로 학습되지 않는 다음과 같은 상태값이 존재함.
running_mean
running_var
num_batches_tracked- PyTorch에서는 이러한 값을 buffer라고 부름.
보통,
- 단순한 구현에서는 parameter에는 EMA를 적용하고,
- buffer는 원본 model에서 그대로 복사하는 방식이 이용됨.
@torch.no_grad()
def update_ema(
ema_model: nn.Module,
model: nn.Module,
decay: float = 0.999,
) -> None:
# 학습 가능한 parameter에 EMA 적용
for ema_param, model_param in zip(
ema_model.parameters(),
model.parameters(),
):
ema_param.mul_(decay)
ema_param.add_(
model_param,
alpha=1.0 - decay,
)
# BatchNorm 등의 buffer는 원본 model에서 복사
for ema_buffer, model_buffer in zip(
ema_model.buffers(),
model.buffers(),
):
ema_buffer.copy_(model_buffer)BatchNorm buffer의 처리 방법은 구현에 따라 달라질 수 있음.
- 원본 model의 buffer를 EMA model로 그대로 복사하는 방법
- Parameter와 함께 buffer에도 EMA를 적용하는 방법
- 학습 종료 후 EMA model의 BatchNorm statistics(통계)를 다시 계산하는 방법
Transformer처럼 BatchNorm 대신 LayerNorm을 사용하는 model에서는 일반적으로 BatchNorm running statistics 문제가 발생하지 않음.
참고: PyTorch 내장 기능 사용
PyTorch의 torch.optim.swa_utils.AveragedModel을 이용하여 Model EMA를 구현할 수도 있음.
get_ema_multi_avg_fn()을 전달하면 EMA 방식의 parameter averaging이 적용됨.
import torch
from torch.optim.swa_utils import (
AveragedModel,
get_ema_multi_avg_fn,
)
model = SimpleModel().to(device)
ema_model = AveragedModel(
model,
multi_avg_fn=get_ema_multi_avg_fn(
decay=0.999,
),
use_buffers=True,
)
학습 loop에서는 원본 model의 optimizer update 이후 update_parameters()를 호출함.
for x, y in train_loader:
x = x.to(device)
y = y.to(device)
optimizer.zero_grad()
prediction = model(x)
loss = criterion(prediction, y)
loss.backward()
optimizer.step()
ema_model.update_parameters(model)
각 코드의 역할은 다음과 같음.
loss.backward()
원본 model의 gradient를 계산하는 과정임.
optimizer.step()
계산된 gradient를 이용하여 원본 model의 weight를 갱신하는 과정임.
ema_model.update_parameters(model)원본 model의 최신 weight를 이용하여 EMA model을 갱신하는 과정임.
EMA model에서는 별도의 backward()가 실행되지 않음.
참고: Linear probing과 fine-tuning에서의 적용
Linear probing에서는 pretrained backbone을 고정하고 새로 추가한 classification head만 학습함.
- 따라서 학습 가능한 parameter가 classification head에 제한됨.
- 학습 step이 적고 classification head의 parameter 수가 작다면 Model EMA의 효과도 제한적일 수 있음.
이러한 이유들 때문에 보통 linear probing에서는 EMA를 사용하지 않는 편임.
반면 fine-tuning에서는 pretrained backbone의 일부 또는 전체 parameter를 학습함.
- 많은 parameter가 mini-batch gradient에 의해 계속 변경되므로, weight 변화가 linear probing보다 클 수 있음.
- 따라서 fine-tuning에서는 Model EMA를 적용하여 학습 중 발생하는 weight 변동을 완화할 가능성이 있음.
다만 EMA의 효과는 다음 조건에 따라 달라질 수 있음.
- Dataset 규모
- Batch size
- Learning rate
- 전체 학습 step 수
- EMA decay
- Data augmentation 강도
- Pretrained model과 downstream task 사이의 차이
- Backbone을 여는 범위
- Fine-tuning stage의 길이
따라서 fine-tuning에 EMA를 적용할 경우 원본 model과 EMA model의 validation performance를 함께 비교하는 것이 적절함.
요약
원본 model
- 학습 data를 입력받는 model
- loss를 계산하는 model
- gradient를 계산하는 model
- optimizer를 통해 weight가 갱신되는 model
EMA model
- 별도의 loss 및 gradient 를 계산하지 않는 model
- 별도의 optimizer를 사용하지 않는 model
- 원본 model의 weight를 지수 이동 평균하여 저장하는 model
Model EMA의 핵심 관계는 다음 세 줄로 정리됨.
loss.backward() # 원본 model의 gradient 계산
optimizer.step() # 원본 model의 weight update
update_ema(ema_model, model) # EMA model의 weight update즉, EMA model은 별도로 학습되는 두 번째 model이 아님.
원본 model의 학습 과정에서 생성되는 weight들을 지수 이동 평균하여 저장하는 shadow model임.
같이보면 좋은 자료들
https://docs.pytorch.org/docs/2.13/generated/torch.optim.swa_utils.get_ema_avg_fn.html
torch.optim.swa_utils.get_ema_avg_fn — PyTorch 2.13 documentation
Rate this Page ★ ★ ★ ★ ★
docs.pytorch.org
https://docs.pytorch.org/docs/2.13/generated/torch.optim.swa_utils.AveragedModel.html
AveragedModel — PyTorch 2.13 documentation
Rate this Page ★ ★ ★ ★ ★
docs.pytorch.org
https://dsaint31.tistory.com/860
[Math] Exponential Moving Average (EMA)
Definitionmoving average (=rolling average or running average)데이터 포인트를 분석하기 위한전체 데이터 집합에서 연속적인 부분 집합들의 평균을 계산하는통계적 방법이다.일반적으로 MA는 다음과 같이 모
dsaint31.tistory.com
'ML' 카테고리의 다른 글
| Patch Embedding - Vision Transformers (0) | 2026.07.19 |
|---|---|
| CNN / Vision Transformer 모델별 Dropout 위치와 Fine-tuning 시 설정 (0) | 2026.07.14 |
| DL Vision Model들에서 Stem이란? (0) | 2026.07.13 |
| Patch Merging of SWIN (0) | 2026.07.11 |
| optuna : Automatic Hyperparameter Optimization Framework (0) | 2026.06.27 |
