Module-aware Architecture Refinement (MAR)

• Module-aware Architecture Refinement (MAR) is a unified framework that refines large language model architectures through linear-time sequence modeling, activation sparsification, and spike-aware bidirectional distillation.
• It replaces quadratic self-attention with state space models and reduces dense FFN energy use using adaptive spiking neurons, achieving substantial energy savings while recovering dense model accuracy.
• Extensive evaluations show that MAR closely approaches dense model performance with lower energy requirements, making it ideal for on-device and resource-constrained applications.

Module-aware Architecture Refinement (MAR) is a unified framework for constructing energy-efficient LLMs that combine linear-time sequence modeling via state space models (SSMs), activation sparsification through adaptive spiking neurons, and spike-aware bidirectional distillation. MAR targets the dominant sources of compute and energy consumption in contemporary neural architectures—quadratic-complexity self-attention and dense feed-forward networks (FFNs)—by refining their implementation at the module level without sacrificing model expressivity or performance. Extensive evaluations demonstrate that MAR restores a large fraction of dense model accuracy while significantly lowering inference energy requirements, outperforming other efficient models both in scale-matched and parameter-larger regimes (Cai et al., 29 Jan 2026).

1. Motivation and Architectural Targets

The principal inefficiencies in standard Transformer-derived LLMs arise from (i) the $O(N^2)$ complexity of self-attention, which impedes scalability to longer sequences, and (ii) the dense FFN sublayer, which, in practice, often dominates energy expenditures for non-extreme sequence lengths (e.g., ≤5,000 tokens). Direct substitutions such as low-rank or quantized FFNs frequently induce excessive accuracy loss. The MAR framework was developed to eliminate the quadratic bottleneck of attention entirely via linear-time models and to drastically reduce FFN costs by sparsifying activations with negligible performance degradation (Cai et al., 29 Jan 2026).

2. Two-stage Refinement Pipeline

MAR applies distinct module-level interventions in a staged architecture:

Stage 1: Linear-Time Sequence Modeling with SSMs

Self-attention layers are replaced by the discrete Mamba-2 SSM modules introduced in the Llamba architecture. These modules model an input sequence $\{u_t\}_{t=0}^{T-1}$ using the recurrence:

$x_t = Ax_{t-1} + Bu_t; \quad y_t = Cx_t + Du_t,$

with learned matrices $A,B,C,D$, and hidden state $x_t \in \mathbb{R}^d$. This approach yields $O(Nd^2)$ computational and memory cost for sequence length $N$ and hidden size $d$, avoiding quadratic scaling in $N$ (Cai et al., 29 Jan 2026).

Stage 2: Activation Sparsification via Spiking Neurons

Sparse computation is achieved by interposing spiking neurons prior to each linear projection—four such placements per decoder layer: inputs/outputs for both SSM and FFN modules. These neurons convert pre-activations to spike trains, so subsequent matrix operations require only accumulate (AC) steps rather than dense multiply-accumulate (MAC), reducing per-operation energy from 4.6 pJ (MAC) to 0.9 pJ (AC), as detailed in the adopted energy model (Cai et al., 29 Jan 2026).

3. Spiking Neuron Design and Information Preservation

Naively introducing binary (0/1) spiking neurons poses two challenges: (a) reduced information density due to infrequent spike emissions, and (b) temporal misalignment between SSMs’ dynamics and discrete spikes. MAR addresses these with the Adaptive Ternary Multi-step Neuron (ATMN):

Adaptive Ternary Multi-step Neuron (ATMN):

• Extends spiking neuron outputs to the ternary domain $s_t \in \{-1,0,1\}$.
• Utilizes an adaptive, neuron-wise learnable threshold $V_{\mathrm{adaptive}} = e^a \ge 0$ (with $a$ learned).
• Stores post-spike residual membrane potential $u_t$, sustaining temporal continuity across activations.

The ATMN’s enhanced representation and learnable thresholds significantly increase per-spike information throughput and firing adaptability, facilitating seamless integration with the continuous-time dynamics of SSMs (Cai et al., 29 Jan 2026).

4. Spike-aware Bidirectional Distillation

