Phân tích kiến trúc Gemma 4 31B: mình đọc source code ba ngày - và đây là những thứ Google không in trên slide marketing

TL;DR: Bài này phân tích kiến trúc Gemma 4 31B (dense model mới của Google DeepMind) dựa trên config.json thật từ HuggingFace Hub và source code modeling_gemma4.py trên GitHub. Bốn ý tưởng kỹ thuật đáng chú ý: (1) hybrid attention 5:1 - 50 layer sliding window + 10 layer global trong tổng 60 layer; (2) Proportional RoPE (p-RoPE) - chỉ xoay 25% head dimension với theta = 1M; (3) Per-Layer Embeddings (PLE) - feature mới cho E2B/E4B, inject per-layer signal; (4) shared KV cache - feature có trong code nhưng 31B không bật. Mỗi claim có screenshot evidence từ primary source (config.json, GitHub, arxiv).

Tuần trước Google DeepMind release Gemma 4. Mình lướt qua blog post chính thức (blog.google) và phản ứng đầu tiên của mình là kiểu reflex: thêm một model nữa, lại benchmark, lại "frontier on device". Mình suýt skip.

Rồi mình thấy một con số làm mình dừng lại: 31B dense, 256K context, ranking #3 open model trên Arena. Một model dense 31B mà ăn được mấy con MoE 400B. Chuyện này không bình thường. Mình tải về, đọc model card, đọc HuggingFace blog (huggingface.co/blog/gemma4), rồi mở luôn modeling_gemma4.py trong repo transformers (source code đây) để xem code.

Ba ngày sau, mình ngồi gõ bài này. Không phải để cheerlead cho Google - mình vẫn nghi ngờ vài thứ sẽ nói ở cuối bài. Mà để kể lại ba ý tưởng kỹ thuật mà mình thấy thật sự thú vị, và một issue trên GitHub mà mình nghĩ ai làm inference engine cũng nên đọc.

Bốn variant của Gemma 4: E2B, E4B, 26B A4B MoE và 31B Dense

Gemma 4 không phải "một model". Nó là bốn model rất khác nhau, và sự khác nhau này không chỉ là scale.

Hai chữ viết tắt cần unpack:

• E2B / E4B: chữ "E" là Effective. Tổng params lớn hơn nhiều nhưng nhờ Per-Layer Embeddings, "phần phải nằm trên VRAM lúc inference" tương đương 2B / 4B. Mình sẽ giải thích sau.
• 26B A4B: A là Active. MoE 25.2B tổng, nhưng mỗi token chỉ "chạm" 3.8B params.

Ba kích cỡ đầu (E2B, E4B, 26B A4B) là về chạy được trên thiết bị nhỏ. Riêng con 31B dense là về đẩy ranking lên cao bằng cách tối ưu attention chứ không scale param. Đây là điểm thú vị nhất với mình - và là phần đầu tiên mình muốn kể.

Ý tưởng 1: Hybrid attention 5:1 (sliding window + global) - tỷ lệ kỳ lạ làm nên 256K context cho Gemma 4 31B

Trước khi nhìn config, mình cần nói nhanh về sliding-window attention là gì, vì nếu chưa làm quen với nó thì phần dưới sẽ khó theo.

Sliding-window attention - giải thích trong một phút

Trong attention "bình thường" (gọi là full attention hoặc global attention), khi model xử lý token thứ N, nó cho phép token đó "nhìn lại" tất cả N-1 token trước nó. Mỗi token tự quyết định nó muốn focus vào những token nào trong toàn bộ history. Đây là cách Transformer gốc hoạt động, và nó là lý do Transformer mạnh - không có thông tin nào bị giới hạn khoảng cách.

Cái giá phải trả: chi phí scale bậc hai theo độ dài context. Để xử lý 100K token với full attention, model phải tính 100K × 100K ≈ 10 tỷ cặp tương quan cho mỗi layer, mỗi attention head. Và với mỗi token mới được generate, model phải lưu lại Key/Value của các token cũ (KV cache, mình sẽ giải thích ngay bên dưới) để không tính lại.

Sliding-window attention đảo ngược trade-off: thay vì cho token nhìn về toàn bộ quá khứ, ta giới hạn nó chỉ được nhìn lại W token gần nhất. W gọi là window size. Với Gemma 4, W = 1024 cho các model lớn.

Hãy tưởng tượng bạn đang đọc một quyển sách. Full attention = bạn có quyền lật về bất kỳ trang nào trong cả quyển bất cứ lúc nào. Sliding window 1024 = bạn chỉ được phép nhìn 1024 từ gần nhất; mọi thứ xa hơn 1024 từ về trước, layer này không thấy. Đối với token thứ 5000, layer sliding-window chỉ thấy được token 3977 đến 4999.

Hai hệ quả quan trọng:

1. Chi phí giảm từ bậc hai xuống tuyến tính. Cho mỗi token mới, layer chỉ tính tương quan với 1024 token trước, không phải toàn bộ history. KV cache cũng chỉ cần lưu 1024 entry cho mỗi position thay vì cả history.
2. Mất thông tin xa. Nếu một fact quan trọng nằm ở token 100 và bạn đang generate ở token 5000, sliding-window layer không thấy nó được. Đây là cái giá thật sự của sliding window - và là lý do Gemma không dùng sliding window cho mọi layer.

Một câu hỏi tự nhiên: nếu sliding-window mất thông tin xa, làm sao model 31B của Gemma 4 vẫn handle được context 256K? Câu trả lời chính là tỷ lệ 5:1 mà mình sắp giải thích - và đó là phần thông minh nhất.

Tỷ lệ 5:1 trong config

Khi nhìn vào config của Gemma 4 31B, có một con số đập vào mắt: 60 layer, trong đó 50 layer là sliding-window attention với window 1024 token, và chỉ 10 layer là full global attention. Tỷ lệ chính xác là 5:1 (phân tích chi tiết trên Substack của Kaitchup).

Để hiểu tại sao tỷ lệ này có ý nghĩa, mình cần kể tiếp về KV cache - kẻ thù số một của long context.

KV cache: tại sao 10 layer global là đắt

Khi generate token mới, model không tính lại attention từ đầu. Nó cache hai tensor - Key và Value - của mỗi token đã xử lý, cho mỗi layer, mỗi attention head. Đây là KV cache, và nó là phần ngốn memory thật sự lúc inference, không phải weight model.

Với full attention, KV cache phải lưu K/V cho tất cả token trong context. Với 256K token, một model 30B có thể cần 60-80GB chỉ riêng KV cache - nhiều hơn cả weight model. Với sliding window 1024, mỗi position chỉ cần 1024 entry → giảm hai bậc.

Bây giờ bạn đã hiểu đủ để thấy tỷ lệ 5:1 là một quyết định toán học cụ thể: 50 layer rẻ + 10 layer đắt, thay vì 60 layer đắt. KV cache giảm xấp xỉ ~83% so với full-attention-everywhere, mà vẫn giữ được cầu nối thông tin xa qua 10 layer global rải đều.

Giải pháp Gemma: hybrid attention với tỷ lệ 5:1

Đây là chỗ mình muốn show evidence trực tiếp. Mình đã pull config.json từ chính repo google/gemma-4-31B-it trên HuggingFace và highlight các field quan trọng:

Nhìn vào array layer_types: đúng pattern sliding × 5, full × 1 lặp lại 10 lần, layer thứ 60 (cuối cùng) là full_attention. Đây không phải mình đoán - đây là enum literal trong file config production của model. Gemma 4 không chọn một hay chọn hai. Nó xen kẽ hai loại layer:

Layer 1: sliding (1024)
Layer 2: sliding (1024)
Layer 3: sliding (1024)
Layer 4: sliding (1024)
Layer 5: sliding (1024)
Layer 6: GLOBAL (full 256K)
Layer 7: sliding (1024)
...
Layer 60: GLOBAL ← layer cuối luôn là global

Với 60 layer, bạn có 50 sliding + 10 global. Chuyện này có nghĩa là gì?

• Hầu hết KV cache là local. 50 trong 60 layer chỉ phải lưu 1024 entry cho mỗi attention head. KV cache giảm khủng khiếp so với full attention everywhere.
• Vẫn có 10 layer "global" rải đều, đảm bảo thông tin từ xa vẫn flow được qua model. Mỗi 6 layer, model có một cơ hội để "đọc lại" toàn bộ context.
• Layer cuối luôn là global, để output cuối được informed bởi full context.

Tỷ lệ 5:1 không phải con số magic. Gemma 3 đã document chính xác pattern này trong technical report (arxiv 2503.19786) - và Gemma 4 kế thừa nó nguyên vẹn:

Gemma 4 giữ nó vì nó thật sự work. Gemini 2.5 Flash, 3 Pro - toàn bộ family Gemini dùng kiến trúc tương tự. Đây không phải research toy.

5:1 chính là tỉ lệ 80:20 trong truyền thuyết. Mình nghĩ họ cũng "bốc thuốc" ra cái tỷ lệ này rồi traning ra kết quả tốt thì giữ lại.

Có một chi tiết nhỏ trong config mà mình muốn note: num_global_key_value_heads: 4 và global_head_dim: 512, trong khi sliding layer dùng num_key_value_heads: 16 và head_dim: 256. Global layer được "mỏng và rộng", local layer "dày và hẹp". Mình chưa thấy tài liệu giải thích lý do, nhưng đoán: global layer cần ít head hơn vì nó đã có toàn bộ context, không cần đa dạng "view"; local layer cần nhiều head để compensate cho việc thiếu information.

