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.

    Source: Introduction · p. 1

  • Method

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

    Source: Introduction · p. 2, Matched-Prefix Policy Anchoring · p. 4, Joint Training Objective · p. 5

  • 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.

    Source: Conclusion · p. 7

  • Takeaways & Limitations

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

    Source: Introduction · p. 2, Ablation Study · p. 7

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

    Source: Teacher-Guided Policy Correction · p. 4

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.Source: Introduction · p. 1
  • Fixed-corpus distillation mismatches autoregressive inference, where one changed token can alter every later self-generated prefix and expose recurrent spiking dynamics.Source: Introduction · p. 1
  • Paired-prefix controls identify output-policy mismatch with the ANN teacher and internal drift across spike rates, hidden states, and membrane potentials.Source: Introduction · p. 1
  • Vanilla OPD applies teacher supervision to self-generated prefixes but can undergo delayed rollout-feedback collapse because policy updates redirect later training distributions.Source: Introduction · pp. 1–2
  • SpikeOPD combines full-KL teacher correction, matched-prefix policy anchoring, and layerwise spike regularization to stabilize on-policy adaptation.Source: Introduction · p. 2
  • 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.Source: Introduction · p. 2

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.Source: Related Work · p. 2
  • On-policy language-model distillation builds on sequence-level distillation, scheduled sampling, and DAgger to learn beyond fixed corpus prefixes.Source: Related Work · p. 2
  • 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.Source: Related Work · p. 2

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

Vanilla OPD Is Insufficient for Stable Distillation

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.Source: Overview of SpikeOPD · p. 4, Joint Training Objective · p. 5

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.Source: Teacher-Guided Policy Correction · p. 4

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.Source: Figure 3 · p. 3, Teacher-Guided Policy Correction · p. 4

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.Source: Matched-Prefix Policy Anchoring · pp. 4–5

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.Source: Preliminaries · p. 4, Layerwise Spike Regularization · p. 5

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.Source: Performance Comparison with KD Baselines · p. 5
  • 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.Source: Performance Comparison with KD Baselines · p. 6

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.Source: Comparison with Post-Training Methods · p. 6

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%.Source: Ablation Study · p. 6
  • 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.Source: Ablation Study · pp. 6–7

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.Source: Initialization Robustness · p. 7

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.Source: Sensitivity Analysis · p. 7

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.Source: Conclusion · p. 7
  • 0.8, 1.7, and 2.9 percentage points: SpikeOPD improves knowledge-distillation performance across the three tested model scales.Source: Conclusion · p. 7
  • SpikeOPD does not increase analytical sparse-computation overhead or require auxiliary models at deployment time.Source: Conclusion · p. 7

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

Two-stage training

Stage 2 protocol

Stage 2 objective

Stage 2 protocol

  • Each update evaluates the frozen teacher, active SNN, and frozen reference on identical detached prefixes using complete vocabulary sums.Source: Update sequence · p. 11

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

Evaluation 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.Source: Diagnostic definitions · pp. 11–12
  • ∆spk(t) > 0 means SpikeOPD has a smaller spike-rate gap than Offline-only on the same stored prefixes and controls.Source: Diagnostic definitions · p. 12
  • 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.Source: Diagnostic definitions · p. 12

Objective properties

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

Reporting caveats

Efficiency model

Reporting caveats

Training boundary

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.Source: Implementation-level softmax-free… · p. 15
  • 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.Source: Implementation-level softmax-free… · pp. 14–15

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.Source: Implementation-level softmax-free… · p. 15

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.Source: Temporal fusion and teacher-to-spike… · p. 15

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.Source: Five aligned targets · pp. 15–16

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.Source: Surrogate-gradient propagation with… · pp. 16–17

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}.Source: G Reproducibility Checklist · p. 17

Related papers

Loading 2608.27857v1…