Per-Layer Embeddings and Parameter Offloading
Why a model's parameter count stopped predicting its accelerator memory footprint, how per-layer embeddings move a large block of weights to cheap memory that only needs a lookup, and what the PCIe or unified-memory path costs when the trick is applied to the wrong tensors.
Gemma 3n ships two models named E2B and E4B. The E stands for effective: the raw parameter counts are roughly 5B and 8B, but both run in about 2 GB and 3 GB of accelerator memory (Google, 2025, Introducing Gemma 3n: the developer guide). The naming looks like marketing and is actually a statement about where the weights live.
The mechanism is Per-Layer Embeddings. A large block of embedding parameters, one set per layer rather than one shared input table, is held in ordinary system memory and computed on the CPU, and only the core transformer weights are resident on the accelerator. Accuracy comes from the full parameter count; footprint comes from the resident subset.
Why embeddings are the right thing to move
Not all parameters are equally expensive to keep off the accelerator. The test is arithmetic intensity: how much compute a tensor participates in per byte you have to move.
A weight matrix in an FFN is multiplied by every token's activation vector. Move it off-chip and you pay the transfer on every forward pass, for every token, and the accelerator stalls waiting. An embedding table is the opposite: for a given token you touch exactly one row. The operation is a gather, the bytes moved are one row rather than the table, and a CPU with normal DRAM bandwidth does it perfectly well.
So the rule for offloading is not "move the biggest tensors", it is "move the tensors with the lowest compute-per-byte, whose access pattern is sparse and predictable". Embedding tables satisfy both conditions. Attention weights satisfy neither.
ALBERT attacked the same block from a different angle in 2019, factorising the \(V \times H\) embedding matrix into \(V \times E\) and \(E \times H\) so that vocabulary size stops dictating hidden size (Lan et al., 2019, ALBERT, arXiv:1909.11942). Factorisation shrinks the parameters; offloading keeps them and relocates them. On a device where system RAM is plentiful and accelerator memory is not, relocating is the better trade.
Capacity you are not paying for in VRAM
Gemma 3n pairs PLE with two other capacity-without-latency techniques, and they are worth naming because they generalise. AltUp widens the token representation while only updating a sub-block of it at each layer, using a predict-and-correct step to refresh the rest, so capacity grows without a proportional latency increase (Baykal et al., 2023, Alternating Updates for Efficient Transformers, NeurIPS 2023, arXiv:2301.13310). LAuReL generalises the residual connection with a small learned combination of previous layer outputs, reporting downstream gains of 2.54% to 20.05% on 1B and 4B pretraining while adding 0.012% and 0.1% more parameters (Menghani et al., 2024, LAuReL: Learned Augmented Residual Layer, arXiv:2411.07501).
The common shape: find a place where the model is under-parameterised relative to its compute budget, and add parameters that do not sit on the critical path of every matmul.
When it breaks
The bus becomes the model's slowest layer. On a discrete GPU across PCIe, a per-token gather from host memory adds a round trip that a 3 GB VRAM saving rarely justifies. The technique is designed for unified-memory devices, phones and laptops where CPU and accelerator share physical RAM, and it degrades badly off that target.
Effective parameters are not comparable parameters. An "E4B" does not benchmark like a 4B and does not cost like an 8B. Quality tracks total parameters, memory tracks resident parameters, and any procurement comparison that uses one number for both will be wrong in one direction or the other. Say which one you mean.
Tooling assumes one weight home. Quantisation scripts, LoRA adapters, checkpoint converters and serving runtimes generally assume every parameter lives with the model. Split-residency models routinely need bespoke loader support, which is why a new architecture of this kind usually arrives with a gap between the reference implementation and the community runtimes.
It does not touch activation memory. Offloading weights leaves the KV cache exactly where it was. On long contexts the cache overtakes the weights as the dominant term, and at that point a PLE-style saving is rounding error next to what cross-layer KV sharing or quantised caching would give you.
References and further reading
Every source this page cites, in the order it cites them. All of them open in a new tab.
- Google, 2025, Introducing Gemma 3n: the developer guide developers.googleblog.com
- Lan et al., 2019, ALBERT, arXiv:1909.11942 arxiv.org
- Baykal et al., 2023, Alternating Updates for Efficient Transformers, NeurIPS 2023, arXiv:2301.13310 arxiv.org
- Menghani et al., 2024, LAuReL: Learned Augmented Residual Layer, arXiv:2411.07501 arxiv.org
6 flashcards for this concept
Click a card to reveal the answer.