Ý tưởng 2: Proportional RoPE (p-RoPE) - Gemma 4 chỉ xoay 25% head dimension và lý do nó work

Đây là phần mà nếu bạn đã từng thử extend context window cho một model open-source nào đó, bạn sẽ thấy quen thuộc và đau đầu.

RoPE là gì nhỉ?

Để hiểu RoPE, mình cần kéo bạn lùi lại một bước và nói về vấn đề nó cố giải quyết.

Vấn đề: Transformer mặc định không biết thứ tự token.

Câu "Mèo cắn chó" và "Chó cắn mèo" có cùng ba từ. Nếu bạn shuffle thứ tự, attention layer truyền thống sẽ tính ra kết quả y hệt. Để thấy tại sao, hãy nhìn công thức attention score giữa token vị trí $i$ (query) và token vị trí $j$ (key):

$$
\text{score}(i, j) = \frac{\mathbf{q}_i^\top \mathbf{k}_j}{\sqrt{d}}, \quad \text{với } \mathbf{q}_i = W_Q\,\mathbf{x}_i, \quad \mathbf{k}_j = W_K\,\mathbf{x}_j
$$

Để ý: công thức này chỉ phụ thuộc vào hai vector embedding $\mathbf{x}_i$ và $\mathbf{x}_j$, hoàn toàn không có $i$ hay $j$ xuất hiện ở vế phải. $W_Q$ và $W_K$ là các ma trận shared chung cho mọi vị trí. Điều này có hệ quả nghiêm trọng: nếu mình swap "mèo" và "chó" trong câu, các vector $\mathbf{x}$ chỉ đổi chỗ - toàn bộ tập dot product giữa chúng không thay đổi gì cả. Attention không có thông tin nào để phân biệt "mèo ở vị trí 1" với "mèo ở vị trí 3".

Cần phải có một cách nào đó tiêm thông tin vị trí vào query/key.

Cách cũ (BERT, GPT-2): tạo một bảng "positional embedding" - vector vị trí 1, vector vị trí 2, vector vị trí 3... - rồi cộng vào input embedding. Vị trí được "ghi đè" lên token. Cách này hoạt động nhưng có 2 vấn đề: (1) muốn extend context dài hơn training thì cần thêm bảng vị trí mới chưa được train, (2) model phải tự học cách disentangle thông tin vị trí khỏi thông tin token - không tự nhiên.

RoPE đảo logic này. Thay vì cộng một vector vị trí vào, RoPE xoay vector query/key đi một góc tỷ lệ với vị trí của token đó. Token ở vị trí 1 → xoay 1 đơn vị góc. Token ở vị trí 2 → xoay 2 đơn vị góc. Token ở vị trí 100 → xoay 100 đơn vị góc.

Phép "xoay" chính xác là gì?

Vì rotation chỉ định nghĩa được trong không gian 2D, RoPE chia head dimension thành các cặp 2D rồi xoay từng cặp độc lập. Với một cặp 2D, ma trận xoay góc $\theta$ là:

$$
R(\theta) = \begin{pmatrix} \cos\theta & -\sin\theta \ \sin\theta & \cos\theta \end{pmatrix}
$$

Đây là ma trận xoay vector ngược chiều kim đồng hồ một góc $\theta$. Áp dụng vào RoPE: với token ở vị trí $m$ và một tần số $\omega$ cố định cho mỗi cặp dimension, ta xoay query và key đi một góc tỷ lệ với vị trí:

$$
\mathbf{q}'_m = R(m\omega)\,\mathbf{q}_m, \qquad \mathbf{k}'_n = R(n\omega)\,\mathbf{k}_n
$$

Nghĩa là: query của token vị trí $m$ bị xoay góc $m\omega$, key của token vị trí $n$ bị xoay góc $n\omega$. Sau đó dot product giữa chúng - vốn là attention score - trở thành:

$$
\mathbf{q}'^{\top}_m \mathbf{k}'_n \;=\; (R(m\omega)\mathbf{q}_m)^{\top} R(n\omega)\mathbf{k}_n \;=\; \mathbf{q}_m^{\top}\, R(m\omega)^{\top} R(n\omega)\, \mathbf{k}_n
$$

Đây là chỗ phép màu xảy ra. Ma trận xoay có một property đẹp: $R(\alpha)^{\top} R(\beta) = R(\beta - \alpha)$ - tích của hai phép xoay là một phép xoay duy nhất bằng hiệu góc. Áp dụng vào trên:

$$
\boxed{\;\mathbf{q}'^{\top}_m \mathbf{k}'_n \;=\; \mathbf{q}_m^{\top}\, R\big((n-m)\omega\big)\, \mathbf{k}_n\;}
$$

