KO | EN

GitLite

Minimalist Open Source Repository Search Engine

gitlite — search

#dsh #dsh-plugin #typescript #ai-agent #developer-tools #claude-code #codex #deepseek-harness #ai-agents #react #mcp #rust

SDM

★ 52 Open GitHub ↗

[IEEE Trans. AI 2024] Spiking Diffusion Models

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<h1 align="center">Spiking Diffusion Models
 
(IEEE Transactions on Artificial Intelligence)</h1>
 
 Jiahang Cao1*,
 Hanzhong Guo2*,
 Ziqing Wang3*,
 Deming Zhou1,
 Hao Cheng1,
 Qiang Zhang1,
 Renjing Xu1†

 
 
 1The Hong Kong University of Science and Technology (Guangzhou)
 
 2The Hong Kong University 
 3Northwestern University 
 
 
 <a href='https://arxiv.org/pdf/2408.16467'>
 <img src='https://img.shields.io/badge/Paper-PDF-red?style=flat&logo=arXiv&logoColor=red' alt='Paper PDF'>
 </a>
 <a href='https://arxiv.org/pdf/2408.16467' style='padding-left: 0.5rem;'>
 <img src='https://img.shields.io/badge/Proceeding-HTML-blue?style=flat&logo=Google%20chrome&logoColor=blue' alt='Proceeding Supp'>
 </a>

This work **SDM** is an extended version of [**SDDPM**](https://github.com/AndyCao1125/SDDPM). We introduce several key improvements:

- **A New Family of Spiking-based Diffusion Models**：This work extends applicability to a wider array of diffusion solvers, including but not limited to DDPM, DDIM, Analytic-DPM and DPM-Solver.
- **Biologically Inspired Temporal-wise Spiking Mechanism (TSM)**: Inspired by biological processes that the neuron input at each moment experiences considerable fluctuations rather than being predominantly controlled by fixed synaptic weights, this module enables spiking neurons to capture more dynamic information. The TSM module can be integrated with existing modules (proposed by SDDPM) to further improve the image generation quality.
- **ANN-SNN Conversion for SDM**: To the best of our knowledge, we make the *first attempt* to utilize an ANN-SNN approach for implementing spiking diffusion models, complete with theoretical foundations.

## TODO Lists
- [x] **[Updated on August 29 2024]** Release SDM direct training (finetuning) codes!
- [x] **[Updated on September 19 2024]** Release SDM ANN-to-SNN conversion codes!

## Requirements
Please see [**SDDPM**](https://github.com/AndyCao1125/SDDPM).

## TSM Finetune
- Pretrained weight of SDDPM  [(this link)](https://drive.google.com/drive/folders/1ATfGPQqlqklDuWtAtll-I9_7bMJuqEmy?usp=drive_link) with spiking timestep=4, CIFAR-10.

Here we provide an example code to finetune the SDM models by inheriting the weights obtained from SDDPM pre-training:

```shell
from TSM import Spk_UNet_TSM

... (First, pretrain the standard SNN UNet)

pretrained_model = Spk_UNet(
      T=args.T, ch=args.ch, ch_mult=args.ch_mult, attn=args.attn,
      num_res_blocks=args.num_res_blocks, dropout=args.dropout, timestep=args.timestep, img_ch=args.img_ch)
  
# Load model
ckpt = torch.load(os.path.join('/your/pretrained_checkpoint'))
pretrained_model.load_state_dict(ckpt['net_model'], strict=True)
pretrained_dict = pretrained_model.state_dict()

net_model = Spk_UNet_TSM(
    T=args.T, ch=args.ch, ch_mult=args.ch_mult, attn=args.attn,
    num_res_blocks=args.num_res_blocks, dropout=args.dropout, timestep=args.timestep, img_ch=args.img_ch)

model_dict = net_model.state_dict()
new_state_dict = OrderedDict()

for name,para in pretrained_dict.items():
    if name in model_dict:
        new_state_dict[name] = para
    
    elif 'conv' and 'weight' in name:
        head = name[:-7]
        new_name = head + '.tsmconv.weight'
        new_state_dict[new_name] = para
        
    elif 'conv' and 'bias' in name:
        head = name[:-5]
        new_name = head + '.tsmconv.bias'
        new_state_dict[new_name] = para
       
net_model.load_state_dict(new_state_dict, strict=False)

print(f'-------Successfully inherit pretrained weights-------')

...(Next, finetune the TSM SDM with the same training code from SDDPM)
```
Please set the number of the finetuning steps as 200k (if 8 GPUs) to 500k (when the number of GPUs is small).

## Sample
Example codes for sampling the images with DDIM solver.

The checkpoint of SDM with `snn_timesteps=8` in CIFAR-10 is released. You can download the checkpoint through [this link](https://drive.google.com/file/d/1Z38wxR-olP_dlL0b-Wa8fwjM611hoO65/view?usp=drive_link).

```shell
cd SDM
CUDA_VISIBLE_DEVICES=0 python sample.py
```

## ANN-to-SNN Conversion
Please see [ANN2SNN](https://github.com/AndyCao1125/SDM/tree/master/ANN2SNN).

## Citation

If you find our work useful, please consider citing:

```
@article{cao2024spiking,
  title={Spiking Diffusion Models},
  author={Cao, Jiahang and Guo, Hanzhong and Wang, Ziqing and Zhou, Deming and Cheng, Hao and Zhang, Qiang and Xu, Renjing},
  journal={IEEE Transactions on Artificial Intelligence},
  year={2024},
  publisher={IEEE}
}
```

## Acknowledgements & Contact
We thank the authors ([pytorch-ddpm](https://github.com/w86763777/pytorch-ddpm), [Fast-SNN](https://github.com/yangfan-hu/Fast-SNN), [spikingjelly](https://github.com/fangwei123456/spikingjelly)) for their open-sourced codes.

For any help or issues of this project, please contact jcao248@connect.hkust-gz.edu.cn.

SDM

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README