Exponential Moving Average trong Deep Learning

MayFest2024 ContentCreator

Chất lượng của mô hình học sâu (deep learning) có liên quan chặt chẽ đến quá trình huấn luyện chúng. Để huấn luyện được mô hình tốt, việc giảm nhiễu (noise) từ quá trình cập nhật ngẫu nhiên (stochastic updates) là cần thiết. Cách chính quy, đã được chứng minh bằng toán học, dùng để giảm nhiễu trong tối ưu hàm lồi phải kể đến (tail) average. Ứng dụng trong học sâu (deep learning), để có được mô hình tốt hơn, kỹ thuật lấy trung bình các bộ trọng số (weights) được tạo ra trong quá trình train là lựa chọn không tồi.

Trong bài viết này, kỹ thuật lấy trung bình Exponential Moving Average (EMA) sẽ được đề cập. EMA được sử dụng như là một trick để cải thiện độ chính xác của model. Ví dụ một cách sử dụng, thay vì sử dụng bộ trọng số (weight) từ training iteration cuối cùng, thì EMA của các bộ trọng số (weights) trong quá trình huấn luyện sẽ được dùng để dự đoán.

Giới thiệu về noise

Trong học sâu (deep learning), cho tập dữ liệu dùng để huấn luyện $D$ bao gồm $N$ điểm dữ liệu, thông thường mục tiêu muốn đạt được là tối ưu bộ trọng số $\Theta$ của mạng nơ-ron để tối thiểu hóa hàm mất mát (loss) $L(D, \Theta)$ . Gọi bộ trọng số (weight) sau khi được tối ưu là $\Theta^{\\*}$ .

Trong thực tế, giá chị của $N$ thường rất lớn, nên việc tối ưu trọng số (weight) mà dùng toàn bộ dữ liệu trong tập huấn luyện là vô cùng "đắt đỏ", đặc biệt là khi dùng các kỹ thuật như gradient descent. Chính vì vậy, thay vì dùng cả bộ dữ liệu, một lượng nhỏ điểm dữ liệu sẽ được dùng để huấn luyện một lúc, gọi là mini-batch, và chúng được lấy một cách ngẫu nhiên từ tập huấn luyện. Điều này sẽ dẫn đến nhiễu (noise). Nhiễu trong quá trình huấn luyện (training noise) có cả mặt tốt, có cả mặt xấu.

Trong gradient descent, đạo hàm (gradient) để cập nhật bộ trọng số (weight) được tính bằng tất cả dữ liệu trong tập huấn luyện. Tuy nhiên như đã đề cập, việc cùng một lúc tính đạo hàm mà dùng tât cả dữ liệu huấn luyện sẽ rất "đắt đỏ", nên nó được ước tính (estimate) bằng đạo hàm (gradient) của mini-batch. Nhờ đó, đạo hàm (gradient) được tính toán nhanh hơn, cũng đi kèm với việc giảm thiểu độ chính xác. Nhưng, việc sử dụng đạo hàm ước tính (estimated gradient), hay (noisy gradient) đôi khi hữu ích trong việc tối ưu, kết quả là đạt được cực trị địa phương tốt hơn so với khi huấn luyện dùng toàn bộ dữ liệu (để tính đạo hàm cùng một lúc). Đôi khi, quá trình tối ưu bộ trọng số (weight) chuẩn bị hội tụ, nhiễu (noise) của đạo hàm (gradient) khi dùng mini-batch có thể làm cho $\Theta^{\\*}$ không phải là cực tiểu, mà $\Theta$ dao động xung quang cực tiểu.

Kĩ thuật giảm nhiễu (noise) như EMA được dùng để khắc phục tình trạng đó.

Giới thiệu về Moving Average

Moving average là kĩ thuật làm mịn (smoothing technique) thường được dùng để giảm nhiễu và biến động trong dữ liệu chuỗi thời gian (time-series data).

