Đặt vấn đề

Có một thời chúng ta nghĩ rằng muốn xây một bộ não lớn hơn thì phải bắt toàn bộ bộ não làm việc nhiều hơn.

Transformer cũng từng sống theo luật đó.

Muốn mô hình thông minh hơn?

Thêm layer.

Muốn thêm kiến thức?

Tăng hidden dimension.

Muốn nhiều parameter hơn?

Làm Feed Forward Network to hơn.

Nghe thì có vẻ đơn giản. Nhưng mỗi lần mô hình phình ra, những con chip ở phía dưới cũng phải nhân những ma trận lớn hơn, bộ nhớ phải chứa nhiều trọng số hơn, dữ liệu phải chạy qua nhiều GPU hoặc TPU hơn, những dòng điện nhỏ xíu phải lao qua hàng tỷ transistor nhanh hơn nữa.

Một tỷ parameter đã nặng.

Mười tỷ nặng hơn.

Một trăm tỷ bắt đầu giống xây nhà bằng xe tải.

Nếu cứ đi theo con đường đó thì mô hình một nghìn tỷ parameter nghe gần giống chuyện:

Muốn chở thêm một người nên quyết định xây thêm một con đường cao tốc.

Nhưng năm 2021, William Fedus, Barret Zoph và Noam Shazeer tại Google đưa ra một ý tưởng khá buồn cười.

Tại sao phải bắt toàn bộ mô hình làm việc?

Nếu chúng ta có một nghìn tỷ parameter nhưng mỗi token chỉ sử dụng một phần rất nhỏ trong số đó thì sao?

Đó chính là Switch Transformer.

Switch Transformer không phải nơi Mixture of Experts (MoE) được phát minh. Ý tưởng MoE đã xuất hiện từ nhiều năm trước. Điều quan trọng hơn là Switch Transformer đã đơn giản hóa MoE đến mức có thể huấn luyện ổn định ở quy mô cực lớn, và nhóm tác giả đã thực sự đưa một Transformer lên khoảng 1,6 nghìn tỷ parameter với 2.048 experts.

Và ở đây lịch sử của Transformer rẽ sang một con đường rất thú vị:

số parameter của mô hình không còn nhất thiết phải tăng cùng tốc độ với lượng computation dành cho mỗi token.

Trước Switch Transformer, mô hình càng to càng phải tính nhiều

Transformer năm 2017 đã giải quyết một vấn đề rất quan trọng của các kiến trúc tuần tự trước đó: thay vì đọc câu từng từ như một đoàn tàu chạy qua từng ga, self-attention cho phép các token xem xét quan hệ với nhau theo cách phù hợp hơn với phần cứng tính toán song song. (Google Research)

Một Transformer block, nếu giản lược rất mạnh, có thể tưởng tượng như:

1Token
2  ↓
3Self-Attention
4  ↓
5Feed Forward Network
6  ↓
7Token mới

Attention giúp các token nhìn nhau.

Còn Feed Forward Network, thường gọi là FFN, thực hiện một lượng rất lớn phép nhân ma trận cho từng token.

Ví dụ đơn giản:

$$ h = \text{ReLU}(xW_1) $$

$$ y = hW_2 $$

Hai ma trận $W_1$ và $W_2$ chứa rất nhiều parameter.

Muốn Transformer lớn hơn, một cách tự nhiên là làm các ma trận này lớn hơn.

Nhưng khi đó có một luật vật lý khá khó chịu:

1nhiều parameter hơn
2        ↓
3ma trận lớn hơn
4        ↓
5nhiều phép nhân hơn
6        ↓
7nhiều FLOPs hơn
8        ↓
9training lâu hơn

Đây là mô hình dense.

Mỗi token bước vào mạng neural giống một vị khách đi vào nhà hàng và bắt toàn bộ đầu bếp trong bếp cùng đứng dậy nấu một tô mì.

Có 1.000 đầu bếp?

Cả 1.000 ông cùng nấu.

Có 10.000 đầu bếp?

Cả 10.000 ông vẫn cùng nấu.

Khá đoàn kết.

Và cũng khá tốn điện.

Mixture of Experts: thuê rất nhiều chuyên gia nhưng không gọi tất cả cùng lúc

Mixture of Experts đưa ra một triết lý khác.

