Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ShishuLM: Lightweight Language Model with Hybrid Decoder-MLP Architecture and Paired Weight Sharing

Shivanshu Kumar, Gopalakrishnan Srinivasan

TL;DR

ShishuLM introduces a hybrid decoder-MLP architecture that omits attention in deeper blocks and shares weights across adjacent MLP layers to reduce parameter count and KV cache requirements. By leveraging linearity in the attention components and interpretability-driven weight distribution analyses, the paper demonstrates that up to ~2/3 of attention layers can be removed with minimal performance loss, while weight sharing in the remaining MLP blocks yields notable memory and latency gains. Evaluations on MobileLLM-125M and -600M show up to 25% memory reduction and up to 40% latency improvement with negligible drops in zero-shot performance on several downstream tasks. The work offers a practical, pre-training-friendly path toward more efficient small-language-model architectures, with limitations tied to model scale and architecture specificity.

Abstract

While the transformer architecture has achieved state-of-the-art performance on natural language processing tasks, these models impose substantial memory and computational overhead. Recent research has identified significant architectural redundancies within these models, presenting opportunities for optimization without compromising performance. Taking insights from research in AI interpretability and inference-time layer pruning, we introduce an efficient language model architecture, referred to as ShishuLM, which reduces both the parameter count and Key-Value (KV) cache requirements. Given the increasing importance of Small Language Models (SLMs) in agentic AI systems, we evaluate our approach on two SLMs of different scales. Our analysis reveals that for moderate-context scenarios, normalization coupled with attention computation is roughly linear with the input, enabling entire transformer blocks to be approximated through Multi-Layer Perceptrons (MLPs). Our results show that ShishuLM provides up to 25% reduction in memory requirements and up to 40% improvement in latency during both training and inference, compared to parent models. Our experimental and analytical findings provide insights towards building more efficient SLM architectures from a pre-training standpoint.

ShishuLM: Lightweight Language Model with Hybrid Decoder-MLP Architecture and Paired Weight Sharing

TL;DR

ShishuLM introduces a hybrid decoder-MLP architecture that omits attention in deeper blocks and shares weights across adjacent MLP layers to reduce parameter count and KV cache requirements. By leveraging linearity in the attention components and interpretability-driven weight distribution analyses, the paper demonstrates that up to ~2/3 of attention layers can be removed with minimal performance loss, while weight sharing in the remaining MLP blocks yields notable memory and latency gains. Evaluations on MobileLLM-125M and -600M show up to 25% memory reduction and up to 40% latency improvement with negligible drops in zero-shot performance on several downstream tasks. The work offers a practical, pre-training-friendly path toward more efficient small-language-model architectures, with limitations tied to model scale and architecture specificity.

Abstract

While the transformer architecture has achieved state-of-the-art performance on natural language processing tasks, these models impose substantial memory and computational overhead. Recent research has identified significant architectural redundancies within these models, presenting opportunities for optimization without compromising performance. Taking insights from research in AI interpretability and inference-time layer pruning, we introduce an efficient language model architecture, referred to as ShishuLM, which reduces both the parameter count and Key-Value (KV) cache requirements. Given the increasing importance of Small Language Models (SLMs) in agentic AI systems, we evaluate our approach on two SLMs of different scales. Our analysis reveals that for moderate-context scenarios, normalization coupled with attention computation is roughly linear with the input, enabling entire transformer blocks to be approximated through Multi-Layer Perceptrons (MLPs). Our results show that ShishuLM provides up to 25% reduction in memory requirements and up to 40% improvement in latency during both training and inference, compared to parent models. Our experimental and analytical findings provide insights towards building more efficient SLM architectures from a pre-training standpoint.
Paper Structure (20 sections, 20 equations, 8 figures, 10 tables)

This paper contains 20 sections, 20 equations, 8 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries: Earth Mover's Distance
  4. ShishuLM: Architecture and Methodology
  5. ShishuMLP
  6. Optimization Methodology
  7. Linearity in the Attention Modules
  8. Similarity in MLP Weight Distribution
  9. Experiments
  10. Results
  11. Performance
  12. Systems level improvements
  13. Conclusion
  14. Limitations
  15. Directions for Future Work
  16. ...and 5 more sections

Figures (8)

  • Figure 1: Vanilla transformer block(left) and a pair of consecutive ShishuMLP blocks with shared weights (right)
  • Figure 2: Schematic representation of the ShishuLM architecture as compared to the Llama architecture
  • Figure 3: Loss curves for all the model variants.
  • Figure 4: Mean squared error for fitting a linear layer in place of input normalisation + self-attention
  • Figure 5: Average cosine similarity in generated tokens
  • ...and 3 more figures