Nhìn vế phải: kết quả chỉ phụ thuộc vào $(n - m)$ - tức là khoảng cách relative giữa hai token, không phải vị trí absolute của chúng. Hai token cách nhau 5 vị trí sẽ luôn cho cùng một dot product modifier, bất kể chúng nằm ở đầu hay cuối câu. Bạn encode position absolute (xoay theo index $m$, $n$) ở input, và automatically thu được position relative ($n - m$) ở output - không phải vì model học được, mà vì hình học của ma trận xoay làm sẵn cho bạn.

Cụ thể hoá bằng hai cặp token có cùng khoảng cách:

• Cặp 1: query ở vị trí 5, key ở vị trí 8 → score phụ thuộc vào hiệu góc xoay $3\omega$
• Cặp 2: query ở vị trí 1005, key ở vị trí 1008 → score cũng phụ thuộc vào $3\omega$
• Hai cặp này có cùng attention score, đúng như mong đợi: model nên xử lý khoảng cách 3 token y như nhau bất kể chúng nằm ở đầu hay cuối câu. Trong hình trên, dù hai mũi tên ở góc absolute khác nhau hoàn toàn, cái cung "gap = $3\omega$" giữa chúng giống hệt nhau - và đó chính là cái duy nhất attention thấy.

Còn về tần số $\omega$: RoPE không dùng một $\omega$ duy nhất. Một head 256-dim được chia thành 128 cặp 2D, mỗi cặp dùng một $\omega$ khác nhau. Cụ thể:

$$
\omega_k = \theta^{-2k/d}, \quad k = 0, 1, \dots, d/2 - 1
$$

với $\theta$ là base frequency (mặc định 10000). Cặp đầu xoay rất nhanh, cặp sau xoay rất chậm - y như đồng hồ có kim giây, kim phút, kim giờ. Cặp xoay nhanh phân biệt token gần nhau; cặp xoay chậm phân biệt token xa nhau. Đây là lý do RoPE encode được nhiều scale khoảng cách trong cùng một head.

Bài toán extrapolation

Cái đẹp của RoPE đi kèm cái dở: các góc xoay model thấy trong training là một range cố định.

Giả sử bạn train ở context 8K. Token vị trí 8000 sẽ có Q/K bị xoay 8000 đơn vị góc - tương đương khoảng 22 vòng xoay đầy đủ. Model học được "các pattern attention sống tốt khi vector bị xoay từ 0 tới ~22 vòng". Bây giờ deploy với context 32K. Token vị trí 32000 bị xoay ~88 vòng - model chưa bao giờ thấy pattern này. Quality giảm thê thảm.

Đây là lý do mọi research extend context window đều phải làm gì đó với RoPE: linear scaling (chia position cho factor để giữ range), NTK-aware scaling (điều chỉnh theta), YaRN (kết hợp nhiều technique), Position Interpolation... Toàn bộ literature quanh RoPE 3 năm qua xoay quanh câu hỏi: làm sao để model handle các góc xoay nó chưa từng thấy?

Gemma 4: chỉ xoay 25% dimension, theta = 1M

Gemma 4 không dùng linear scaling như Gemma 3. Nó chọn một cách khác mà trong code gọi là Proportional RoPE (p-RoPE): chỉ apply RoPE cho 25% đầu của head dimension, để 75% còn lại không bị xoay. Theta = 1,000,000 (so với 10K mặc định).

Đây không phải con số mình đoán. Mở thẳng config.json của google/gemma-4-31B-it, kéo xuống block rope_parameters:

"rope_parameters": {
"full_attention": {
"partial_rotary_factor": 0.25,
"rope_theta": 1000000.0,
"rope_type": "proportional"
},
"sliding_attention": {
"rope_theta": 10000.0,
"rope_type": "default"
}
}

Ba dòng quan trọng:
- partial_rotary_factor: 0.25 → chỉ 25% head dimension được xoay (với head_dim 256 = 64 dimension đầu)
- rope_theta: 1000000.0 → base frequency $\theta = 10^6$, lớn gấp 100 lần default 10K
- rope_type: "proportional" → áp dụng Gemma4RotaryEmbedding với rope_type = "proportional", có thể trace được trong modeling_gemma4.py và ROPE_INIT_FUNCTIONS trong transformers

Và để ý thêm: hai loại layer dùng hai cấu hình RoPE khác nhau. full_attention (10 layer global) dùng p-RoPE 25% với theta 1M; sliding_attention (50 layer local) dùng RoPE default - full dimension, theta 10K. Đây là một detail mà không có blog HuggingFace nào nói trực tiếp - chỉ có config nói rõ ràng.

Tại sao chỉ xoay một phần?