Giả sử chúng ta có một bệnh viện với:

• bác sĩ tim mạch
• bác sĩ da liễu
• bác sĩ thần kinh
• bác sĩ mắt
• bác sĩ xương
• bác sĩ nhi
• vài trăm ông khác đang ngồi uống cà phê

Một bệnh nhân đau mắt đi vào.

Không cần gọi 2.048 bác sĩ xếp hàng khám người đó.

Chúng ta cần một router đứng ở cửa:

1Bệnh nhân
2    ↓
3 Router
4    ↓
5Bác sĩ phù hợp

Trong Transformer cũng vậy.

Thay vì một FFN duy nhất:

1Token
2  ↓
3FFN

ta tạo ra rất nhiều FFN:

1                    ┌── Expert 1
2                    ├── Expert 2
3Token → Router ─────├── Expert 3
4                    ├── ...
5                    └── Expert N

Mỗi FFN được gọi là một expert.

Đây là conditional computation:

Mô hình có thể sở hữu rất nhiều parameter, nhưng với một input cụ thể chỉ một phần parameter được kích hoạt.

Paper Switch Transformer mô tả mục tiêu này khá trực tiếp: tăng số lượng parameter trong khi giữ FLOPs trên mỗi example gần như không đổi.

Đây chính là cái khe nhỏ mà những con quái vật nghìn tỷ parameter sau này có thể chui qua.

Nhưng MoE trước đó khá phiền

Ý tưởng MoE không mới.

Paper Switch Transformer nhắc lại những công trình MoE từ thập niên 1990 và đặc biệt là sparse MoE của Shazeer và cộng sự năm 2017. Trước Switch, GShard cũng đã đưa Transformer MoE lên quy mô hàng trăm tỷ parameter.

Nhưng MoE mang theo một túi vấn đề.

Router phải chọn expert.

Token phải được chuyển giữa các accelerator.

Expert này có thể nhận quá nhiều token.

Expert kia có thể ngồi ngáp.

Training có thể mất ổn định.

Và một thiết kế MoE phổ biến lúc đó sử dụng Top-k routing.

Ví dụ Top-2:

1Token "con mèo"
2       ↓
3     Router
4      ↙   ↘
5Expert 7  Expert 23

Cùng một token được gửi tới hai expert.

Output có dạng:

$$ y = \sum_{i \in T} p_i(x)E_i(x) $$

Trong đó:

• $x$ là representation của token
• $E_i$ là expert thứ $i$
• $p_i(x)$ là xác suất router gán cho expert đó
• $T$ là tập các expert được chọn

Top-2 có cái hay của nó.

Nhưng token phải chạy đến hai nơi, hai FFN phải tính toán, rồi kết quả phải được ghép lại.

Ở quy mô vài nghìn accelerator, chuyện một vector nhỏ chạy từ chip này sang chip kia không còn nhỏ nữa.

Hàng tỷ token đang chạy.

Mỗi token kéo theo những tensor.

Những tensor đi xuyên qua interconnect.

Các accelerator gửi dữ liệu qua lại.

Và trong một hệ thống phân tán khổng lồ, đôi khi thứ giết tốc độ không phải phép nhân.

Mà là chờ hàng xóm gửi dữ liệu sang.

Switch Transformer: thôi, chọn một ông thôi

Đây là phần vừa đơn giản vừa quan trọng nhất của Switch Transformer.

Nhóm tác giả đặt câu hỏi:

Tại sao một token phải đi qua nhiều expert?

Thử chỉ chọn một expert xem sao.

Router tính:

$$ h(x) = W_r x $$

sau đó chuyển logits thành xác suất:

$$ p_i(x)=\frac{e^{h_i(x)}}{\sum_{j=1}^{N} e^{h_j(x)}} $$

Và cuối cùng:

$$ i^* = \arg\max_i p_i(x) $$

Chọn đúng một expert.

Output trở thành:

$$ y = p_{i^}(x)E_{i^}(x) $$

Hết.

Xong thuật toán, quá dễ.

Trong paper, nhóm tác giả gọi chiến lược $k=1$ này là Switch layer. Họ báo cáo rằng việc chỉ route mỗi token tới một expert vẫn giữ được chất lượng, đồng thời giảm computation của router, giảm expert capacity cần thiết và giảm communication cost.

