Zyphra

🌈 Abstract

The article discusses the challenges of training hybrid models and presents the Zyphra Training Cookbook, which aims to enable other technical groups to build their own hybrid models. It focuses on practical engineering work and architectural innovations to maximize performance per parameter and inference FLOP, particularly for smaller models deployed on devices with strict power and memory constraints.

🙋 Q&A

[01] Model Architectures

1. Why do the authors think hybrid architectures offer the best model quality per training/inference FLOPs?

• Dense transformers have shortcomings such as high computational cost of attention at long sequence lengths and lack of efficiency compared to alternative sequence mixers like Mamba and RWKV.
• Mixture of Experts (MoE) architectures have the inference latency of their forward-pass parameters and require all parameters to be loaded into VRAM, limiting inference to distributed GPU clusters.
• State Space Models (SSMs) offer a more efficient alternative to attention, but may require significantly more tokens to match the performance of attention-based models on tasks like in-context learning and long-context reasoning.
• Dense hybrid architectures combine the strengths of dense transformers and SSMs - they don't introduce the memory overhead of MoEs, maintain the cross-sequence dependencies of attention, and have inference latency close to pure SSMs.

[02] Architectural Innovations

1. What are the key architectural innovations in the Zamba models?

• Zamba uses a parameter-sharing scheme where a single transformer block consisting of an attention and an MLP block is re-used multiple times throughout the network, comprising the only attention in the network.
• This increases performance per parameter at the expense of additional FLOPs, but the inherent efficiency of the Mamba backbone results in an architecture that outperforms transformers in both equi-token and equi-flop conditions.
• Concatenating the original text embeddings with the current layer embeddings at every shared attention block provided a significant boost to performance per parameter.
• Applying LoRAs to the shared layers in Zamba2 models further specializes the shared blocks, improving performance at a very small parameter cost.

2. What insights does the success of the Zamba architecture provide?

• Even when attention is used rarely, there is still great redundancy in the attention parameters, suggesting attention is primarily needed to "remind" the network of the past sequence in a few stereotyped ways, rather than to perform novel sequence mixing operations at every attention block.
• The Zamba architecture exploits this regularity to reduce the parameter count of the model for a given level of performance.

[03] Training Approaches

1. What training approaches did the authors find effective?

• A curriculum training approach of increasing the proportion of higher quality tokens (e.g. fact-rich documents) towards the end of training can significantly improve performance.
• Maintaining a high "replay fraction" of tokens from the original pre-training dataset (50-70%) while incorporating annealing data (30-50%) is important to stabilize training and maintain performance.
• Multiple epochs of annealing data do not harm performance, but beyond 2 epochs provide little additional improvement.