Hàm Kích Hoạt Cổng Tuyến Tính SwiGLU (Shazeer 2020) LLaMA

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

• Trong Transformer ban đầu, tầng FFN dùng 2 ma trận: FFN(x) = ReLU(x W₁ + b₁) W₂ + b₂.
• Shazeer (2020) chứng minh cấu trúc Cổng Tuyến Tính Gated Linear Units (GLU) vượt trội hơn hẳn. Biến thể tốt nhất là SwiGLU:

SwiGLU(x) = (Swish(x W_gate) ⊗ x W_up) W_down Trong đó:

• Swish(z) = z · σ(z) = SiLU(z).
• ⊗ là phép nhân từng phần tử (Hadamard product).
• Tầng ẩn thường có chiều rộng d_ffn = ⌊ 23 · 4 · d_model ⌋ = ⌊ 83 d_model ⌋ để giữ nguyên tổng số tham số so với FFN 2 tầng truyền thống.

Yêu cầu

Xây dựng lớp SwiGLUFFN(torch.nn.Module):

class SwiGLUFFN(nn.Module):
    def __init__(self, d_model: int, d_ffn: int):
        pass

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        pass

• Sử dụng 3 phép biến đổi tuyến tính không có bias (bias=False theo chuẩn LLaMA):
• w_gate: Linear (d_model -> d_ffn)
• w_up: Linear (d_model -> d_ffn)
• w_down: Linear (d_ffn -> d_model)

Ràng buộc

• Thời gian chạy tối đa: 2000ms.
• 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ả.