PaperlayerView exact revision on arXiv

Source-linked AI summary

SpikeOPD: Stable On-Policy Distillation for Autoregressive Spiking Language Models

Enqiao Lu, Xingrui Yu, Yiwei Fu, Zhenglin Wan, Pengfei Zhou, Wangbo Zhao, Muqing Jian, Xueyi Zhang, Yang You, Ivor Tsang

arXiv:2608.27857v1cs.AI

TL;DR

Fixed-prefix ANN-to-SNN distillation leaves autoregressive SNNs exposed to prefix-source mismatch and internal spiking drift. SpikeOPD adapts on self-generated prefixes with teacher correction, reference-policy anchoring, and spike regularization, improving accuracy across three scales while preserving sparse computation. The controlled evidence also shows that teacher-only Vanilla OPD can suffer delayed rollout-feedback collapse.

  • Problem

    Fixed corpus prefixes used for ANN-to-SNN distillation differ from self-generated prefixes used during autoregressive inference, producing output-policy mismatch and internal spiking-dynamics drift.

  • Method

    SpikeOPD uses full-KL teacher correction on self-generated prefixes, matched-prefix anchoring to a frozen reference SNN, and layerwise spike regularization.

  • Results

    SpikeOPD improves corresponding KD SNN average accuracy by 0.8, 1.7, and 2.9 points at 0.125B, 0.35B, and 1.3B, respectively, without increasing analytical sparse-compute overhead.

  • Takeaways & Limitations

    On-policy coverage alone does not ensure stable adaptation; constraining policy movement and spiking dynamics supports stable autoregressive distillation.

  • Takeaways & Limitations

    Teacher-only full-KL Vanilla OPD can redirect later rollouts and undergo delayed rollout-feedback collapse, motivating an explicit policy-movement constraint.

Abstract

from arXiv · show

Spiking neural networks (SNNs) offer a path to energy-efficient language modeling through sparse encoding and event-driven computation, but training capable spiking language models from scratch remains difficult. A practical alternative is ANN-to-SNN migration through knowledge distillation (KD), where a pretrained artificial neural network (ANN) teacher supervises an SNN student. Existing migration approaches distill on fixed corpus prefixes, whereas autoregressive inference conditions on self-generated prefixes, creating prefix-source mismatch. It manifests as output-policy mismatch with the ANN teacher and internal spiking-dynamics drift between self-generated and matched corpus prefixes. On-policy distillation (OPD) offers a natural way to mitigate both manifestations by continuing teacher supervision on self-generated prefixes. We evaluate a teacher-only full-KL variant, Vanilla OPD, via a controlled stress test and observe it may suffer from delayed rollout-feedback collapse. This result shows that on-policy coverage alone does not ensure stable adaptation. Motivated by these findings, we propose SpikeOPD, a stable on-policy distillation framework for autoregressive SNNs that learns from self-generated prefixes while maintaining rollout stability. It applies full-KL teacher correction to reduce output-policy mismatch, while matched-prefix policy anchoring constrains policy departure from the frozen reference SNN on the same prefixes. Layerwise spike regularization further limits firing-rate deviations during on-policy adaptation. Across three model scales, SpikeOPD improves average accuracy over the corresponding KD SNNs by 0.8, 1.7, and 2.9 points at 0.125B, 0.35B, and 1.3B, respectively, while preserving their sparse-compute profiles.

Introduction

Autoregressive SNN migration must address the shift from fixed corpus prefixes during distillation to self-generated prefixes during inference. SpikeOPD combines teacher correction with policy and spike-dynamics constraints, improving accuracy across model scales while preserving sparse computation.

  • Large generative SNNs remain difficult to train from scratch, motivating transfer from pretrained ANN language models without losing language ability or sparse activity.
  • Fixed-corpus distillation mismatches autoregressive inference, where one changed token can alter every later self-generated prefix and expose recurrent spiking dynamics.
  • Paired-prefix controls identify output-policy mismatch with the ANN teacher and internal drift across spike rates, hidden states, and membrane potentials.
  • Vanilla OPD applies teacher supervision to self-generated prefixes but can undergo delayed rollout-feedback collapse because policy updates redirect later training distributions.
  • SpikeOPD combines full-KL teacher correction, matched-prefix policy anchoring, and layerwise spike regularization to stabilize on-policy adaptation.
  • 0.8, 1.7, and 2.9 average points improve over corresponding KD SNNs at 0.125B, 0.35B, and 1.3B, respectively, across eight tasks.