Nếu Top-2 giống:

“Không biết hỏi ai nên gửi email cho hai phòng ban.”

thì Switch Transformer nói:

“Chọn một phòng thôi. Đừng CC cả công ty.”

Đây là thay đổi cực nhỏ trên sơ đồ.

Nhưng khi nhân nó lên hàng tỷ token và hàng nghìn accelerator, thay đổi nhỏ ấy biến thành một con sông computation.

Một token bước vào Switch Transformer như thế nào?

Giả sử có câu:

1Con mèo đang ngủ.

Sau tokenizer, cứ giả sử ta có các token:

1["Con", "mèo", "đang", "ngủ"]

Và một Switch layer có bốn expert.

Router có thể quyết định:

1"Con"  → Expert 2
2"mèo"  → Expert 4
3"đang" → Expert 1
4"ngủ"  → Expert 4

Expert 3 lần này chẳng làm gì.

Batch khác đến, Expert 3 có thể lại nhận rất nhiều token.

Điểm quan trọng nằm ở đây.

Mô hình sở hữu parameter của:

1Expert 1
2Expert 2
3Expert 3
4Expert 4
5...
6Expert 2048

nhưng token "mèo" không chạy qua tất cả 2.048 expert.

Nó chỉ chạy qua một expert được router chọn.

Đây là lý do phải phân biệt hai khái niệm:

• total parameters: toàn bộ parameter mà mô hình sở hữu
• active parameters / computation per token: phần thực sự tham gia xử lý một token

Hai con số trước đây thường đi cùng nhau.

MoE bắt đầu kéo chúng ra xa nhau.

Switch Transformer kéo mạnh hơn nữa.

1,6 nghìn tỷ parameter nhưng không tính 1,6 nghìn tỷ parameter cho mỗi token

Đây có lẽ là phần quan trọng nhất của cả bài.

Nhìn vào một số model trong paper:

Các số liệu này nằm trong Table 9 của paper. Switch-C sử dụng 2.048 experts và có khoảng 1,571 nghìn tỷ parameter.

Khoan.

Một mô hình 1.571T parameter lại có FLOPs/sequence thấp hơn T5-XXL 11B trong cấu hình được paper so sánh?

Đúng.

Bởi vì 1.571T là số parameter đang tồn tại.

Không phải tất cả parameter đều thức dậy khi một token đi ngang qua.

Hãy tưởng tượng một thành phố có 2.048 thư viện.

Thành phố sở hữu hàng tỷ cuốn sách.

Nhưng khi một cậu bé muốn biết:

“Khủng long có biết bơi không?”

cậu bé không cần khiêng toàn bộ 2.048 thư viện về nhà.

Router chỉ cần chỉ:

1→ thư viện số 742

Năng lực lưu trữ tri thức của thành phố có thể tăng rất lớn.

Nhưng công việc cần làm cho từng câu hỏi không nhất thiết tăng tương ứng.

Đây là sự tách rời quan trọng:

$$ \text{Model Capacity} \neq \text{Compute per Token} $$

Một mô hình có thể rất lớn về capacity, trong khi vẫn tương đối nhỏ về active computation.

Nhưng nếu mọi token đều thích cùng một expert thì sao?

Bây giờ xuất hiện một vấn đề rất con người.

Giả sử có 100 quầy bán vé.

Router phát hiện quầy số 7 hơi ngon.

Thế là:

11.000 người → quầy số 7
299 quầy còn lại → ngồi chơi

Hệ thống chết.

Trong MoE, hiện tượng tương tự gọi là load imbalance.

Một số expert có thể nhận quá nhiều token, trong khi các expert khác gần như không được dùng.

Switch Transformer thêm một auxiliary loss nhằm khuyến khích router phân phối token tương đối đều giữa các expert.

Paper viết loss này dưới dạng:

$$ L_{\text{balance}}=\alpha N\sum_{i=1}^{N}f_i P_i $$

Trong đó:

• $N$ là số expert
• $f_i$ là tỷ lệ token thực sự được gửi tới expert $i$
• $P_i$ là lượng xác suất trung bình router dành cho expert $i$
• $\alpha$ điều khiển mức quan trọng của loss phụ

