🤖 machine learning

Multi-Mixer Models: Flexible Sequence Modeling with Shared Representations

The paper introduces Oryx, a hybrid architecture that dynamically switches between efficient linear recurrent mixers and quadratic attention mechanisms along the token sequence while sharing over 90% of parameters, achieving superior performance and long-context retrieval capabilities compared to single-mixer baselines.

Original authors: Kevin Y. Li, Asher Trockman, Ananda Theertha Suresh, Ziteng Sun

Published 2026-05-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Kevin Y. Li, Asher Trockman, Ananda Theertha Suresh, Ziteng Sun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to write a long story. To do this well, you need two different tools in your toolbox:

The "Super-Scanner" (Softmax Attention): This tool is amazing at looking back at everything you've written so far to find specific details, connect ideas, and understand the whole context. It's like having a librarian who can instantly pull up every single book you've ever read to find the perfect quote. However, this librarian is slow and expensive; the more books you have, the longer it takes them to search, and they need a huge shelf to hold all the books.
The "Speed-Runner" (Linear Recurrent Models): This tool is incredibly fast and efficient. It doesn't look back at the whole library; instead, it keeps a small, compact summary in its head and updates it as it reads. It's like a runner who memorizes the gist of the story as they go. It's cheap and fast, but sometimes it forgets the specific details needed to answer tricky questions.

For a long time, AI models had to choose one or the other. If they wanted speed, they lost detail. If they wanted detail, they lost speed.

Enter "Oryx": The Flexible Storyteller

The paper introduces a new model called Oryx that refuses to choose. Instead, it acts like a chameleon that can switch between being a "Super-Scanner" and a "Speed-Runner" while it is writing the story.

Here is how it works, using simple analogies:

1. The Shared Brain (Shared Representations)

Usually, if you switch from a librarian to a runner, you have to hand them a completely new set of notes. They don't speak the same language.

Oryx solves this by giving both the "Scanner" and the "Runner" the same notebook.

They both read the same words and write down the same key facts (called "shared representations").
Because they are writing in the same notebook, they understand each other perfectly.
This means the model can switch from the slow, detailed Scanner to the fast Runner (or back again) without losing its place or forgetting what happened. It's like a worker who can instantly switch from doing detailed accounting to running a delivery, because they are using the same clipboard for both tasks.

2. The "Chunked" Training (Learning to Switch)

How do you teach a model to switch gears so smoothly? The researchers used a training method called "Chunked Mixed-Mode Training."

Imagine you are training a student for a test. Instead of making them study the whole book using only one method, you break the book into small chapters (chunks).

For Chapter 1, you tell the student: "Use the Scanner to read this."
For Chapter 2, you say: "Switch to the Runner for this one."
For Chapter 3, go back to the Scanner.

By practicing this back-and-forth switching during training, the model learns that it's okay to change tools in the middle of a sentence. It learns that the "Shared Notebook" works for both tools.

3. The Results: Best of Both Worlds

The paper tested Oryx on various tasks and found some impressive things:

It's as smart as the best: When Oryx uses the "Scanner" mode, it is just as good at understanding language as the best traditional models.
It's as fast as the best: When it uses the "Runner" mode, it is efficient and fast.
The Magic Switch: The coolest part is that Oryx can use the Runner to read through a massive amount of text (like a long document) to save time and memory, and then instantly switch to the Scanner to answer a specific question about that text.
- The paper found that Oryx could process a huge story using the fast "Runner" for 90% of the text, and only switch to the slow "Scanner" for the last 10% (the part where the question is asked).
- Even with this tiny bit of "Scanner" time, it performed just as well as models that used the "Scanner" for the entire story.

Why This Matters (According to the Paper)

The paper suggests that we don't need to pick a side anymore. We can build models that are hybrids. They can be efficient and cheap for most of the work, but switch to being powerful and detailed exactly when they need to be.

The authors call this "Sequence-Axis Hybridization." Instead of building a model that is half-scanner and half-runner (like a car with one wheel), Oryx is a vehicle that can drive on wheels for most of the trip and switch to a jet engine for the steep hill, all while using the same steering wheel and dashboard.

In short: Oryx is a flexible AI that shares its memory between a fast mode and a smart mode, allowing it to switch between them instantly without getting confused, making it both efficient and highly capable.

Technical Summary: Multi-Mixer Models (Oryx)

Problem Statement

Softmax attention is the dominant mechanism in modern large language models (LLMs) but incurs quadratic computational costs and linear memory requirements relative to sequence length due to the maintenance of a full key-value (KV) cache. While linear recurrent models (e.g., linear attention and state space models like Mamba-2) offer linear compute and constant memory, they often lag behind attention-based models in tasks requiring strong retrieval or in-context learning. Existing hybrid architectures attempt to mitigate these trade-offs by statically interleaving attention and recurrent blocks across layers (inter-layer) or fusing them within a single layer (intra-layer). However, these approaches fix the computational pathway per token or layer, failing to adapt the modeling mechanism to the specific needs of different segments within a sequence.

