L$^3$: Large Lookup Layers

TL;DR

The paper tackles the inefficiencies of dynamic routing in Mixtures-of-Experts by introducing L$^3$, a Large Lookup Layer that generalizes token embedding tables to decoder layers with static routing and context-aware aggregation. It couples a systems-friendly architecture with an information-theoretic LZW-based embedding allocation to balance speed and model quality, enabling CPU-offloaded inference and efficient training. Empirical results show that L$^3$ layers, tested up to 2.6B active parameters, outperform iso-FLOP dense models and iso-sparse MoEs on language modeling and downstream tasks, demonstrating a practical, hardware-friendly axis of sparsity. The work highlights how static routing and large embedding caches can cache information and reduce compute without sacrificing model performance, suggesting a viable path for scalable sparse transformers in real-world deployments.

Abstract

Modern sparse language models typically achieve sparsity through Mixture-of-Experts (MoE) layers, which dynamically route tokens to dense MLP "experts." However, dynamic hard routing has a number of drawbacks, such as potentially poor hardware efficiency and needing auxiliary losses for stable training. In contrast, the tokenizer embedding table, which is natively sparse, largely avoids these issues by selecting a single embedding per token at the cost of not having contextual information. In this work, we introduce the Large Lookup Layer (L$^3$), which unlocks a new axis of sparsity by generalizing embedding tables to model decoder layers. L$^3$ layers use static token-based routing to aggregate a set of learned embeddings per token in a context-dependent way, allowing the model to efficiently balance memory and compute by caching information in embeddings. L$^3$ has two main components: (1) a systems-friendly architecture that allows for fast training and CPU-offloaded inference with no overhead, and (2) an information-theoretic embedding allocation algorithm that effectively balances speed and quality. We empirically test L$^3$ by training transformers with up to 2.6B active parameters and find that L$^3$ strongly outperforms both dense models and iso-sparse MoEs in both language modeling and downstream tasks.

Figures (11)

• Figure 1: L$^3$ scaling. Adding sparsity with L$^3$ layers significantly improves performance over iso-FLOP dense models ($\star$) and iso-FLOP, iso-sparse MoEs ($\times$).
• Figure 2: The basic architecture of a L$^3$ layer for a single token. The token ID $t$ is used to select embeddings $K_t, V_t$. Embeddings are aggregated with attention from the token hidden state $x$. The attention output is combined with the residual stream (hidden state $x$) to produce the L$^3$ output. Embeddings are assigned to token IDs with an information-theoretic LZW-based algorithm.
• Figure 3: Efficient training with L$^3$ . Since L$^3$ only performs channel mixing, all tokens in a batch can be sorted to form a block-diagonal "attention mask" for L$^3$ embedding aggregation. This allows L$^3$ to use fast kernels such as MegaBlocks or FlexAttention. The only overhead is sorting tokens, which can be done ahead of time, and inverting the sort, which is trivial.
• Figure 4: Fast offloaded inference with L$^3$ . Since the set of needed L$^3$ parameters is known when a token is generated, L$^3$ parameters can be offloaded and prefetched during pre-L$^3$ compute.
• Figure 5: Histogram of embedding count for 710K embeddings, a 180K vocab BPE tokenizer, and Algorithm \ref{['alg:lzw']} with $k=512$. "Common" tokens such as "then" get 512 embeddings while "rare" tokens such as "orm" get 1 embedding.