Thinking Slow, Fast: Scaling Inference Compute with Distilled Reasoners

• URL: http://arxiv.org/abs/2502.20339v1
• Date: Thu, 27 Feb 2025 18:08:16 GMT
• Title: Thinking Slow, Fast: Scaling Inference Compute with Distilled Reasoners
• Authors: Daniele Paliotta, Junxiong Wang, Matteo Pagliardini, Kevin Y. Li, Aviv Bick, J. Zico Kolter, Albert Gu, François Fleuret, Tri Dao,
• Abstract summary: Recent advancements have demonstrated that the performance of large language models (LLMs) can be significantly enhanced by scaling computational resources at test time.<n>This raises a fundamental question: can models with lower complexity leverage their superior generation throughput to outperform similarly sized Transformers for a fixed computational budget?<n>To address this question and overcome the lack of strong subquadratic reasoners, we distill pure and hybrid Mamba models from pretrained Transformers.
• Score: 72.37408197157453
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Recent advancements have demonstrated that the performance of large language models (LLMs) can be significantly enhanced by scaling computational resources at test time. A common strategy involves generating multiple Chain-of-Thought (CoT) trajectories and aggregating their outputs through various selection mechanisms. This raises a fundamental question: can models with lower complexity leverage their superior generation throughput to outperform similarly sized Transformers for a fixed computational budget? To address this question and overcome the lack of strong subquadratic reasoners, we distill pure and hybrid Mamba models from pretrained Transformers. Trained on only 8 billion tokens, our distilled models show strong performance and scaling on mathematical reasoning datasets while being much faster at inference for large batches and long sequences. Despite the zero-shot performance hit due to distillation, both pure and hybrid Mamba models can scale their coverage and accuracy performance past their Transformer teacher models under fixed time budgets, opening a new direction for scaling inference compute.

Related papers

• M1: Towards Scalable Test-Time Compute with Mamba Reasoning Models [72.75501495786297]
  We introduce a novel hybrid linear RNN reasoning model, M1, built on the Mamba architecture. Experimental results show that M1 not only outperforms previous linear RNN models but also matches the performance of state-of-the-art DeepSeek R1 distilled reasoning models.
  arXiv  Detail & Related papers  (2025-04-14T17:38:25Z)
• LazyDiT: Lazy Learning for the Acceleration of Diffusion Transformers [79.07412045476872]
  Diffusion Transformers have emerged as the preeminent models for a wide array of generative tasks. We show that performing the full of the model at each diffusion step is unnecessary, as some computations can be skipped by lazily reusing the results of previous steps. We propose a lazy learning framework that efficiently leverages cached results from earlier steps to skip redundant computations.
  arXiv  Detail & Related papers  (2024-12-17T01:12:35Z)
• The Mamba in the Llama: Distilling and Accelerating Hybrid Models [76.64055251296548]
  We show how to distill large Transformers into linear RNNs by reusing the linear projection weights from attention layers with academic GPU resources.<n>The resulting hybrid model achieves performance comparable to the original Transformer in chat benchmarks.<n>We also introduce a hardware-aware speculative decoding algorithm that accelerates the inference speed of Mamba and hybrid models.
  arXiv  Detail & Related papers  (2024-08-27T17:56:11Z)
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces [31.985243136674146]
  Foundation models are almost universally based on the Transformer architecture and its core attention module. We identify that a key weakness of such models is their inability to perform content-based reasoning. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even blocks (Mamba) As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics.
  arXiv  Detail & Related papers  (2023-12-01T18:01:34Z)
• Fourier Transformer: Fast Long Range Modeling by Removing Sequence Redundancy with FFT Operator [24.690247474891958]
  Fourier Transformer is able to significantly reduce computational costs while retain the ability to inherit from various large pretrained models. Our model achieves state-of-the-art performances among all transformer-based models on the long-range modeling benchmark LRA. For generative seq-to-seq tasks including CNN/DailyMail and ELI5, by inheriting the BART weights our model outperforms the standard BART.
  arXiv  Detail & Related papers  (2023-05-24T12:33:06Z)
• Long Range Arena: A Benchmark for Efficient Transformers [115.1654897514089]
  Long-rangearena benchmark is a suite of tasks consisting of sequences ranging from $1K$ to $16K$ tokens. We systematically evaluate ten well-established long-range Transformer models on our newly proposed benchmark suite.
  arXiv  Detail & Related papers  (2020-11-08T15:53:56Z)
• AxFormer: Accuracy-driven Approximation of Transformers for Faster, Smaller and more Accurate NLP Models [4.247712017691596]
  AxFormer is a framework that applies accuracy-driven approximations to create optimized transformer models for a given downstream task. Our experiments show that AxFormer models are up to 4.5% more accurate, while also being up to 2.5X faster and up to 3.2X smaller than conventional fine-tuned models.
  arXiv  Detail & Related papers  (2020-10-07T23:29:34Z)
• Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers [94.43313684188819]
  We study the impact of model size in this setting, focusing on Transformer models for NLP tasks that are limited by compute. We first show that even though smaller Transformer models execute faster per iteration, wider and deeper models converge in significantly fewer steps. This leads to an apparent trade-off between the training efficiency of large Transformer models and the inference efficiency of small Transformer models.
  arXiv  Detail & Related papers  (2020-02-26T21:17:13Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.