Để trả lời câu này, mình lùi về paper "Round and Round We Go! What makes Rotary Positional Encodings useful?" của Barbero et al., được công bố tại ICLR 2025 (arxiv 2410.06205). Bốn trong năm tác giả đến từ Google DeepMind - chính team Gemma. Họ làm một thứ rất cụ thể: mở Gemma 7B ra, đo xem các tần số RoPE khác nhau thực sự đang được model dùng để làm gì. Kết quả counter-intuitive:

Cái tên paper =))

Hai dòng được highlight nói rõ: Gemma học cách dùng highest frequencies để xây các pattern attention "positional" (đếm khoảng cách), nhưng lại prefer dùng lowest frequencies - và các tác giả nghi ngờ chúng được dùng để carry semantic information, không phải position. Tức là không phải mọi cặp dimension trong RoPE đều đang làm việc về vị trí - phần lớn cặp tần số thấp đang làm semantic memory thuần tuý.

Khi bạn force RoPE lên toàn bộ head dimension (như RoPE gốc), bạn đang xoay cả các cặp đang làm semantic - gây nhiễu chúng một cách không cần thiết. P-RoPE chỉ xoay 25% đầu (các cặp tần số cao) và để 75% còn lại - tức các cặp tần số thấp - yên một chỗ làm việc semantic của chúng. Đây không phải hand-wave, đây là design choice trực tiếp suy ra từ kết quả empirical trong paper Barbero.

Bằng cách chỉ xoay 64 trong 256 dimension đầu (25%), Gemma giữ được benefit của RoPE cho long-range modeling, nhưng để 192 dimension còn lại làm "semantic memory" thuần túy. Theta cao (1M) đảm bảo rằng ngay cả 25% được xoay cũng không lặp lại trong khoảng 256K token.

Nó áp dụng chỉ cho global layer. Sliding window layer dùng RoPE chuẩn (full dimension, theta nhỏ hơn) vì context của chúng cap tại 1024 - không có vấn đề extrapolation.

Tới đoạn này, có một câu hỏi mà mình đã ngồi với khá lâu: tại sao không ai làm thế này sớm hơn? Câu trả lời mình đoán là vì hầu hết mọi người tin rằng RoPE phải được apply uniformly mới hoạt động đúng - đó là cách paper gốc viết. Phải có ai đó đủ tự tin để break assumption đó. Gemma team có vẻ là một trong số không nhiều team làm chuyện này ở scale frontier.

Ý tưởng 3: Per-Layer Embeddings (PLE) - feature mới của Gemma 4 E2B/E4B và lý do nó là phần khó debug nhất

Đây là phần mà mình thấy thú vị nhất, và cũng là phần ít được giải thích nhất trong tất cả tài liệu Google publish.

Vấn đề: embedding layer rất tốn memory cho model nhỏ

Khi bạn build một model nhỏ (E2B, E4B), một sự thật khó chịu hiện ra: embedding table chiếm tỷ trọng cực lớn trong tổng param. Với vocab 262K (Gemma dùng vocab này từ Gemma 1) và hidden size, ví dụ, 2048, embedding table = 262K × 2048 ≈ 537M param. Cho một model "2B", riêng input + output embedding đã ăn trên 1B param.

Nhiều model giải quyết bằng cách tie input và output embedding (share weight). Gemma làm điều đó. Nhưng vẫn còn vấn đề: hầu hết "intelligence" của embedding bị "frontload" - toàn bộ thông tin về token được dồn vào một vector ở đầu, rồi từ đó chảy qua 35 layer. Mỗi layer phải tự suy luận xem trong vector này có thông tin gì còn relevant.

Per-Layer Embeddings (PLE) đảo ngược logic này.

PLE: bơm thông tin token vào mỗi layer, một cách nhỏ giọt

Ý tưởng: ngoài embedding chính, ta có một bảng embedding phụ thứ hai, tạo ra một vector nhỏ riêng cho mỗi layer, dựa trên token ID. Mỗi layer khi forward sẽ nhận thêm signal từ vector này, cộng thẳng vào residual stream.

Code trong modeling_gemma4.py (class Gemma4TextModel) làm hai bước:

1. get_per_layer_inputs(input_ids, inputs_embeds) - lookup từ bảng PLE, ra tensor shape [B, S, num_layers, ple_dim]. Mỗi token có num_layers vector nhỏ, mỗi vector dành cho một layer cụ thể.
2. project_per_layer_inputs(inputs_embeds, per_layer_inputs) - combine với một context-aware projection từ embedding chính:
   python context_proj = per_layer_model_projection(inputs_embeds) / sqrt(hidden_size) context_proj = RMSNorm(reshape(context_proj, [B, S, num_layers, ple_dim])) final = (context_proj + per_layer_inputs) / sqrt(2)

Layer per_layer_model_projection được khai báo ở line 1533, với scale factor per_layer_model_projection_scale = hidden_size ** -0.5 tức $1/\sqrt{d_{\text{hidden}}}$.