Methodology: The Oryx Architecture

The authors propose Oryx, a sequence-axis hybrid architecture that enables flexible switching between different sequence mixers (e.g., quadratic softmax attention and linear recurrent mechanisms) throughout a single token sequence.

Core Design Principles

Shared Representations via Tied Projections:
- Oryx unifies attention and linear models under an associative memory view where both maintain key-value (KV) associations.
- The architecture employs tied projections for keys ( $W_K$ ) and values ( $W_V$ ) across mixer modes. This allows a single forward pass to update both the attention KV cache and the linear recurrent state simultaneously using shared representations.
- Mixer-Specific Queries: Unlike keys and values, query projections ( $W_Q$ ) are kept distinct for each mixer. Empirical results indicate that sharing query weights degrades performance, suggesting that attention and recurrent mechanisms require different readout vectors to extract mode-specific information from the shared state.
Architectural Components:
- The block incorporates components critical to linear model performance, such as short convolutions and multiplicative gates, applied to the shared keys and values.
- To ensure compatibility, Oryx enforces a unified head structure (Multi-Head Attention, MHA) across both mixers, even though linear models like Mamba-2 typically use Multi-Value (MVA) structures.
- The block maintains both the KV cache and the linear recurrent state at every timestep, updating them jointly.
Chunked Mixed-Mode Training:
- To enable robust mode switching during inference, Oryx is trained using a chunked mixed-mode strategy. Training sequences are partitioned into fixed-length chunks (e.g., 128 tokens), where each chunk is randomly assigned a mixer mode (attention or linear).
- This training regime forces the model to learn compatible representations that allow seamless transitions between modes, preventing the performance degradation often seen when switching from linear to attention mechanisms in sequence-level mixed training.

Key Contributions

Sequence-Axis Hybridization: Oryx introduces a new axis for hybridization, allowing the model to switch mechanisms dynamically across the token sequence rather than being constrained by fixed layer or block architectures.
Parameter Efficiency: The design ties at least 90% of parameters (specifically $W_K, W_V, W_G, W_O$ ) across mixer modes, enabling attention and recurrent modes to operate over shared internal representations without significant parameter overhead.
Flexible Inference: The model supports inference where different segments of a sequence (e.g., context prefill vs. generation) can be processed by different mixers. For instance, a model can use linear recurrence for efficient context processing and switch to attention for retrieval-heavy generation.

Empirical Results

Experiments were conducted on models ranging from 130M to 1.4B parameters, using Mamba-2 and Gated DeltaNet (GDN) as the linear mixers.

Language Modeling: Under fixed token budgets, Oryx models (both Oryx-TM and Oryx-TG variants) achieve comparable or better performance than single-mixer baselines (pure Transformer, pure Mamba-2, pure GDN). At the 1.4B scale, Oryx variants outperform their respective baselines by at least 0.7 percentage points on averaged language modeling tasks.
Retrieval Capabilities:
- Oryx preserves the retrieval capabilities of its constituent mixers.
- Crucially, when using a mixed-inference mode (processing context with linear recurrence and generation with attention), Oryx achieves performance comparable to the Transformer baseline while processing <10% of tokens in attention mode.
- On real-world retrieval tasks, mixed-mode Oryx surpasses linear baselines by 8.6 percentage points (Oryx-TG) and 13.5 percentage points (Oryx-TM).
- On synthetic "Needle-in-a-Haystack" (NIAH) tests, Oryx outperforms linear baselines by 38.6 to 40.3 percentage points, demonstrating that shared representations effectively support long-context retrieval even when the linear mixer handles the majority of the sequence.
Mode Switching: Perplexity analysis shows that after switching between mixers (e.g., from attention to Mamba-2 or vice versa), the model's perplexity rapidly converges to the baseline of the selected mode, indicating that the shared states are compatible and the transition is seamless.

Significance and Claims

The paper claims that Oryx demonstrates that attention and linear recurrent mechanisms can share internal representations despite their differing methodologies. This finding challenges the assumption that these mechanisms require entirely separate state spaces.

The significance of this work lies in:

Validating Sequence-Axis Hybridization: It establishes that dynamically switching mixers within a sequence is a viable and effective strategy, offering a new design space for adaptive compute.
Bridging the Retrieval Gap: It shows that linear models can achieve retrieval performance comparable to Transformers if they can leverage attention modes for specific segments, without requiring a full attention mechanism for the entire sequence.
Future Directions: The authors suggest this architecture opens opportunities for learned routing (e.g., using RL to select mixers based on FLOP savings) and speculative decoding, where linear modes could draft reasoning traces for verification by attention modes.

The authors remain modest, noting that maintaining multiple states (KV cache and linear state) incurs memory overhead, though this is dominated by the KV cache at long context lengths. They also note that the training dynamics of multi-mixer models remain largely unexplored and may require specialized objectives as scale increases.