MAR: Efficient Large Language Models via Module-aware Architecture Refinement

TL;DR

The paper tackles the high energy demands of large language models by proposing Module-aware Architecture Refinement (MAR), a two-stage framework that first adopts State Space Models for linear-time sequence processing and then sparsifies activations with spiking neurons to reduce FFN costs. It introduces Adaptive Ternary Multi-step Neuron (ATMN) to boost information density and Spike-aware Bidirectional Distillation Strategy (SBDS) to reconcile temporal dynamics and retain performance under constraints. Through extensive experiments on Llamba-1B, MAR restores teacher-level accuracy across multiple benchmarks while substantially lowering inference energy, and it outperforms larger efficient models of comparable scale. These results demonstrate MAR’s potential to enable practical, energy-efficient LLMs suitable for deployment under resource limitations.

Abstract

Large Language Models (LLMs) excel across diverse domains but suffer from high energy costs due to quadratic attention and dense Feed-Forward Network (FFN) operations. To address these issues, we propose Module-aware Architecture Refinement (MAR), a two-stage framework that integrates State Space Models (SSMs) for linear-time sequence modeling and applies activation sparsification to reduce FFN costs. In addition, to mitigate low information density and temporal mismatch in integrating Spiking Neural Networks (SNNs) with SSMs, we design the Adaptive Ternary Multi-step Neuron (ATMN) and the Spike-aware Bidirectional Distillation Strategy (SBDS). Extensive experiments demonstrate that MAR effectively restores the performance of its dense counterpart under constrained resources while substantially reducing inference energy consumption. Furthermore, it outperforms efficient models of comparable or even larger scale, underscoring its potential for building efficient and practical LLMs.

Figures (5)

• Figure 1: Energy consumption of LLaMA-3.2 Attention/FFN versus Llamba’s bick2025llamba Discrete Mamba-2. Crossovers occur at sequence lengths 2470 and 9330.
• Figure 2: Overview of the MAR framework. This figure illustrates how a dense attention-based model is transformed into a sparse and efficient linear sequence model through a two-stage optimization process.
• Figure 3: Overview of the SBDS: the student model (right) learns from the teacher model (left) through feature-level and logit-level losses.
• Figure 4: Illustration of spiking integration in the Discrete Mamba-2 (left) and FFN (right). The left and right sides correspond to the Discrete Mamba-2 and the Feed Forward Network in the student model shown in Figure \ref{['fig:distill']}, respectively.
• Figure 5: Total energy consumption of different decoders at varying sequence lengths.