Nhóm tác giả sử dụng $\alpha = 10^{-2}$ trong các thí nghiệm chính để cân bằng tải mà không lấn át cross-entropy loss.

Nói nôm na:

Router được học rằng đừng có cả lớp 40 đứa cùng chạy tới hỏi một cô giáo trong khi 39 giáo viên khác đang ngồi nhìn trần nhà.

Expert Capacity: mỗi chuyên gia cũng chỉ có hai tay

Ngay cả khi có load balancing, phân phối token không bao giờ hoàn hảo.

Vì vậy mỗi expert có một giới hạn gọi là expert capacity.

Paper định nghĩa:

$$ \text{Expert Capacity}=\frac{\text{Tokens per Batch}}{\text{Number of Experts}}\times\text{Capacity Factor} $$

Ví dụ:

• batch có 10.000 token
• 100 experts
• capacity factor = 1.25

thì:

$$ \frac{10{,}000}{100}\times 1.25=125 $$

Mỗi expert có chỗ cho khoảng 125 token.

Capacity factor lớn hơn tạo thêm khoảng đệm khi router phân phối không đều.

Nhưng buffer càng lớn thì càng tốn memory và communication.

Nếu expert đầy, một số token có thể bị bỏ qua tại Switch layer và đi tiếp qua residual connection. Trong các thí nghiệm của paper, tỷ lệ dropped token thường dưới 1% khi load balancing được cấu hình phù hợp.

Đây là một ví dụ rất đẹp của engineering.

Không có:

“Mọi thứ phải hoàn hảo.”

Chỉ có:

“Cho hệ thống một cái hành lang đủ rộng để đám đông đi qua mà không cần xây cả sân vận động.”

Tại sao Switch Transformer lại huấn luyện nhanh hơn?

Đến đây cần làm rõ một câu rất dễ bị viết sai:

Switch Transformer không đơn giản là “mô hình 1,6T chạy nhanh hơn mọi Transformer nhỏ hơn”.

Speedup trong paper được đo trong các cấu hình và baseline cụ thể.

Với các mô hình Switch-Base dựa trên T5, nhóm tác giả báo cáo mức pre-training speedup hơn 7 lần trong một số thí nghiệm khi giữ FLOPs/example tương đương. Một Switch-Base 64 experts đạt cùng chất lượng với T5-Base chỉ trong khoảng 1/7 thời gian training trong thí nghiệm trên 32 TPUv3 cores.

Ở quy mô lớn hơn, Switch-C đạt khoảng 4x speedup tới một mức perplexity cố định so với T5-XXL với cùng compute budget trong thí nghiệm được công bố.

Tức là chữ “nhanh” ở đây nên hiểu là:

đạt một mức chất lượng pre-training nhất định bằng ít bước/thời gian hơn trong cấu hình so sánh của paper.

Không phải cứ thay FFN bằng MoE rồi chương trình tự nhiên chạy nhanh gấp bảy.

Paper thậm chí nhấn mạnh rằng tốc độ cuối cùng còn phụ thuộc implementation và low-level optimization.

Nhưng nguyên nhân kiến trúc phía sau speedup thì rất đáng chú ý.

Switch Transformer có thể tăng model capacity bằng cách thêm experts mà không buộc mỗi token chạy qua toàn bộ số experts đó.

Nhiều parameter hơn.

Nhiều chỗ để học hơn.

Nhưng lượng compute cho từng token không tăng theo cùng tỷ lệ.

Một nghìn tỷ parameter bắt đầu trở nên khả thi

Switch-C là ví dụ cực đoan nhất.

Paper công bố cấu hình:

1Switch-C
2≈ 1.571 trillion parameters
32,048 experts
4Top-1 routing

Điều đáng chú ý không phải chỉ là con số 1,6 nghìn tỷ.

Nếu mục tiêu chỉ là khoe parameter, người ta có thể tạo một khối parameter khổng lồ rồi để nó nằm đó.

Điều quan trọng là Google đã xây được một hệ thống thực sự pre-train được bằng cách kết hợp:

• data parallelism
• model parallelism
• expert parallelism
• sparse activation
• distributed routing

Các expert có thể nằm trên những thiết bị khác nhau. Token được router phân loại rồi gửi tới accelerator đang chứa expert tương ứng.

