STEM: Scaling Transformers with Embedding Modules

TL;DR

STEM replaces the up-projection of a gated FFN with a per-token embedding lookup from a layer-local table, while leaving the gate and down projections dense and shared across tokens. The result is an extreme-sparsity Transformer variant where each token’s “expert” is statically chosen by token ID — no learned router, no all-to-all communication, and the giant lookup table can sit on CPU and be asynchronously prefetched. Empirically the recipe trains without the loss spikes that plague fine-grained MoEs, beats dense baselines on downstream tasks, and removes roughly one-third of FFN parameters from the per-token activation budget.

Key claims

The FFN up-projection can be replaced by a token-indexed embedding lookup E[token_id] while keeping the gate and down projections as ordinary dense matrices shared across tokens [§3].
This static, token-ID routing is equivalent to a hash-layer MoE in spirit but removes the learned router entirely, so there is no load-balancing loss and no expert-frequency drift during training [§2, §3].
Decoupling per-token capacity from FLOPs and from cross-device communication lets the embedding tables live in node-local CPU memory and be asynchronously prefetched, avoiding the all-to-all traffic that fine-grained MoE incurs [Abstract, §3].
STEM does not exhibit the loss spikes that fine-grained MoE models typically suffer from at extreme sparsity [§4, Fig. 5(a)].
The learned per-token embeddings spread out angularly (low pairwise cosine similarity) — the authors interpret this as reduced representational interference / better addressability of the parametric memory [§4].
STEM improves downstream performance over dense baselines while reducing per-token FLOPs and parameter accesses, eliminating roughly one-third of FFN parameters from the activated set [Abstract].
STEM positions itself as a direct extension of Gemma’s Per-Layer Embedding (PLE), but unlike PLE — which adds an embedding block alongside a full FFN — STEM actually replaces part of the FFN, making the embedding the substitute rather than a supplement [§2].

Method

In a gated FFN (SwiGLU-style), each token’s hidden state x is mapped to down(gate(W_g · x) ⊙ (W_u · x)). STEM keeps W_g and down as dense layers shared by all tokens, but replaces W_u · x with a token-indexed lookup E_ℓ[token_id] from a layer-local embedding table E_ℓ. Each token thus has its own “up-projection vector” per layer — a static, fine-grained form of sparsity, since the choice of which row of E_ℓ to read is determined entirely by the token ID (no router network, no top-k softmax, no load-balancing auxiliary loss).

The systems consequence is that E_ℓ is huge but accessed in a strictly sparse, predictable way: token IDs in the upcoming batch are known, so the relevant rows can be prefetched from CPU memory before the layer fires. This pushes the bulk of “parameter capacity” off the GPU and onto cheaper RAM, decoupling parameter count from both per-token FLOPs and cross-device bandwidth. STEM is presented as building on Per-Layer Embedding from Gemma 3n and on hash-layer MoE / MoWE; the closest precursor is PLE, but PLE keeps the FFN intact and just adds an embedding block, whereas STEM substitutes part of the FFN.

Results

STEM trains stably under extreme sparsity, where comparable fine-grained MoEs show loss spikes [Fig. 5(a)].
STEM beats dense baselines on downstream evaluations while reducing per-token FLOPs and activated parameter count — about one-third of the FFN parameters are no longer hit on each forward pass [Abstract].
Embedding-space analysis shows STEM’s learned per-token vectors have large angular spread (low pairwise cosine similarity), which the authors connect to better information storage in the parametric memory [§4].

Why it’s interesting

This is the second sparsity-via-embedding-table paper to land in #research-external in a month — the first was DeepSeek’s Engram, which @rosinality and @Haoxiang both flag as the natural comparison point. Both papers argue that the standard MoE recipe (learned router + top-k experts + load-balance loss + all-to-all) is doing more work than it needs to, and that a static token-indexed lookup recovers most of the capacity benefit while sidestepping training instability and communication overhead. STEM and Engram differ on what they index — STEM uses the raw token ID, Engram uses a hashed local N-gram — but the underlying bet is the same: parametric memory should be a lookup, not a routed computation.

For Luma, this is most relevant as a possible scaling axis for the text-side of T2V conditioning stacks (where the wiki is already tracking frozen LLM text encoders like LTX-2’s Gemma-3 connector) and, more speculatively, as an inspiration for analogous mechanisms in DiTs — e.g. token-ID-style lookups for the (small, finite) class tokens or modality tokens that DiT blocks process. The architectural minimalism of “drop the up-projection, add an embedding table” makes this cheap to try.