SwiGLU

SwiGLU (đơn vị tuyến tính có cổng Swish) là một biến thể có cổng của lớp feed-forward trong Transformer. Cơ chế này chiếu vector của mỗi token hai lần, đưa một phép chiếu qua hàm kích hoạt SiLU (Swish), nhân từng phần tử của hai phép chiếu rồi chiếu tích trở lại chiều rộng của model: FFN(x) = (SiLU(xW) ⊙ xV)W₂. Noam Shazeer giới thiệu phương pháp này trong bài báo năm 2020 GLU Variants Improve Transformer, và từ đó nó trở thành khối feed-forward trong các mô hình ngôn ngữ lớn (LLM) như PaLM, Llama, Mistral và Qwen. Khác với ReLU hoặc GELU, vốn tác động độc lập lên từng số, SwiGLU chứa trọng số được học, khiến nó trở thành một lớp nhỏ riêng biệt.

Cách SwiGLU hoạt động#

Mạng feed-forward (FFN) tiêu chuẩn của Transformer mở rộng vector của mỗi token bằng một ma trận trọng số, áp dụng hàm kích hoạt rồi nén trở lại bằng ma trận thứ hai. SwiGLU thay ma trận đầu tiên và hàm kích hoạt bằng một đơn vị tuyến tính có cổng (GLU), gồm hai phép chiếu song song của cùng một đầu vào:

• Phép chiếu cổng: xW đi qua SiLU, x · σ(x), tạo ra các giá trị gần bằng không với đầu vào âm mạnh và gần bằng chính đầu vào với đầu vào dương.
• Phép chiếu giá trị: xV giữ dạng tuyến tính và mang nội dung.
• Tích từng phần tử: Nhân hai phép chiếu cho phép cổng tăng, giảm hoặc đưa gần về không từng đặc trưng ẩn tùy theo đầu vào; phép chiếu đầu ra W₂ ánh xạ kết quả trở lại chiều rộng của model.

Ý tưởng cổng bắt nguồn từ GLU của Dauphin và cộng sự, những người đã dùng cổng sigmoid trong các mô hình ngôn ngữ tích chập có cổng vào năm 2016. Bài báo của Shazeer thử nghiệm cùng cấu trúc với các hàm cổng khác nhau và đặt tên từng biến thể theo hàm kích hoạt: GLU dùng sigmoid, ReGLU dùng ReLU, GEGLU dùng GELU, SwiGLU dùng Swish, còn lớp song tuyến tính không dùng hàm kích hoạt. Theo codebase T5, các lớp này bỏ các thành phần bias.

Vì sao SwiGLU thay thế ReLU và GELU trong LLM#

Các lớp feed-forward có cổng có ba ma trận trọng số thay vì hai, vì vậy Shazeer giảm chiều rộng ẩn xuống hệ số 2/3 để giữ nguyên số lượng tham số và phép tính; trong model T5-base, số đơn vị giảm từ 3,072 xuống 2,048. Với cùng kích thước, SwiGLU đạt log-perplexity trên tập giữ lại là 1.636 sau 524,288 bước huấn luyện, so với 1.677 của ReLU và 1.679 của GELU; phương pháp này cũng đạt điểm trung bình cao hơn trên các benchmark hiểu ngôn ngữ GLUE và SuperGLUE. Bài báo không đưa ra lý thuyết giải thích mức cải thiện; kết luận cho rằng thành công của các kiến trúc này "cũng như mọi điều khác, là nhờ lòng nhân từ của thần linh."

Kết quả này cũng đúng với các model lớn. PaLM của Google áp dụng SwiGLU vào năm 2022, còn LLaMA của Meta thay ReLU bằng SwiGLU vào năm 2023, với chiều rộng ẩn bằng 2/3 × 4d, làm tròn lên bội số của 256; cách tính này cho ra 11,008 đơn vị ẩn ở model 7B có 4,096 chiều. Llama 2, Llama 3, Mistral và Qwen tiếp tục dùng thiết kế này; trong các model hỗn hợp chuyên gia (MoE) như Mixtral, mỗi chuyên gia là một khối feed-forward SwiGLU.

Triển khai SwiGLU bằng PyTorch#

Đoạn code PyTorch dưới đây xây dựng một lớp feed-forward ReLU tiêu chuẩn và một lớp SwiGLU cho model 512 chiều, rồi so sánh kích thước của chúng:

import torch
import torch.nn.functional as F
from torch import nn

d_model = 512
x = torch.randn(4, d_model)  # four token vectors

# Standard transformer FFN: two matrices around a 4x wider ReLU hidden layer
relu_ffn = nn.Sequential(nn.Linear(d_model, 2048, bias=False), nn.ReLU(), nn.Linear(2048, d_model, bias=False))

# SwiGLU FFN: three matrices, hidden width cut to about 2/3 of 2048
w_gate = nn.Linear(d_model, 1365, bias=False)
w_value = nn.Linear(d_model, 1365, bias=False)
w_out = nn.Linear(1365, d_model, bias=False)
y = w_out(F.silu(w_gate(x)) * w_value(x))  # (SiLU(xW) * xV) W2

relu_params = sum(p.numel() for p in relu_ffn.parameters())
swiglu_params = sum(m.weight.numel() for m in (w_gate, w_value, w_out))
print(f"ReLU FFN:   output {tuple(relu_ffn(x).shape)}, {relu_params:,} parameters")
print(f"SwiGLU FFN: output {tuple(y.shape)}, {swiglu_params:,} parameters")

Kết quả in ra là ReLU FFN: output (4, 512), 2,097,152 parameters và SwiGLU FFN: output (4, 512), 2,096,640 parameters. Với 1,365 đơn vị ẩn, xấp xỉ 2/3 của 2,048, ba ma trận SwiGLU có gần như cùng số trọng số với hai ma trận ReLU, vì vậy mọi khác biệt về chất lượng đều đến từ cơ chế cổng. Nhiều cách triển khai gộp ma trận cổng và ma trận giá trị thành một phép chiếu rồi chia đôi đầu ra, nhờ đó chỉ cần chạy một phép nhân ma trận lớn thay vì hai phép nhân nhỏ.

SwiGLU và các thuật ngữ liên quan#

• SiLU (Swish): SiLU là hàm tác động từng phần tử bên trong cổng. SwiGLU là toàn bộ lớp có cổng với ba phép chiếu.
• GLU và GEGLU: Các biến thể cùng nhóm sử dụng cổng sigmoid hoặc GELU. Trong các thử nghiệm của Shazeer, GEGLU đạt kết quả ngang SwiGLU với log-perplexity là 1.633; các model Gemma của Google sử dụng GEGLU.
• Attention: Attention cũng tạo ra trọng số phụ thuộc dữ liệu, nhưng trộn thông tin giữa các token. SwiGLU điều tiết các đặc trưng độc lập bên trong từng token.
• Model thị giác: Các lớp feed-forward có cổng cũng xuất hiện trong model thị giác và đa phương thức. YOLOE, model phát hiện từ vựng mở do các nhà nghiên cứu tại Đại học Thanh Hoa phát triển và chạy trong package Ultralytics, tinh chỉnh embedding văn bản của prompt lớp bằng khối SwiGLU dư; khối này chỉ chạy một lần mỗi lần gọi set_classes() và được gộp bỏ khi export nên không làm tăng chi phí xử lý mỗi khung hình.