Lúc này một Transformer không còn giống một bộ não nguyên khối.

Nó giống một thành phố.

Attention là hệ thống đường sá chung.

Các layer dense là hạ tầng dùng chung.

Còn những expert giống hàng nghìn khu công nghiệp, mỗi khu chứa một lượng parameter khổng lồ.

Một token chạy vào thành phố.

Router bật một cái công tắc.

Một khu sáng lên.

2.047 khu còn lại tối.

Và chính chữ Switch nằm ở đó.

Huấn luyện một con quái vật sparse không hề dễ

Nghe đến đây có thể thấy MoE giống cheat code:

1while model_not_smart:
2    add_more_experts()

Tiếc là đời không dễ vậy.

Sparse models có thể khó train hơn dense Transformer.

Router phải dùng softmax.

Routing có hard decision.

Một số expert có thể nhận quá nhiều token.

Gradient có thể trở nên mất ổn định.

Đặc biệt khi chuyển sang precision thấp để tăng tốc training, những vấn đề số học bắt đầu xuất hiện.

Switch Transformer sử dụng một thủ thuật khá đẹp gọi là selective precision.

Phần lớn model vẫn chạy bằng bfloat16.

Nhưng computation nhạy cảm trong router được chuyển sang float32.

Sau đó kết quả lại được cast về bfloat16 trước communication.

Trong thí nghiệm của paper, pure bfloat16 bị diverge, trong khi selective precision giữ độ ổn định gần float32 nhưng tốc độ gần với bfloat16.

Ngoài ra nhóm tác giả giảm initialization scale của Transformer xuống 10 lần, từ 1.0x xuống 0.1x, nhằm cải thiện stability khi số experts tăng rất lớn.

Đây là một chi tiết quan trọng.

Một ý tưởng kiến trúc hay chưa đủ.

Muốn biến nó thành mô hình nghìn tỷ parameter, phải giải quyết cả một tầng vật lý bên dưới:

 1routing
 2↓
 3tensor layout
 4↓
 5memory
 6↓
 7precision
 8↓
 9communication
10↓
11accelerator
12↓
13network

AI cuối cùng vẫn không sống trong mây.

Nó sống trong transistor.

Cái Switch thực sự thay đổi là cách chúng ta nghĩ về “mô hình lớn”

Trước sparse MoE, chúng ta dễ hình dung:

$$ \text{Model lớn hơn} \Rightarrow \text{Compute lớn hơn} $$

Switch Transformer cho thấy một trục scaling khác:

$$ \text{Model Capacity} \uparrow $$

trong khi cố giữ:

$$ \text{FLOPs per Token} \approx \text{constant} $$

Đây mới là phần có ảnh hưởng lâu dài.

Sau đó, khi nhìn một mô hình MoE có hàng trăm tỷ hay hàng nghìn tỷ parameter, chỉ nhìn total parameters đã không còn đủ.

Ta phải hỏi tiếp:

 1Có bao nhiêu parameter tổng cộng?
 2
 3Mỗi token kích hoạt bao nhiêu expert?
 4
 5Top-1 hay Top-2?
 6
 7Mỗi expert lớn bao nhiêu?
 8
 9Bao nhiêu parameter thực sự active?
10
11Communication giữa các GPU ra sao?
12
13Load balancing thế nào?

Một model 1T parameter chưa chắc cần computation tương đương một dense model 1T.

Hai con quái vật cùng ghi trên trán chữ:

11 TRILLION PARAMETERS

có thể là hai sinh vật hoàn toàn khác nhau.

Một con bắt cả nghìn tỷ parameter đứng dậy làm việc.

Một con sở hữu nghìn tỷ parameter nhưng mỗi token chỉ gọi một nhóm nhỏ.

Nhìn bảng thông số mà không biết điều này rất dễ bị con số parameter lừa.

Switch Transformer không chứng minh rằng càng nhiều parameter càng thông minh

Có một điểm nữa cũng nên nói cho công bằng.

Paper tự ghi nhận rằng upstream pre-training tốt hơn không phải lúc nào cũng chuyển thành downstream performance tốt hơn tương ứng.