Hai signal được combine: một signal token-identity thuần (lookup từ bảng PLE) và một signal context-aware (projection từ main embedding, chứa thông tin sentence). Hệ số $\tfrac{1}{\sqrt{2}}$ là để variance không nổ khi cộng hai signal độc lập có cùng variance - tổng của hai biến ngẫu nhiên độc lập có variance $\sigma^2$ sẽ có variance $2\sigma^2$, chia cho $\sqrt{2}$ kéo về lại $\sigma^2$.

Confirm thêm: trong config của 31B mình đã pull, field hidden_size_per_layer_input: 0 - tức 31B Dense không dùng PLE. Đây thuần là feature cho E2B/E4B, đúng như HuggingFace blog mô tả.

Vì mỗi vector PLE rất nhỏ (ple_dim thường khoảng 256), tổng size của bảng PLE = vocab × num_layers × ple_dim - vẫn đáng kể, nhưng bảng này không cần nằm trên VRAM hot. Nó có thể stream từ disk hoặc CPU memory, vì lookup là sparse và predictable. Đây chính là lý do "effective param" khác "total param": tổng 5.1B nhưng phần thực sự cần trên GPU lúc forward chỉ 2.3B.

Vì sao đây là phần khó debug

Nếu bạn implement Gemma 4 inference từ đầu, đây là chỗ bạn sẽ stuck. Có một issue trên GitHub mình nghĩ ai đụng vào model này nên đọc: issue #45206 - "Gemma4: PLE implementation is underdocumented and config is misleading". Người mở issue chỉ ra rằng:

• Tên field trong config (per_layer_model_projection, ple_dim, num_kv_shared_layers) không có docstring giải thích.
• Hai scale factor 1/sqrt(hidden_size) và 1/sqrt(2) không có comment giải thích lý do.
• Pipeline tổng thể chỉ có thể hiểu được bằng cách đọc forward và trace từng tensor.

Mình đã thử reimplement PLE để chắc mình hiểu. Mất khoảng 4 giờ với output sai một cách subtle - không crash, chỉ là logits hơi lệch. Cuối cùng mình tìm ra: mình đã quên RMSNorm ở giữa. Nếu bạn chỉ đọc blog post của Google, bạn sẽ không biết cái RMSNorm đó tồn tại.

Đây là điểm mình muốn nói thẳng: Google đang publish weight và code, nhưng không publish technical report chính thức cho Gemma 4. Khác với Gemma 3 (có paper trên arxiv), Gemma 4 ra mắt ngày 2/4/2026 và đến hôm nay (13/4) vẫn chưa có paper. Các chi tiết như tỷ lệ 5:1, p-RoPE 25%, PLE pipeline - toàn bộ là reverse-engineered từ code và blog post. Đây không phải open-source theo nghĩa đầy đủ. Đây là open-weight với obscured methodology.

Ý tưởng 4: Shared KV cache (num_kv_shared_layers) - feature có trong code Gemma 4, nhưng 31B không bật

Đây là một surprise mà mình chỉ phát hiện khi đã đọc config thật, không phải blog post.

HuggingFace blog viết về Gemma 4 nói: "Last N layers reuse key/value tensors from earlier layers" - nghe như đây là một feature chung của Gemma 4. Trong code modeling_gemma4.py có tham số num_kv_shared_layers thật, và logic forward có rẽ nhánh để N layer cuối skip KV projection và reuse K/V từ layer "không-shared" gần nhất cùng loại attention (sliding hoặc full).

Nhưng: trong config.json của google/gemma-4-31B-it (xem screenshot ở phần trên), giá trị là:

"num_kv_shared_layers": 0

Tức là model 31B Dense không kích hoạt feature này. Code support nó, nhưng config cụ thể của 31B set N = 0. Có thể các variant khác (E2B/E4B/26B A4B) bật nó với giá trị khác - mình chưa pull config của chúng để confirm, sẽ làm trong bài follow-up.

Bài học mình rút ra cho chính mình: đừng generalize từ blog post sang config. Một feature được mô tả trong release blog có thể chỉ áp dụng cho một subset variant. Nếu mình dừng ở blog HuggingFace mà không pull config thật, mình đã viết sai trong bài này. Đây cũng là lý do mình đặt tên section là "feature có trong code, nhưng 31B không bật" - để bạn đọc thấy luôn cái dây giữa "feature exists" và "feature is on for the model you're benchmarking".

Câu hỏi mở: tại sao 31B không bật KV sharing trong khi nhỏ hơn (E2B/E4B) có thể có? Đoán: với 31B, mỗi layer K/V đắt giá (16 head, head_dim 256) và quality regression có thể nhiều hơn benefit. Với model nhỏ chạy on-device, ép memory mới là priority. Nhưng đây vẫn chỉ là đoán - chỉ có khi đọc paper (nếu Google có publish) mới biết.

