Spike-aware Bidirectional Distillation Strategy

• The SBDS methodology uses bidirectional KL divergence and feature-level pre-norm alignment to synchronize spiking student outputs with dense teacher distributions.
• It addresses challenges of low information density and temporal misalignment in spiking neural networks when integrated with modern SSMs or dense LLM teachers.
• Empirical evaluations show that incorporating ATMN, reverse KL, and pre-norm features markedly improves accuracy, reducing the gap to dense teacher performances.

The Spike-aware Bidirectional Distillation Strategy (SBDS) is a training paradigm designed for neural architectures with spiking components, aiming to bridge the information and temporal mismatches that arise when integrating Spiking Neural Networks (SNNs) with modern State Space Models (SSMs) or dense teacher models. SBDS promotes alignment between discrete, sparse spiking student models and their smooth, high-density teacher counterparts by leveraging both bidirectional distributional matching and feature-level supervision. While inspired by principles of biologically plausible learning and prior bidirectional distillation frameworks, SBDS is particularly tailored for large-scale, energy-efficient models such as those employing Module-aware Architecture Refinement (MAR) and Adaptive Ternary Multi-step Neurons (ATMN) (Cai et al., 29 Jan 2026, Lv et al., 24 Sep 2025).

1. Motivation and Background

Integrating spiking neurons in deep sequence models, such as LLMs with SSM-based backbones, offers significant energy efficiency due to temporal and activation sparsity. However, two inherent challenges impede direct application:

1. Low Information Density: Binary spiking neurons output only $\{0, 1\}$, discarding negative signals and producing minimal information per time step.
2. Temporal Misalignment: Dense SSM teachers produce continuous, token-level hidden states, whereas spiking students operate with multiple sub-token micro-steps, leading to bursty and coarse-grained temporal outputs.

Standard one-way knowledge distillation (KL divergence from teacher to student) is inadequate because it cannot address these mismatches; spiking students fail to approximate the dense teacher’s distribution, especially under sparse or bursty firing regimes (Cai et al., 29 Jan 2026).

SBDS was developed to provide richer supervisory signals that explicitly handle both information sparsity and desynchronized output timing by incorporating a bidirectional KL-based loss and feature-level pre-normalization alignment.

2. Core Methodology

SBDS employs a two-way distillation approach between a fixed dense teacher (e.g., a Llamba LLM) and a spiking student (e.g., MAR + ATMN + SSM):

• Forward (Teacher→Student): Standard KL divergence aligns the student’s probability distribution $q$ with the teacher’s softmax $p$ over token classes. This ensures the student's outputs cover the same support as the teacher, even if the firing rate is low.
• Reverse (Student→Teacher): A reverse KL term incentivizes the student to prioritize tokens where the teacher exhibits high confidence, compensating for the student’s output sparsity.
• Feature-level Pre-Norm Alignment: At each model layer, immediately after the first RMSNorm operation, the $\ell_2$ distance between the teacher and student’s normalized hidden states is minimized, aligning their internal representations despite differences in activation types.

At each training iteration, all loss computations and parameter updates are applied only to the student; the teacher remains frozen (Cai et al., 29 Jan 2026).

3. Mathematical Formulation

Let $T$ be the number of spiking micro-steps per token, $M$ the sequence length, $L$ the number of layers, and $D$ the vocabulary size. For time $t=0,\ldots,T-1$ and position $m=0,\ldots,M-1$:

Logit-level Bidirectional KL Loss:

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

where $p$ and $q$ are teacher and student softmax outputs; $\alpha, \beta$ control the weight of forward and reverse KL terms.

Feature-level Pre-Norm Loss:

$$\mathcal{L}_2(h^j, h^k) = \lVert \mathrm{PreNorm}(h^j) - \mathrm{PreNorm}(h^k) \rVert_2$$

with $h^j$ (teacher) and $h^k$ (student) denoting hidden states post-RMSNorm.

Total SBDS Loss:

$$\mathcal{L}_\mathrm{SBDS} = \frac{1}{T M} \sum_{t=0}^{T-1}\sum_{m=0}^{M-1} \mathcal{L}_1(p_{t,m} \Vert q_{t,m}) + \frac{1}{T L} \sum_{t=0}^{T-1} \sum_{\ell=0}^{L-1} \mathcal{L}_2(h^j_{t,\ell}, h^k_{t,\ell})$$

This composite loss (denoted as $\mathcal{L}_\mathrm{distill}$ in the source) drives both distributional and representational alignment (Cai et al., 29 Jan 2026).

