[Paper Review] Mechanistic Design and Scaling of Hybrid Architectures

Michael Poli, Armin W. Thomas+10 more|arXiv (Cornell University)|Mar 26, 2024

Architecture and Computational Design4 citations

TL;DR

This paper introduces Mechanistic Architecture Design (MAD), a fast, scalable pipeline that uses synthetic token manipulation tasks—like compression and recall—to predict the scaling behavior of hybrid deep learning architectures. By training small models on these tasks, MAD identifies high-performing architectures that outperform SOTA models (e.g., Transformers, Mamba, Hyena) in both compute-optimal and overtrained regimes, with strong correlation between MAD scores and actual scaling performance.

ABSTRACT

The development of deep learning architectures is a resource-demanding process, due to a vast design space, long prototyping times, and high compute costs associated with at-scale model training and evaluation. We set out to simplify this process by grounding it in an end-to-end mechanistic architecture design (MAD) pipeline, encompassing small-scale capability unit tests predictive of scaling laws. Through a suite of synthetic token manipulation tasks such as compression and recall, designed to probe capabilities, we identify and test new hybrid architectures constructed from a variety of computational primitives. We experimentally validate the resulting architectures via an extensive compute-optimal and a new state-optimal scaling law analysis, training over 500 language models between 70M to 7B parameters. Surprisingly, we find MAD synthetics to correlate with compute-optimal perplexity, enabling accurate evaluation of new architectures via isolated proxy tasks. The new architectures found via MAD, based on simple ideas such as hybridization and sparsity, outperform state-of-the-art Transformer, convolutional, and recurrent architectures (Transformer++, Hyena, Mamba) in scaling, both at compute-optimal budgets and in overtrained regimes. Overall, these results provide evidence that performance on curated synthetic tasks can be predictive of scaling laws, and that an optimal architecture should leverage specialized layers via a hybrid topology.

Motivation & Objective

• Address the high cost and inefficiency of architectural prototyping in deep learning due to vast design spaces and long training cycles.
• Develop a systematic, scalable method to evaluate and optimize new architectures before full-scale training.
• Investigate whether performance on curated synthetic tasks can predict real-world scaling behavior of emerging architectures.
• Explore hybrid architectures combining computational primitives (e.g., convolutions, recurrences, MoEs) to improve scaling efficiency and performance.
• Establish a framework for fast, compute-efficient architecture search grounded in mechanistic understanding of model capabilities.

Proposed method

• Design a suite of synthetic token manipulation tasks—compression, recall, memorization—targeted at probing specific model capabilities like in-context learning and sequence retention.
• Apply MAD to evaluate architectures at small scale (e.g., 2-block models) using normalized state dimension to ensure fair comparison across recurrent and non-recurrent primitives.
• Use isolated proxy tasks to compute aggregate MAD scores, representing an architecture’s capability profile across multiple unit tests.
• Train over 500 language models (70M–7B parameters) across diverse architectures to perform compute-optimal and state-optimal scaling law analyses.
• Correlate MAD scores with compute-optimal perplexity on The Pile to validate predictive power, using iso-FLOP groups for consistent evaluation.
• Design hybrid architectures via striping—interleaving blocks of different computational primitives (e.g., Mamba + Mixture-of-Experts) to leverage specialized capabilities.

Experimental results

Research questions

• RQ1 Can synthetic, small-scale token manipulation tasks serve as reliable predictors of large-scale model performance and scaling laws?
• RQ2 Do hybrid architectures combining specialized computational primitives (e.g., gated convolutions, recurrences) outperform uniform architectures like Transformers in scaling efficiency?
• RQ3 Is there a strong correlation between performance on MAD synthetics and actual compute-optimal perplexity in real-world language modeling?
• RQ4 Can the MAD pipeline guide iterative architecture improvement by identifying which primitives excel at specific sub-tasks?
• RQ5 How do state-optimal scaling laws differ from compute-optimal ones, and what implications do they have for inference efficiency?

Key findings

• Aggregate MAD scores show strong linear correlation with compute-optimal perplexity across all FLOP budgets, validating MAD as a predictive proxy for scaling behavior.
• Hybrid architectures—particularly striped variants combining Mamba and MoE or other primitives—achieved superior scaling performance compared to SOTA models like Transformer++, Hyena, and Mamba.
• The best-performing architectures identified via MAD outperformed state-of-the-art models in both compute-optimal regimes and overtrained settings.
• Performance on the memorization task showed little variation across architectures and did not correlate with scaling performance, suggesting limited discriminative power for this task.
• The state-optimal scaling law analysis revealed that fixed-state blocks (e.g., Mamba) can achieve better efficiency in inference memory and latency, highlighting trade-offs in architectural design.
• Iterative improvements guided by MAD—such as refining Hyena variants—translated directly into measurable gains in perplexity on The Pile at 2×10¹⁹ FLOPs, confirming the pipeline’s practical utility.