Vision encoder: phần ít được nói đến

Tất cả model Gemma 4 đều multimodal. Vision encoder của 31B/26B có ~550M param (E2B/E4B có ~150M). Nó hỗ trợ:

• Variable aspect ratio - không phải resize về 224×224 như cũ.
• Configurable image token budget: 70, 140, 280, 560, hoặc 1120 token. Bạn chọn trade-off speed/quality.
• Learned 2D position kết hợp với multi-dimensional RoPE - tức là RoPE không chỉ encode "vị trí trong sequence" mà còn encode "vị trí trên grid hai chiều".

Trong benchmark MMMU Pro, 31B đạt 76.9%, vượt mọi open model khác cùng size. Cá nhân mình thì cẩn thận với benchmark visual reasoning - chúng dễ bị contaminate hơn text - nhưng số đó không phải nhỏ.

E2B và E4B còn có thêm audio encoder (USM-style conformer, kế thừa từ Gemma 3n) cho speech recognition và understanding, max 30 giây audio. Mình chưa test nhưng sẽ là chủ đề một bài khác.

Đo memory thực tế: 31B chạy trên một 3090?

Vì mình đang chỉ thử nghiệm Gemma trên Colab và API (1 click) nên chưa có self host thử. Vì vậy xin mượn phân tích của Kaitchup (link) ở context 262K, bf16:

Số quan trọng nhất: 31B dense, full 256K context, AWQ 4-bit quantization, vừa 24GB VRAM. Tức là một con RTX 3090 hoặc 4090 single GPU chạy được. Cách đây 18 tháng, "31B param + 256K context trên một consumer GPU" là bốc phét.

KV cache cho 31B ở 262K context chiếm khoảng 20.78 GiB (theo phân tích Kaitchup, dựa trên config thật). Cái 5:1 attention pattern là lý do duy nhất con số này không nổ thành 80GB - và như mình đã note ở section trên, 31B không dùng shared KV cache để giảm thêm. Nếu bạn vẫn tin "context length không quan trọng vì RAG giải quyết hết", thì hãy nghĩ lại - vì bây giờ context length rẻ tới mức bạn có thể bỏ một codebase cỡ trung vào prompt.

Vài thứ mình vẫn chưa tin

Để bài này không thành cheerleading, mình muốn ghi lại ba điều mình còn skeptical:

1. Arena Elo có thể bị thao túng bằng cách chiều ý người chấm. Ranking #3 open model cho 31B nghe rất hay, nhưng Arena là human preference - người dùng nhìn hai câu trả lời cạnh nhau rồi click cái "thấy thích hơn". Nó đo "model nào nghe tự tin, format đẹp, có bullet rõ ràng, mở bài lịch sự" gần như nhiều bằng "model nào trả lời đúng". Một model được fine-tune để chiều cảm tình người chấm - dùng nhiều markdown, hedge cẩn thận, agree với assumption ẩn của câu hỏi - có thể leo Elo cao mà không thực sự mạnh hơn về reasoning. Mình muốn xem private eval của các team độc lập, đặc biệt khi Gemma sẽ phải đối đầu Llama 4 và Qwen 3 trong vài tháng tới.
2. PLE và effective param có thể đang oversell efficiency. "5.1B total, 2.3B effective" nghe như miễn phí, nhưng bảng PLE phải nằm đâu đó - và nếu bạn không có infrastructure để stream embedding từ CPU/disk hiệu quả, "effective param" không cứu bạn. Trên một con phone Android low-end thật sự, mình chưa thấy benchmark độc lập đáng tin.
3. Không có technical report chính thức. Mình đã nói ở trên nhưng đáng nhắc lại: trong một tuần kể từ release, vẫn không có paper. So với Llama, Mistral, Qwen - đều có technical report kèm release - đây là điểm lùi đáng kể về tính reproducibility.

Mình mang gì về sau ba ngày này

Ba ý tưởng kỹ thuật mà Gemma 4 đóng đinh trong đầu mình:

• Hybrid attention với tỷ lệ 5:1 không phải compromise - nó là design choice đúng cho long context. Nếu bạn đang build inference engine, đây là pattern bạn nên optimize cho.
• RoPE không cần áp dụng uniformly. P-RoPE 25% là một idea đơn giản đến mức mình ngạc nhiên không ai làm sớm hơn. Mình đoán sẽ thấy nó xuất hiện trong nhiều open model khác trong 6 tháng tới.
• PLE là một hướng tổ chức embedding mới mà mình muốn thử trong một project nhỏ. Conceptually, nó tách bạch "token identity" và "token-in-context" thành hai pathway song song - và đó là một mental model đẹp về cách model có thể save memory.