Related Work

Spiking language-model research combines surrogate-gradient training, ANN-to-SNN conversion, and teacher-student transfer. On-policy distillation extends supervision beyond fixed corpus prefixes through varied divergences, rollout sources, and auxiliary controls.

  • Deep SNNs are commonly developed through surrogate-gradient training, ANN-to-SNN conversion, or teacher-student transfer.
  • On-policy language-model distillation builds on sequence-level distillation, scheduled sampling, and DAgger to learn beyond fixed corpus prefixes.
  • Recent methods vary divergence and rollout design, including reverse KL, mixed on- and off-policy samples, generalized OPD, entropy-aware selection, stored trajectories, and representation alignment.

Empirical Analysis

Controlled paired-prefix diagnostics show that self-generated prefixes amplify output-policy and internal spiking mismatches, while stress tests reveal delayed collapse in teacher-only Vanilla OPD. SpikeOPD reduces these discrepancies and remains stable through explicit policy and dynamics controls.

  • Prefix-Source Mismatch in Knowledge Distillation: RQ1 pairs self-generated and corpus prefixes by prompt, position, and length using a common fixed prefix bank.
  • Prefix-Source Mismatch in Knowledge Distillation: Figure 1 compares teacher-to-SNN KL on self-generated versus corpus prefixes and reports their prompt-internal Self−Corpus gap.
  • Prefix-Source Mismatch in Knowledge Distillation: Figure 1(A) shows higher KL for KD than SpikeOPD on both self-generated and corpus prefixes.
  • Prefix-Source Mismatch in Knowledge Distillation: GKD(t) grows from near zero to about 0.9, whereas SpikeOPD remains smaller, reducing both absolute teacher KL and the paired prefix-source effect.
  • Prefix-Source Mismatch in Knowledge Distillation: Figure 2 measures spike-rate, hidden-state, and aggregate membrane-potential gaps under matched prompts, positions, and lengths; smaller values are better.
  • Prefix-Source Mismatch in Knowledge Distillation: The KD gap grows with rollout position; aggregated hidden-state and membrane-potential gaps decrease from 0.102 and 3.46 for KD to 0.095 and 2.67 for SpikeOPD.
  • Vanilla OPD Is Insufficient for Stable Distillation: Vanilla OPD uses only full-KL teacher supervision on self-generated prefixes, while SpikeOPD remains lower in both output and internal diagnostics.
  • Vanilla OPD Is Insufficient for Stable Distillation: Ten matched-seed Vanilla OPD runs eventually collapse after repeated text shifts later training contexts, despite early updates remaining close to KD.

Method

SpikeOPD adapts a KD-trained spiking language model on self-generated prefixes using teacher correction, matched-prefix policy anchoring, and layerwise spike regularization. These components target output-policy mismatch, policy drift, and internal firing-rate instability while updating only the active SNN.

  • SpikeOPD overview: SpikeOPD combines teacher-guided policy correction, matched-prefix anchoring, and layerwise spike regularization to address prefix-source mismatch and stabilize internal dynamics.The method operates on self-generated prefixes and updates only the active SNN parameters.
  • Teacher correction: The active SNN samples self-generated prefixes, while the frozen ANN teacher evaluates those same prefixes and supplies full-distribution supervision.Gradients pass through the active SNN policy and spike paths, while sampled token identities are detached.
  • Stability motivation: Vanilla OPD can enter delayed rollout-feedback collapse because policy corrections alter later prefixes and training distributions.The controlled stress test reports collapse after early updates remain close to KD.
  • Matched-prefix policy anchoring: The frozen reference SNN anchors the active model on identical self-generated prefixes, preserving transferred language behavior without blocking teacher-guided correction.The reference is a frozen copy of the initial KD checkpoint.
  • Spike regularization: Layerwise regularization targets selected layers because output matching alone can leave firing rates free to drift or allow layerwise compensation.The method penalizes deviations from both a viable firing-rate interval and the frozen reference SNN rates.

Experiments