Transferring task performance from a dense teacher (Llamba) to a sparsified, spiking MAR student is complex due to spike train burstiness and altered representational geometry. MAR introduces the Spike-aware Bidirectional Distillation Strategy (SBDS):

• Logit-level bidirectional loss: Blends the standard $\mathrm{KL}$ divergence with a reversed term balancing the teacher’s and student’s predictive distributions, parametrized by $(\alpha, \beta) = (0.2, 0.7)$ for best empirical performance:

$$\mathcal{L}_1(p\Vert q) = \sum_{k=0}^{D-1} [\alpha p(k) - \beta q(k)] [\log p(k) - \log q(k)]$$

• Feature-level pre-normalization alignment: Matches hidden activations immediately after RMSNorm, minimizing $L_2$ discrepancy between teacher and student representations:

$$\mathcal{L}_2(h^j_{t,l}, h^k_{t,l}) = \|\mathrm{PreNorm}(h^j_{t,l}) - \mathrm{PreNorm}(h^k_{t,l})\|_2$$

• Overall distillation objective: Averages logit and feature losses across time, token, and layer dimensions.

Ablation results indicate each distillation augmentation—ternary neuron (ATMN), bidirectional logit loss, and pre-norm feature alignment—substantially contributes to performance restoration (Cai et al., 29 Jan 2026).

5. Experimental Assessment and Benchmarking

The Llamba-1B (Mamba-2, 1.4 B parameters) serves as the foundation for MAR model construction, trained on 7 B tokens (GenQA, OpenHermes 2.5, InfinityInstruct) for one epoch. Performance is evaluated via zero-shot accuracy on PIQA, BoolQ, Winogrande, HellaSwag, ARC-Easy, and ARC-Challenge. Energy costs adopt estimates from Horowitz (2014).

Model	Size	Spiking?	Avg. Zero-Shot Acc.
LLaMA	1.3B	No	61.80%
Llamba (teacher)	1.4B	No	61.88%
Bi-Mamba	1.3B	No	49.38%
SmoothQuant	1.3B	No	54.93%
TinyLLaMA	1.3B	No	55.91%
SpikeLLM	7.0B	Yes	52.48%
MAR (Ours)	1.4B	Yes	57.20%

MAR exhibits substantial parameter-efficiency, matching or outperforming non-spiking models of similar scale and exceeding a 7 B parameter spiking LLM baseline (SpikeLLM). Ablations demonstrate the necessity of each architectural and training refinement: incorporating ATMN (from 46.28% to 55.20%), adding reverse KL (55.46%), and pre-norm alignment (57.20%). Feature alignment post-normalization was found suboptimal (Cai et al., 29 Jan 2026).

Energy consumption measurements confirm that MAR’s total energy scales much more favorably with sequence length than LLaMA and Llamba, achieving up to 30–40% lower energy use at longer contexts (≤4,000 tokens).

6. Design Insights and Practical Implications

The empirical results signify that module-level interventions—specifically, integrating linear-time SSMs in lieu of self-attention, augmenting FFNs with adaptive spiking neurons, and applying targeted bidirectional spike-aware distillation—enable energy reductions without incurring major accuracy penalties. Key mechanisms include: (1) $O(N)$ sequence modeling removing attention bottlenecks; (2) ternary, adaptively-thresholded spikes preserving information density; (3) pre-norm, reverse KL-based distillation addressing temporal and representational mismatches. This composition supports the creation of deployable, low-power LLMs suitable for on-device or resource-constrained applications, achieving full or near-full performance under strict energy and latency constraints (Cai et al., 29 Jan 2026).

7. Comparative Context and Significance

Module-aware Architecture Refinement advances the field of efficient language modeling by demonstrating that principled, module-targeted interventions can closely approach dense model accuracy in much more energy- and resource-impoverished environments. Benchmarks confirm that MAR’s performance is robust to sequence length scaling—reflecting the practical gains from linearized sequence processing and sparse activation protocols—and that it outperforms both non-spiking and larger spiking competitors when parameter count and data regime are held constant. These findings indicate that MAR constitutes a substantive framework for constructing the next generation of practical, high-performance, resource-efficient LLMs (Cai et al., 29 Jan 2026).