SmoothSpike

Official code for SmoothSpike: Spiking Transformer with Learnable Hadamard Transformation.

SmoothSpike targets the spike saturation problem in spiking Transformers. In a bounded time window, a spiking neuron can emit at most T spikes, so distinct high-amplitude inputs may collapse to the same spike count. This causes information homogenization across layers and weakens fine-grained semantic discrimination.

The main idea is to smooth the pre-activation inputs of selected spiking neurons with orthogonal Hadamard-style transforms. The paper first motivates randomized Hadamard transforms, which spread input energy across channels and bound the maximum transformed coordinate with high probability. It then makes these transforms learnable by initializing them from randomized Hadamard matrices and projecting them to orthogonal matrices with Newton-Schulz matrix-sign iterations.

At inference time, the learned orthogonal transforms can be fused into adjacent linear weights, so the deployed model keeps the original spike-driven computation path without extra transform modules.

Requirements

• Python 3.9 or newer.
• PyTorch 2.0 or newer with CUDA. Install the CUDA build that matches your machine from pytorch.org; this is a prerequisite for fast-hadamard-transform.
• Install this repo's Python dependencies:

pip install -r requirements.txt

• Install fast-hadamard-transform, which provides the Hadamard kernels used by the project:

git clone https://github.com/Dao-AILab/fast-hadamard-transform.git
cd fast-hadamard-transform
pip install .

Repository Contents

• spikingbert_rot.py: SmoothSpike BERT model used by the pretrained checkpoints.
• spikingbert_rot_inf.py: inference model after per-layer rotation matrices are fused into linear weights.
• convertor.py: fuses the learned per-layer H2 and H3 rotations into the checkpoint weights for faster inference.
• spiking_pretrain_rot.py: masked-language-model pretraining entry for spikingbert_rot.BertForMaskedLM.
• finetune_spiking_rot.py: GLUE-style finetuning entry for spikingbert_rot.BertForSequenceClassification.
• hadamard_utils.py, utils.py: local helpers required by the model and training scripts.
• tokenizer_files/: local tokenizer/config files for offline loading.
• checkpoints/smoothspike-bert-base/: pretrained checkpoint directory.
• checkpoints/smoothspike-bert-base-fused/: fused inference checkpoint directory generated by convertor.py.

The cls.predictions.decoder.* tensors are tied weights and may be omitted from raw safetensors metadata. from_pretrained ties them correctly during loading.

Pretrained Weights

The pretrained unfused checkpoint is hosted on ModelScope:

kailai1104/SmoothSpike

Download it from the repository root so the checkpoint lands at the expected local path:

modelscope download \
  --model kailai1104/SmoothSpike \
  checkpoints/smoothspike-bert-base/model.safetensors \
  --local_dir .

Only the unfused pretrained checkpoint is hosted there. The fused inference checkpoint is generated locally with convertor.py.

Method Summary

SmoothSpike inserts learnable orthogonal transforms around selected spiking branches:

• H1 is shared across residual-connected branches to keep representation spaces aligned.
• H2 is applied to the value projection branch.
• H3 is a block-diagonal MLP transform with four equally sized blocks, reducing training cost.
• Query/key LIF neurons are left untransformed because they show milder saturation empirically.

The model uses a pre-norm architecture with RMSNorm. After the RMSNorm scale is absorbed into spiking thresholds, RMSNorm becomes scale-free and orthogonally equivariant. This property is what allows SmoothSpike transforms to be absorbed into linear weights during inference.

In the paper, SmoothSpike improves the Spikingformer GLUE average from 66.8 to 75.0, an 8.2 point gain, while preserving spike-driven efficiency. The same mechanism is implemented here for the BERT-style spiking language model.

Weight Fusion For Faster Inference

During training/pretraining, spikingbert_rot.py keeps the learnable transform parameters in the checkpoint. At inference, convertor.py folds the per-layer transforms into neighboring linear layers:

• input-side transform: H @ W -> W_fused
• output-side inverse transform: W @ H.T -> W_fused
• per-layer H2 and H3 parameters are removed after fusion
• global bert.H1 is kept because spikingbert_rot_inf.py still applies the embedding-side H1 and final encoder-output H1.T

Generate the fused checkpoint:

python convertor.py \
  --input checkpoints/smoothspike-bert-base \
  --output checkpoints/smoothspike-bert-base-fused \
  --config tokenizer_files

Expected conversion summary:

Saved fused weights to checkpoints/smoothspike-bert-base-fused/model.safetensors
Saved fusion report to checkpoints/smoothspike-bert-base-fused/fusion_report.json
Removed 24 keys
Kept bert.H1: True
Remaining per-layer H2/H3 keys: 0

Load the fused checkpoint with the inference model:

from transformers import BertConfig
from spikingbert_rot_inf import BertForMaskedLM

config = BertConfig.from_pretrained("tokenizer_files", local_files_only=True)
config.T = 4
config._attn_implementation = "eager"

model = BertForMaskedLM.from_pretrained(
    "checkpoints/smoothspike-bert-base-fused",
    config=config,
    local_files_only=True,
)
model.eval()

Do not load the fused checkpoint with spikingbert_rot.py; that model still expects explicit per-layer H2 and H3 parameters. Do not load the unfused checkpoint with spikingbert_rot_inf.py; that model expects the fused linear weights.

Pretraining Data

The pretraining script has two different data stages. The raw dataset loader is still present in spiking_pretrain_rot.py, but the expensive tokenization/chunking stage is currently commented out. The run that produced checkpoints/smoothspike-bert-base trained from a precomputed Hugging Face DatasetDict cache:

tokenized_datasets = load_from_disk("data/128_tokenized_data")
train_dataset = tokenized_datasets["train"]
eval_dataset = tokenized_datasets["validation"]

That cache is the important pretraining data artifact:

data/128_tokenized_data
  train:      137,271,688 examples
  validation:   1,627,657 examples

features:
  input_ids:            List[int32], length 128
  token_type_ids:       List[int8],  length 128
  attention_mask:       List[int8],  length 128
  special_tokens_mask:  List[int8],  length 128

There are no stored MLM labels in this cache. Labels are created online by:

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=args.mlm_probability)

The default --mlm_probability is 0.15, so masking is random and dynamic each time a batch is collated.

Raw Sources

The original pretraining run used these local raw datasets:

--dataset_name \
  data/raw/STORIES \
  data/raw/bookcorpus \
  data/raw/cc_news \
  data/raw/openwebtext \
  data/raw/wikipedia

Observed raw dataset shapes on this server:

STORIES      train 945,354      validation 946      test 947       column: text
bookcorpus   train 74,004,228                                    column: text
cc_news      train 708,241                                      columns include text
openwebtext  train 20,610                                       column: text
wikipedia    train 6,407,814                                    columns include text

For datasets without a validation split, the loader in spiking_pretrain_rot.py uses the first min(1000, len(train)) raw rows as that dataset's validation contribution. Then all train splits are concatenated and all validation contributions are concatenated.

Tokenization And Chunking

The commented preprocessing block in spiking_pretrain_rot.py shows how the cache was intended to be made:

1. Pick the text column:

text_column_name = "text" if "text" in raw_datasets["train"].column_names else raw_datasets["train"].column_names[0]

2. Tokenize with the BERT tokenizer and keep special_tokens_mask:

def tokenize_function(examples):
    return tokenizer(examples[text_column_name], return_special_tokens_mask=True)

3. Concatenate tokenized documents in each map batch and split into fixed-length blocks:

def group_texts(examples):
    concatenated_examples = {k: list(chain(*examples[k])) for k in examples}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    total_length = (total_length // max_seq_length) * max_seq_length
    return {
        k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
        for k, t in concatenated_examples.items()
    }

For the copied checkpoints, max_seq_length was 128, which is why the cache is named 128_tokenized_data and every example is a 128-token block.

Regenerating The Cache

If you need to regenerate the dataset cache, use the same tokenizer/config and save the processed DatasetDict before training:

from itertools import chain
from datasets import DatasetDict, concatenate_datasets, load_from_disk
from transformers import AutoTokenizer

raw_paths = [
    "data/raw/STORIES",
    "data/raw/bookcorpus",
    "data/raw/cc_news",
    "data/raw/openwebtext",
    "data/raw/wikipedia",
]

tokenizer = AutoTokenizer.from_pretrained("tokenizer_files", local_files_only=True)
max_seq_length = 128
all_datasets = [load_from_disk(path) for path in raw_paths]
raw_datasets = DatasetDict({
    "train": concatenate_datasets([d["train"] for d in all_datasets]),
    "validation": concatenate_datasets([
        d["validation"] if "validation" in d else d["train"].select(range(min(1000, len(d["train"]))))
        for d in all_datasets
    ]),
})

column_names = raw_datasets["train"].column_names
text_column_name = "text" if "text" in column_names else column_names[0]

def tokenize_function(examples):
    return tokenizer(examples[text_column_name], return_special_tokens_mask=True)

def group_texts(examples):
    concatenated = {k: list(chain(*examples[k])) for k in examples}
    total_length = len(concatenated[list(examples.keys())[0]])
    total_length = (total_length // max_seq_length) * max_seq_length
    return {
        k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
        for k, t in concatenated.items()
    }

tokenized = raw_datasets.map(
    tokenize_function,
    batched=True,
    num_proc=32,
    remove_columns=column_names,
    desc="Running tokenizer on every text",
)

tokenized = tokenized.map(
    group_texts,
    batched=True,
    num_proc=32,
    desc="Grouping texts in chunks of 128",
)

tokenized.save_to_disk("data/128_tokenized_data")

After the cache exists, spiking_pretrain_rot.py trains from load_from_disk(args.tokenized_dataset_path). Use --tokenized_dataset_path data/128_tokenized_data or pass your own processed dataset cache.

Pretraining

The original run that produced checkpoints/smoothspike-bert-base used:

python spiking_pretrain_rot.py \
  --dataset_name data/raw/STORIES \
                 data/raw/bookcorpus \
                 data/raw/cc_news \
                 data/raw/openwebtext \
                 data/raw/wikipedia \
  --model_name_or_path bert-base-uncased \
  --per_device_train_batch_size 64 \
  --per_device_eval_batch_size 64 \
  --learning_rate 2e-4 \
  --max_train_steps 800000 \
  --num_warmup_steps 5000 \
  --output_dir ./checkpoints/smoothspike-bert-base \
  --max_seq_length 128 \
  --tokenized_dataset_path data/128_tokenized_data \
  --checkpointing_steps 50000 \
  --preprocessing_num_workers 32 \
  --with_tracking \
  --report_to wandb

The logged distributed setup was 8 GPUs, so the effective train batch size was 64 * 8 = 512 because gradient accumulation was 1.

Finetuning

Use finetune_spiking_rot.py for sequence classification finetuning. The script imports BertForSequenceClassification from spikingbert_rot.py, so it uses the SmoothSpike model definition and can initialize from the pretrained checkpoint directory.

Example:

python finetune_spiking_rot.py \
  --task_name sst2 \
  --model_name_or_path checkpoints/smoothspike-bert-base \
  --per_device_train_batch_size 32 \
  --per_device_eval_batch_size 32 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --output_dir outputs/sst2

Citation

@inproceedings{zhou2026smoothspike,
  title = {SmoothSpike: Spiking Transformer with Learnable Hadamard Transformation},
  author = {Zhou, Zijian and Wei, Wenjie and Liang, Yu and Li, Jialin and Belatreche, Ammar and Cao, Honglin and Wang, Shuai and Zhang, Malu and Yang, Yang and Li, Haizhou},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
  year = {2026}
}