Experiments show that SpikeOPD improves retained language accuracy across model scales and controlled post-training comparisons while preserving sparse execution and rollout stability. Ablations and sensitivity studies indicate that stable adaptation depends on jointly controlling teacher correction, policy movement, and internal spiking dynamics.

  • Multi-scale validation: The eight-task averages are 35.5±0.02%, 40.0±0.19%, and 44.3±0.12%, with firing rates of 0.17 ± 0.006, 0.17 ± 0.000, and 0.19 ± 0.010.Each uncertainty is the sample standard deviation across three runs.
  • Multi-scale validation: SpikeOPD improves the KD average by 0.8, 1.7, and 2.9 points at 0.125B, 0.35B, and 1.3B, respectively.The gains become more pronounced over the tested scale range.
  • Comparison with Post-Training Methods: SpikeOPD is the only compared method that raises the 0.125B KD average from 34.7% to 35.5% while keeping rollout statistics close to the initial checkpoint.The controlled comparison uses one checkpoint and a 500-update budget.
  • Ablation Study: The full configuration achieves the highest average accuracy of 35.5%, while removing any component degrades performance.Removing matched-prefix policy anchoring produces the largest decline, from 35.5% to 31.7%.
  • Ablation Study: Matched-prefix policy anchoring is the primary stability constraint because unconstrained policy movement can amplify early deviations in subsequent prefixes.The frozen reference limits this feedback on the same visited prefixes without blocking teacher correction.
  • Initialization Robustness: SpikeOPD improves all three matched offline checkpoints by 0.6–0.8 points, although initialization maturity still determines absolute performance.The result indicates that offline alignment and on-policy adaptation provide complementary benefits.
  • Sensitivity Analysis: Stable on-policy adaptation benefits from bounded self-generated context and selective internal constraints, with intermediate optima at K = 32 and ∆L = 3.Short rollouts provide insufficient exposure, while excessively long rollouts may accumulate feedback from the evolving policy.

Conclusion

SpikeOPD combines full-KL teacher correction, prefix-matching policy anchoring, and layer-wise spike regularization on self-generated prefixes. Across three scales, it improves knowledge-distillation performance without increasing analytical sparse-computation overhead or requiring auxiliary deployment models.

  • SpikeOPD combines full-KL teacher correction, prefix-matching strategy anchoring, and layer-wise spike regularization on self-generated prefixes.
  • 0.8, 1.7, and 2.9 percentage points: SpikeOPD improves knowledge-distillation performance across the three tested model scales.
  • SpikeOPD does not increase analytical sparse-computation overhead or require auxiliary models at deployment time.

Models, data, and two-stage training

The study evaluates matched OPT teachers and causal BiSpikCLM students across three scales using a two-stage training pipeline. Stage 2 adapts an active SNN on self-generated prefixes with fixed teacher, reference, and activity-stability terms.

  • Models and data: The experiments use matched OPT teachers and causal BiSpikCLM students at 125M, 0.35B, and 1.3B parameters.
  • Two-stage training: Stage 1 performs offline spike-aware ANN-to-SNN distillation for 4,000 updates, after which the checkpoint initializes the active Stage 2 SNN.
  • Stage 2 protocol: Stage 2 samples 480-token FineWeb prompts and generates K = 32 tokens at temperature 1.0 before scoring visited prefixes.
  • Stage 2 objective: The default regularizer monitors layers Sspk = {3, 6, 9, 12}, while the layer-sensitivity study changes this set explicitly.
  • Stage 2 objective: The objective combines full-KL teacher correction, matched-prefix policy anchoring, and layerwise spike regularization with fixed weights throughout each run.
  • Stage 2 protocol: Each update evaluates the frozen teacher, active SNN, and frozen reference on identical detached prefixes using complete vocabulary sums.

Controlled baselines and evaluation

