MoVE: Mixture of Value Embeddings -- A New Axis for Scaling Parametric Memory in Autoregressive Models

TL;DR

MoVE introduces Mixture of Value Embeddings, a global, differentiable memory bank shared across all attention layers that decouples parametric memory from network depth in autoregressive transformers. By using a per-head soft routing mechanism, MoVE dynamically mixes retrieved memory concepts into the Value stream, enabling memory capacity to scale via the number of embedding slots $M$ without increasing depth or FLOPs significantly. Empirical results in text and image generation show consistent gains over standard and layer-wise memory baselines, and MoVE extends to compressed latent spaces in MLA, illustrating broad architectural versatility. The approach opens a new scaling axis for generative models, yielding memory-dense systems that can achieve higher fidelity and perplexity improvements within similar compute budgets, while highlighting areas for future hardware-aware optimization and hybrid architectures.

Abstract

Autoregressive sequence modeling stands as the cornerstone of modern Generative AI, powering results across diverse modalities ranging from text generation to image generation. However, a fundamental limitation of this paradigm is the rigid structural coupling of model capacity to computational cost: expanding a model's parametric memory -- its repository of factual knowledge or visual patterns -- traditionally requires deepening or widening the network, which incurs a proportional rise in active FLOPs. In this work, we introduce $\textbf{MoVE (Mixture of Value Embeddings)}$, a mechanism that breaks this coupling and establishes a new axis for scaling capacity. MoVE decouples memory from compute by introducing a global bank of learnable value embeddings shared across all attention layers. For every step in the sequence, the model employs a differentiable soft gating mechanism to dynamically mix retrieved concepts from this bank into the standard value projection. This architecture allows parametric memory to be scaled independently of network depth by simply increasing the number of embedding slots. We validate MoVE through strictly controlled experiments on two representative applications of autoregressive modeling: Text Generation and Image Generation. In both domains, MoVE yields consistent performance improvements over standard and layer-wise memory baselines, enabling the construction of "memory-dense" models that achieve lower perplexity and higher fidelity than their dense counterparts at comparable compute budgets.

Figures (10)

• Figure 1: Architectural Comparison. Left: Standard Attention, where the Value stream is a linear projection of the local hidden state $X$. Center: Layer-wise Value Embeddings (LaVE), a baseline that augments attention with a layer-specific memory vector gated by the input. Right:MoVE (Ours), which decouples memory from depth by introducing a global bank of value embeddings. Crucially, both layers fetch value embeddings from this central shared bank, employing a differentiable router to dynamically mix these global concepts into the value stream.
• Figure 2: Conceptual Illustration of Memory-Augmented Multi-Head Latent Attention. Left: Standard MLA compresses Key-Value heads into a low-rank latent vector $\mathbf{c}_{KV}$ to minimize cache footprint. Center: MLA + LaVE, which augments the latent space with layer-specific parameters. Right: MLA + MoVE (Ours), which injects a global shared memory into the latent space.
• Figure 3: Routing visualization for the polysemous word "bank" (Financial Institution vs. River Bank). In the Long context (bottom row), the routing stabilizes: the intra-meaning variation (Control Diff) decreases to 0.0513, while the inter-meaning variation (Semantic Diff) remains high at 0.1040, exhibited by distinct vertical banding.
• Figure 4: Routing visualization for "crane" (Construction Machine vs. Bird). The model maintains a consistent routing separation between the animate and inanimate senses across all context lengths.
• Figure 5: Routing visualization for "bat" (Sports Equipment vs. Animal). In the Short context, the routing is noisy: the variation between sentences with the same meaning (Control) is actually higher than the variation between different meanings (Semantic). However, as context increases (Medium/Long), the signal stabilizes and the Semantic difference correctly dominates.