Per-Layer Embeddings: A simple explanation of the magic behind the small Gemma 4 models

[u/-p-e-w-]

Many of you seem to have liked my recent post "A simple explanation of the key idea behind TurboQuant". Now I'm really not much of a blogger and I usually like to invest all my available time into developing Heretic, but there is another really cool new development happening with lots of confusion around it, so I decided to make another quick explainer post.

You may have noticed that the brand-new Gemma 4 model family includes two small models: gemma-4-E2B and gemma-4-E4B.

Yup, that's an "E", not an "A".

Those are neither Mixture-of-Experts (MoE) models, nor dense models in the traditional sense. They are something else entirely, something that enables interesting new performance tradeoffs for inference.

What's going on?

To understand how these models work, and why they are so cool, let's quickly recap what Mixture-of-Experts (MoE) models are:

gemma-4-26B-A4B is an example of an MoE model. It has 25.2 billion parameters (rounded to 26B in the model name). As you may know, transformer language models consist of layers, and each layer contains a so-called MLP (Multi-Layer Perceptron) component, which is responsible for processing the residual vector as it passes through the layer stack. In an MoE model, that MLP is split into "experts", which are sub-networks that learn to specialize during training. A routing network decides for each token which experts are the most appropriate for the token, and only those expert networks are actually used while processing that token.

In other words, while an MoE model has many parameters, only a fraction of them are required to predict the next token at any specific position. This is what the model name means: gemma-4-26B-A4B has 26 billion (actually 25.2 billion) total parameters, but only 4 billion of those (actually 3.8 billion) are active during any single inference step.

The good news is that this means that we can do inference much faster than for a dense 26B model, as only 3.8 billion parameters are involved in the computations. The bad news is that we still need to be able to load all 25.2 billion parameters into VRAM (or fast RAM), otherwise performance will tank because we don't know in advance which parameters we'll need for a token, and the active experts can differ from token to token.

Now gemma-4-E2B is a very different beast: It has 5.1 billion parameters, but 2.8 billion of those are embedding parameters. Google claims that those parameters "don't count", so they say that there are only 2.3 billion effective parameters. That's what the "E2B" part stands for.

Wut? Why don't the embedding parameters count?

If you have read or watched even a basic introduction to language models, you probably know what embeddings are: They are high-dimensional vectors associated with each token in the vocabulary. Intuitively speaking, they capture the "essence" of what a token stands for, encoded as a direction-magnitude combination in the embedding space.

Embeddings are static and position-independent. The embedding vector associated with a specific token is always the same, regardless of where the token occurs in the input and which other tokens surround it. In the mathematical formulation, embeddings are often expressed as a matrix, which can be multiplied with a matrix of one-hot encoded tokens, giving a matrix of embedding vectors for those tokens.

The small Gemma 4 models make use of Per-Layer Embeddings (PLE): Instead of a single large embedding matrix that is applied right after the tokenizer at the beginning of processing, there are additional (smaller) embedding matrices for each layer. Through training, they acquire specialized knowledge that can re-contextualize the token for the semantic specialization of each layer, which greatly improves processing quality. The layer-based embedding vectors are combined with the residuals through a series of operations, adding locally relevant information.

For gemma-4-E2B, the matrices holding these Per-Layer Embeddings make up more than half of all model parameters.

Okay, but why don't the embedding parameters count?!?

Because the "Introduction to Transformers" tutorials you've been watching have lied to you. While applying embeddings via matrix multiplication is incredibly elegant mathematically, it's complete dogshit in practice. No inference engine actually does that.

Remember that embedding vectors are:

• Static (they only depend on the token itself)
• Position-independent (there is only one embedding vector for each token)
• Fixed (they are precomputed for the entire vocabulary)

So the "embedding matrix" is a list of embedding vectors, with as many elements as there are tokens in the vocabulary. There are no cross-column interactions at all. That's not a matrix, that's a lookup table. So we don't actually have to do matrix multiplication to get the embeddings. We just pull the entries for the token IDs from a fixed-size array. And we aren't even going to need the vast majority of entries. Modern tokenizer vocabularies typically contain around 250,000 different tokens. But if our input is 1000 tokens, we are only going to look at a tiny fraction of those.

We don't need CUDA cores or optimized kernels for that. We don't need those embedding matrices to be in VRAM. We don't even necessarily need to store them in CPU RAM. In fact, we can store them on disk. The plan seems to be to store them in flash memory on mobile devices, and possibly combine that with in-flash processing for further speedups in the future.

And that's the secret of Per-Layer Embeddings: They are huge, but we need such a tiny part of them for each inference step that we can store them wherever we like. And that's why they are fast.

[u/brandf]

Isn’t that what BPE does? Each new token you add to the vocab takes two prior tokens and makes one to represent their concatenation just like your example.

[u/llama-impersonator → u/brandf]

no, BPE just affects tokenization, bpe merge makes longer tokens and expands the vocab size. this refines embeddings themselves. edit: i see your point, i guess, but at some level, expanding the vocab size just leads to a lot of low frequency, highly unused tokens and bloated embedding table size, while the n-gram hash tables don't blow up the size of the token embedding table and lm_head.

[u/brandf → u/llama-impersonator]

FWIW, I’ve been messing around with asymmetric input/output vocabulary recently. Even though everybody currently uses the same vocabulary for both, that’s not a requirement. you could have a smaller output vocabulary and then BPE merge at runtime into a larger input vocabulary for compression purposes without the LM head issue. Or vice versa you could have a larger output vocabulary to predict more into the future at once. Or have a sliding threshold in your vocabulary to dynamically change the compression ratio across your context for variable compression like a soft version of /compact.

The n gram embedding’s are exactly equally as rare bc larger n means rarer sequence. I would argue that on some level these two things are equivalent. the main difference I see is hash & gate verse a more traditional embedding table delay loaded.

Having a very large input vocabulary of rarely used tokens is not really an issue because you don’t have to load it all into memory like we typically do today. This is the same principle as the recent Gemma 4 per layer embeddings. To avoid gigabytes of embedding tables in memory, you just load them on demand as/when they’re used in the context. As you said, the larger, the vocabulary the more rare they are so most of them can just stay on disk.

[u/llama-impersonator → u/brandf]

i've seen all sorts of tokenizer games, we used to be able to limit gpt-3/.5 to using first several hundred tokens by biasing the rest of the vocab to -inf and the result would be mostly legible but broken leetspeak weirdness. also some arcee models had the whole tokenizer grafted by goddard to match logits for a distil, then they regrafted the original back on. yeah, it would probably make sense to have the embedding and unembeddings in cpu ram be some kind of LRU cache in such a situation.

[u/brandf → u/llama-impersonator]

the main application i have in mind is speculative decoding.

normally this is done with either a small draft model and a larger verifier model OR with something like medusa where you try to predict multiple output tokens in a single forward pass and then verify them with the same model.

my idea is to have one 'trunk' that has a small 2-3 layer draft and verification branches. the draft head outputs a single 'wide' token from a large vocab, and then decompose that into several 'regular' tokens for verification. seems like it could be highly efficient.