The evaluation combines controlled post-training baselines, zero-shot benchmark accuracy, generation-stability metrics, and matched-prefix diagnostics. It also establishes exact full-vocabulary KL evaluation while limiting gradient claims to a detached-prefix snapshot objective.

  • Controlled baselines: Controlled 125M comparisons share the same checkpoint, temporal setting, and 500-update budget across Offline SpAD, Continued SpAD, and SFT.
  • Controlled baselines: Lightning OPD replays cached trajectories, while EOPD switches KL direction using teacher entropy and omits matched-prefix reference-policy regularization.
  • Evaluation metrics: Accuracy is reported on eight zero-shot benchmarks with their unweighted average, alongside adjacent repetition, Dist.-4, and maximum-token-run metrics.
  • Diagnostics: Paired-prefix diagnostics compare self-generated prefixes with length-matched corpus prefixes across spike rates, hidden states, membrane potentials, and teacher–student KL.
  • Diagnostics: ∆spk(t) > 0 means SpikeOPD has a smaller spike-rate gap than Offline-only on the same stored prefixes and controls.
  • Diagnostics: 0.102, 0.100, and 0.095: hidden-state gaps are reported for Offline-only, Vanilla OPD, and SpikeOPD, respectively; membrane-potential gaps are 3.46, 3.30, and 2.67.
  • Objective properties: The complete teacher correction evaluates every vocabulary coordinate exactly, without vocabulary truncation or sampled teacher tokens.
  • Objective properties: The trajectory-average gradient is unbiased for a detached-prefix snapshot objective, not a global convergence guarantee for the nonconvex SNN.

E Analytical Efficiency Reporting

Analytical efficiency is estimated with an operation-based 45-nm energy model that combines dense MACs and spike-driven ACs. The reported values are arithmetic proxies rather than wall-plug or on-chip measurements.

  • Efficiency model: The efficiency protocol estimates operations and energy under a 45-nm model rather than measuring wall-plug or on-chip power.
  • Efficiency model: Dense MACs cover embedding and LM-head computations, while SFSA and SFFN blocks use spike-driven AC counts.
  • Reporting caveats: The displayed aggregate Rate is rounded and cannot be multiplied by a single FLOP total to reconstruct operations or energy.
  • Efficiency model: The aggregate OPs and energy calculations combine dense MAC counts with spike-driven AC counts that have different energy constants.
  • Reporting caveats: Reported OPs, Rate, and Energy should be interpreted as arithmetic proxies under the shared analytical boundary, not hardware measurements.
  • Training boundary: Stage 1 initialization is implementation detail for the retained checkpoint; SpikeOPD discards the Stage 1 losses during Stage 2.

Implementation-level softmax-free spiking attention

The implementation uses stateful, softmax-free spiking attention with explicit causal masking and spike-domain teacher targets. Alignment losses and finite-horizon surrogate-gradient analysis support offline training while preserving the student’s temporal dynamics.

  • Attention execution: Binary queries and keys produce integer coincidence counts, which pass through the attention neuron before masking, spike conversion, and output projection.The public path uses unscaled coincidence counts and performs no softmax.
  • Attention execution: The public execution path avoids a separate 1/√dh operation by using unscaled coincidence counts and a corresponding membrane-scale parameterization.The scaled and unscaled forms are equivalent only after transforming the complete membrane scale.
  • Attention supervision: Causal, padding, and target-validity masks exclude future entries, padding locations, and positions without next-token supervision.Teacher attention maps are paired with student layers and heads through depth- and head-mapping functions.
  • Teacher representation: Teacher tensors are converted into empirical spike-rate targets by repeating each scalar as constant current through the student’s subtractive-reset LIF dynamics.At Ts = 4, this target is a finite-horizon proxy rather than an assumed asymptotic equality with ANN values.
  • Training objectives: The objective combines embedding, attention, feature, and token-level supervision using explicit alignment branches, width maps, LayerNorm, and soft- and hard-token losses.For shape-matched models, width maps reduce to identities while LayerNorm remains in the continuous alignment transform.
  • Gradient propagation: With subtractive reset, the temporal Jacobian retains reset derivatives; for Ts = 4, temporal amplification is finite and bounded by max(1, κ̄^3).The eligibility trace is likewise bounded by MX(1 + κ̄ + κ̄^2 + κ̄^3) without requiring contraction.

G Reproducibility Checklist

The reproducibility record fixes the initialization and adaptation settings, including tokenization, surrogate sharpness, gradient clipping, KL training, simulation length, anchoring, and monitored layers.

  • G Reproducibility Checklist: Stage 1 uses shared tokenization, explicit alignment masks, ξ = 2, and global gradient clipping at 0.7; Stage 2 uses full-vocabulary KL with Ts = 4 and ρ = 1.The default spike-regularization layer set is {3, 6, 9, 12}.
Loading 2608.27857v1…