Shared projection parameters allow switching mixers during generation
Excited to share last summer's work at Google Research!
Most hybrid models today are static: each token sees the same interleaved pattern of your favorite linear model and attention.
Oryx instead varies the model used across the sequence through shared representations.
1/
Many users praised Google's Oryx dynamic hybrid model for sequences as cool work and a nice read while offering supportive technical feedback on its attention mechanisms.
No Digg Deeper questions have been answered for this story yet.
Rather than interleaving layers naively, a more fine-grained approach to hybrid models is to allow hybridization across the sequence models within a single layer. The fact that softmax attention and linear attention use similar underlying projection parameters allows switching between different mixers in a single generation, for the best of both worlds.
Excited to share last summer's work at Google Research!
Most hybrid models today are static: each token sees the same interleaved pattern of your favorite linear model and attention.
Oryx instead varies the model used across the sequence through shared representations.
1/
@atrost3122 Great question! I think it would probably be quite engineering heavy since ideally the queries to attention would be << than that of the linear.
FlexAttention would probably be pretty helpful.
https://arxiv.org/abs/2603.22325 is also very relevant! @ZyphraAI has amazing work!
Paper: http://arxiv.org/abs/2605.28769
This work was done with my wonderful collaborators @ashertrockman @th33rtha @sziteng
2/
Despite attention and linear models using distinct states, they are connected via the associative memory viewpoint.
Because both operate by querying key-value pairs, we can share their core representations by tying >90% of weights! Only queries are untied for flexibility.
3/
To unlock the full capabilities of both modes, shared Oryx blocks use the gate, short conv, and norm found in many linear models.
We test our design with both Mamba-2 & Gated DeltaNet as the linear mechanism (with minor differences between variants) and find both work well!
4/
One of the coolest results is the mixer/mode switching of Oryx models!
Both Mamba-2 and GDN variants can prefill context in linear mode then generate in attention mode, recovering much of the retrieval capabilities pure linear baselines lose for NIAH and real-world tasks.
6/
Surprisingly, Oryx retains strong performance when using just one mechanism for the entire sequence.
At 1.4B scale, both attention and linear modes of Oryx outperform their respective pure baselines by 0.7 points or more, even with linear models disadvantaged by state size!
5/
Many open questions remain to be explored!
For instance, how can we tune a base Oryx model to route tokens automatically?
What are the synergies/trade-offs with current hybrid paradigms?
Why does representation sharing work, when does it arise, and how does it break down?
7/7
@kevinyli_ Cool work and a nice read! For token-level routing prospects, have you thought about how to handle heterogeneous batches where different tokens get routed to different mixers? Seems like an engineering challenge but I wonder if simple solutions exist
@kevinyli_ @ZyphraAI Makes sense yeah. And neat I'll check that out thanks!