Dạng đơn giản nhất của moving average là simple moving average (SMA), được tính bằng giá trị trung bình của $k$ điểm dữ liệu trước đó. Cụ thể, cho chuỗi dữ liệu $p_1, p_2, p_3, ..., p_t, ...$ , SMA trên $k$ điểm dữ liệu gần nhất so với điểm dữ liệu $p_t$ , được ký hiệu là ${SMA}_{t}$ :

$\begin{equation} {SMA}_t = \frac { 1 } { k } \sum _{ i = t-k+1 } ^ { t } p _ { i } \end{equation}$

${SMA}_t$ có thể tính bằng ${SMA}_{t-1}$ . Điều này hữu ích khi dùng SMA để tính toán với streaming data:

$\begin{equation} {SMA}_t = {SMA}_{t-1} + \frac{1}{k} (p_t - p_{t-k}) \end{equation}$

Về mặt toán học, moving average là một phép toán, dùng để tính trung bình của một số lượng nhất định các điểm dữ liệu trước đó.

Giới thiệu về Exponential Moving Average (EMA)

Trong học sâu (deep learning), Exponential Moving Average (EMA) tính trung bình có trọng số của các bộ trọng số (weights). Cụ thể, cho dãy các trọng số thu được trong quá trình huấn luyện $\Theta_1, \Theta_2, \Theta_3, ..., \Theta_t, ...$ , trong đó $\Theta_t$ là bộ trọng số (weight) thu được sau training iteration thứ $t$ , EMA của trọng số (weight) tại thời điểm thứ $t$ (training iteration thứ $t$ ) ký hiệu là ${EMA}_t$ .

Công thức của EMA thường được tính như sau:

$\begin{equation} EMA_{t} = \begin{cases} \Theta_1 & \text{if $iteration = 1$}\\ \alpha * {EMA}_{t-1} + (1-\alpha) * \Theta_t & \text{if $iteration = t$} \end{cases} \end{equation}$

Trong đó, $\alpha \in[0, 1)$ , được gọi là EMA decay. Thông giá trị của $\alpha$ được chọn là $0.99, 0.999, ...$

Để hiểu rõ hơn ý nghĩa phía sau của EMA, cùng nhau phân tích $(3)$ .

$\begin{equation} \begin{aligned} EMA_{t} & = (1-\alpha) [\Theta_t + \alpha\Theta_{t-1} + \alpha^2\Theta_{t-2} + ... + \alpha^{t-2}\Theta_2] + \alpha^{t-1}\Theta_1 \\ & \approx (1-\alpha) [\Theta_t + \alpha\Theta_{t-1} + \alpha^2\Theta_{t-2} + ... + \alpha^{t-2}\Theta_2 + \alpha^{t-1}\Theta_1] \\ &= (1-\alpha) \sum_{i=1}^t \alpha^{t-i} \Theta_{i} \end{aligned} \end{equation}$

Trong đó, tổng của các hệ số xấp xỉ 1 khi $t \rightarrow \infty$

Một vài nhận xét của bản thân tác giả: Giá trị của $EMA_t$ chịu ảnh hưởng nhiều hơn bởi các weights gần với training iteration $t$ hơn.

Code EMA trong deep learning

Mẫu code tham khảo về cách EMA được dùng trong deep learning.

class ModelEmaV2(nn.Module):
    def __init__(self, model, decay=0.9999, device=None):
        super(ModelEmaV2, self).__init__()
        # make a copy of the model for accumulating moving average of weights
        self.module = deepcopy(model)
        self.module.eval()
        self.decay = decay
        self.device = device  # perform ema on different device from model if set
        if self.device is not None:
            self.module.to(device=device)

    def _update(self, model, update_fn):
        with torch.no_grad():
            for ema_v, model_v in zip(self.module.state_dict().values(), model.state_dict().values()):
                if self.device is not None:
                    model_v = model_v.to(device=self.device)
                ema_v.copy_(update_fn(ema_v, model_v))

    def update(self, model):
        self._update(model, update_fn=lambda e, m: self.decay * e + (1. - self.decay) * m)

    def set(self, model):
        self._update(model, update_fn=lambda e, m: m)