Ví dụ, Switch-C 1.6T có một số kết quả fine-tuning không tốt bằng Switch-XXL nhỏ hơn, dù Switch-C sở hữu nhiều unique parameters hơn rất nhiều. Nhóm tác giả chỉ ra rằng Switch-XXL thực hiện khoảng 10 lần FLOPs/token so với Switch-C trong so sánh đó, cho thấy mối quan hệ giữa parameter count, compute và downstream quality phức tạp hơn câu:

“To hơn = thông minh hơn.”

Đây cũng là bài học vẫn còn nguyên giá trị.

Parameter là capacity.

Compute là khả năng sử dụng capacity đó.

Data quyết định model được học thứ gì.

Routing quyết định phần nào của model được gọi.

Architecture quyết định thông tin được biến đổi thế nào.

Training recipe quyết định con quái vật có học được hay nổ tung giữa đường.

Không có một con số duy nhất đại diện cho trí thông minh.

Vì sao Switch Transformer là một cột mốc quan trọng?

Switch Transformer không phát minh Transformer.

Cũng không phát minh Mixture of Experts.

Nó thậm chí không phải công trình đầu tiên đưa MoE vào Transformer.

Nhưng nó làm một việc rất quan trọng:

đơn giản hóa conditional computation đủ nhiều để biến sparse Transformer thành một con đường scaling thực tế.

Từ:

1token → nhiều experts

thành:

1token → một expert

Một cái argmax.

Một cái công tắc.

Một thay đổi tưởng như bé đến mức sinh viên đọc paper có thể hỏi:

“Ủa vậy thôi hả?”

Ừ.

Nhưng khoa học nhiều khi rất kỳ.

Không phải phát minh lớn nào cũng cần thêm mười tầng toán học.

Có những lúc tiến bộ đến từ việc nhìn một cỗ máy đang có quá nhiều bánh răng rồi hỏi:

“Bánh răng này có thật sự cần không?”

Switch Transformer bỏ bớt một phần phức tạp của MoE, giảm routing computation và communication, cải thiện stability, rồi dùng khoảng trống vừa tạo ra để đẩy Transformer tới quy mô hơn một nghìn tỷ parameter. (Journal of Machine Learning Research)

Và từ đó, câu hỏi của ngành AI thay đổi.

Không còn chỉ là:

“Chúng ta có thể xây model lớn tới đâu?”

Mà trở thành:

“Chúng ta có thể xây model lớn tới đâu mà mỗi token không phải đánh thức toàn bộ con quái vật?”

Tổng kết

Nếu phải gói Switch Transformer vào một ý duy nhất thì mình sẽ chọn câu này:

Đừng bắt toàn bộ bộ não suy nghĩ cho mọi câu hỏi. Hãy xây một bộ não rất lớn, rồi chỉ đánh thức phần cần thiết.

Transformer dense mở rộng bằng cách làm phần đang hoạt động lớn hơn.

Switch Transformer mở thêm một con đường khác: làm tổng capacity lớn hơn bằng hàng nghìn experts, nhưng mỗi token chỉ kích hoạt một expert trong mỗi Switch layer.

Nhờ Top-1 routing, load balancing, expert capacity, selective precision và expert parallelism, nhóm Google đã pre-train mô hình khoảng 1,6 nghìn tỷ parameter với 2.048 experts, đồng thời báo cáo speedup lên tới 7x trong một số cấu hình Switch-Base và khoảng 4x so với T5-XXL ở thí nghiệm quy mô lớn tương ứng.

Nhưng di sản lớn nhất của Switch Transformer có lẽ không nằm ở con số 1,6T.

Nó nằm ở một phương trình mới trong đầu chúng ta:

$$ \text{Total Parameters} \neq \text{Active Compute} $$

Một nghìn tỷ parameter có thể nằm im trong bóng tối.

Một token xuất hiện.

Router nhìn nó.

Một expert sáng lên.

Hàng nghìn expert khác ngủ tiếp.

Và cái công tắc nhỏ bé đó đã giúp Transformer bước qua cánh cửa của kỷ nguyên mô hình nghìn tỷ tham số.

Tham khảo

• William Fedus, Barret Zoph, Noam Shazeer — Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, Journal of Machine Learning Research, 2022.
• William Fedus, Barret Zoph, Noam Shazeer — Switch Transformers, arXiv:2101.03961.
• Google Research — Transformer: A Novel Neural Network Architecture for Language Understanding.