4. Algorithmic Implementation

A typical SBDS training step consists of the following sequence:

1. Teacher Forward Pass: Compute fixed teacher activations and logits.
2. Student Forward Pass: Run the spiking student model (incorporating ATMN + SSM), producing per-time, per-position logits and hidden states.
3. Logit-level Loss: Calculate $\mathcal{L}_1$ over all time steps and positions.
4. Feature-level Loss: Compute $\mathcal{L}_2$ at every layer and time step.
5. Backward Pass and Update: Sum losses, backpropagate through the student, and update student parameters.

A high-level pseudocode excerpt from (Cai et al., 29 Jan 2026) is shown below:

for each minibatch of token sequences X do # Teacher forward (fixed) H_teacher, logits_teacher = TeacherModel(X) # Student forward (spiking) H_student, logits_student = StudentModel(X) # Compute bidirectional KL loss L_logit = 0 for t in 0..T-1, m in 0..M-1: p = softmax(logits_teacher[t,m]) q = softmax(logits_student[t,m]) L_logit += sum_k [α p[k] - β q[k]]*(log p[k] - log q[k]) L_logit /= (T*M) # Compute feature-level pre-norm loss L_feat = 0 for t in 0..T-1, ℓ in 0..L-1: hj = PreNorm(H_teacher[t,ℓ]) hk = PreNorm(H_student[t,ℓ]) L_feat += ||hj - hk||₂ L_feat /= (T*L) # Total loss, backprop Loss = L_logit + L_feat Loss.backward() optimizer.step() optimizer.zero_grad() end for

Key hyperparameters include $(\alpha, \beta)=(0.2, 0.7)$, ATMN unroll length $T$ (commonly $T=4$), and batch configurations to fit hardware constraints (Cai et al., 29 Jan 2026).

5. Integration with MAR and ATMN

Within the MAR pipeline, SBDS facilitates efficient training of SSM-based sequence models in which the dense FFNs are replaced by ATMN-enabled SNN components. SBDS leverages:

• Spike Outputs: The ternary spike output (from ATMN) forms the basis for the student’s softmax logits.
• Internal Membrane States: Feature alignment is enforced on internal membrane voltages after RMSNorm, which are more semantically aligned with the dense teacher than the spike bins themselves.

SBDS thus ensures that both the final distributional outputs and the semantic trajectory through the network’s layers remain well-matched between spiking student and dense teacher models, mitigating the discrepancies imposed by quantized temporal dynamics and sparse activation regimes (Cai et al., 29 Jan 2026).

6. Empirical Evaluation and Ablation Studies

Ablation and sensitivity studies demonstrate the practical impact of each SBDS component:

Configuration	Average Acc. (%)
Binary neurons + KL only	46.28
+ ATMN	55.20
+ Reverse-KL term	55.46
+ Pre-norm feature loss	57.20

The inclusion of ATMN spikes, reverse KL, and especially pre-norm alignment substantially reduces the performance gap to the dense teacher baseline (61.88%). Sensitivity to $(\alpha, \beta)$ hyperparameters shows peak accuracy with $(0.2, 0.7)$. Feature-level evaluation indicates pre-norm alignment is more effective than post-norm or combined pre/post (Cai et al., 29 Jan 2026).

Empirical results suggest that SBDS is a critical enabler for LLMs and sequence models with energy-efficient spiking modules, yielding robust performance under severe computational constraints.

7. Relation to Biologically Plausible Spike-based Distillation

Bidirectional distillation with spike-based signals also underpins biologically plausible learning paradigms, such as BSD (Lv et al., 24 Sep 2025). Analogous mechanisms appear:

• Bidirectional Pathways: BSD utilizes feedforward (stimulus→concept) and feedback (concept→stimulus) SNNs, enforcing alignment via local contrastive (ReCo) losses at each layer.
• Biological Constraints: BSD enforces asymmetric weights, local synaptic plasticity, unsigned learning signals, and simultaneous updates, avoiding global error signals and signed error propagation.
• Surrogate Gradients: Both SBDS and BSD rely on surrogate gradients for efficient learning with spiking non-linearities.

While SBDS is specialized for aligning SNN-SSM hybrids with dense LLMs, its foundational principle—a bidirectional, spike-aware transfer of information between teacher and student—is shared with biologically motivated frameworks such as BSD, highlighting a broader methodological convergence in energy-efficient and plausibility-constrained spiking systems (Cai et al., 29 Jan 2026, Lv et al., 24 Sep 2025).