Còn câu hỏi mình mang theo trong vài tuần tới: liệu phần khó hiểu của PLE có phải là tính năng, hay là bug có tổ chức? Một thiết kế cần 4 giờ để reverse-engineer cho một người đọc đã quen với transformer code thì có thể là dấu hiệu của insight sâu, hoặc dấu hiệu của technical debt được publish ra ngoài. Mình chưa quyết được.

Nếu bạn đã đọc tới đây - cảm ơn bạn. Nếu bạn nghĩ mình hiểu sai chỗ nào, đặc biệt phần PLE và shared KV cache, mình rất muốn được bạn chỉ cho. Mình đã có thói quen là reverse-engineer nó từ code, chứ không hoàn toàn tin vào paper, nên nếu có sai khác với paper chắc là code implement láo đấy =))

Nửa li latte chắp bút cho những dòng cuối cùng. Cảm ơn bạn đã đọc đến đây - mình biết đây là một bài dài và kỹ thuật và khá tiêu tốn não để đọc nhỉ.

Bình.

FAQ: các câu hỏi phổ biến về Gemma 4

Gemma 4 là gì?

Gemma 4 là thế hệ mới nhất của dòng model mở Google DeepMind, release ngày 2 tháng 4, 2026. Nó kế thừa research từ Gemini 3, gồm bốn variant: E2B (~2.3B effective params), E4B (~4.5B effective params), 26B A4B (MoE 25.2B tổng, 3.8B active), và 31B Dense (30.7B params). Tất cả đều multimodal (text + image; E2B/E4B còn có audio).

Gemma 4 31B khác gì Gemma 3 27B?

Gemma 4 31B dùng cùng hybrid attention 5:1 như Gemma 3, nhưng thay đổi RoPE: chuyển từ linear scaling sang Proportional RoPE (p-RoPE) với 25% head dimension và theta = 1M. Context window tăng từ 128K lên 256K. Vocab giữ nguyên 262K.

Gemma 4 31B chạy được trên GPU nào?

Với AWQ 4-bit quantization, Gemma 4 31B full 256K context vừa 24GB VRAM - tức là một RTX 3090 hoặc 4090 consumer GPU chạy được. Unquantized bf16 cần ~61GB. KV cache chiếm ~21GB ở 256K context, nhờ hybrid attention 5:1.

Per-Layer Embeddings (PLE) trong Gemma 4 hoạt động thế nào?

PLE là feature của E2B/E4B (không phải 31B). Ngoài bảng embedding chính, model có một bảng embedding phụ sinh ra vector nhỏ riêng cho mỗi layer dựa trên token ID. Vector này được cộng vào residual stream của từng layer, inject thông tin token-specific dọc theo depth. Bảng PLE không cần nằm trên VRAM hot → đây là lý do "effective params" nhỏ hơn "total params".

Proportional RoPE (p-RoPE) là gì?

P-RoPE chỉ áp dụng rotation vị trí lên 25% đầu của head dimension thay vì toàn bộ. Dựa trên paper Round and Round We Go! (Barbero et al., ICLR 2025) - paper tìm ra rằng Gemma học cách dùng highest frequencies cho position, lowest frequencies cho semantic. P-RoPE giữ các cặp tần số cao làm position, để 75% còn lại (tần số thấp) yên cho công việc semantic.

Gemma 4 có dùng Mixture of Experts (MoE) không?

Có - một variant. Gemma 4 26B A4B là MoE với 128 experts, 8 active + 1 shared expert cho mỗi token, total 25.2B / active 3.8B. Ba variant còn lại (E2B, E4B, 31B) đều dense.

Source code Gemma 4 ở đâu?

Implementation chính thức nằm trong HuggingFace transformers: src/transformers/models/gemma4/modeling_gemma4.py. Config cụ thể cho từng variant có trong repo HuggingFace Hub của model, ví dụ google/gemma-4-31B-it/config.json.

Gemma 4 có technical report không?

Tính đến ngày viết bài này (13/04/2026), chưa có technical report chính thức trên arxiv - khác với Gemma 3 (arxiv 2503.19786). Mọi chi tiết kiến trúc mình nói trong bài đều reverse-engineered từ config.json, source code, và HuggingFace blog post.

Nguồn tham khảo chính:

• HuggingFace blog: Welcome Gemma 4
• Source code: transformers/models/gemma4/modeling_gemma4.py
• Issue #45206: PLE underdocumented
• Kaitchup: Gemma 4 architecture & memory analysis
• Google AI: Gemma 4 model card
• Gemma 3 Technical Report (arxiv 2503.19786) - tham khảo vì Gemma 4 chưa có paper chính thức
• Barbero et al. - Round and Round We Go! What makes Rotary Positional Encodings useful? ICLR 2025 (arxiv 2410.06205) - justification cho p-RoPE 25%
• google/gemma-4-31B-it/config.json - config thật, primary source
• Google DeepMind: Gemma 4 page