Cài đặt Thuật toán Tối ưu Hóa AdamW (Decoupled Weight Decay) từ Gốc

Mục tiêu kiến thức

• Hiểu lỗi sai kinh điển trong Adam gốc: Adam gốc cộng L2 weight decay trực tiếp vào gradient trước khi chia cho √(v_t). Điều này khiến các trọng số có gradient lớn lại bị phạt ít hơn các trọng số có gradient nhỏ.
• Thuật toán AdamW (Loshchilov & Hutter, 2017) tách rời hoàn toàn (Decoupled):
• Tăng bước đếm thời gian: t ← t + 1.
• Cập nhật moment bậc 1: m_t = β₁ m_t-1 + (1 - β₁) g_t.
• Cập nhật moment bậc 2: v_t = β₂ v_t-1 + (1 - β₂) g_t².
• Hiệu chỉnh độ chệch (Bias Correction):

m_t = m_t1 - β₁^t, v_t = v_t1 - β₂^t

• Cập nhật trọng số kết hợp Decoupled Weight Decay:

w_t ← w_t-1 - α ( m_t√(v_t) + ε + λ w_t-1 ) (hoặc tương đương w_t ← w_t-1(1 - α λ) - α m_t√(v_t) + ε).

Yêu cầu

Xây dựng lớp AdamWOptimizer:

class AdamWOptimizer:
    def __init__(self, params: list[np.ndarray], lr: float = 1e-3, betas: tuple[float, float] = (0.9, 0.999), eps: float = 1e-8, weight_decay: float = 0.01):
        pass

    def step(self, grads: list[np.ndarray]) -> None:
        pass

• Cập nhật tham số tại chỗ (in-place mutation).

Ràng buộc

• Thời gian chạy tối đa: 3000ms.
• Giới hạn bộ nhớ: 256MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.