{"id":17714323,"url":"https://github.com/nus-hpc-ai-lab/speed","last_synced_at":"2025-09-05T21:41:20.106Z","repository":{"id":229495702,"uuid":"767335399","full_name":"NUS-HPC-AI-Lab/SpeeD","owner":"NUS-HPC-AI-Lab","description":"SpeeD: A Closer Look at Time Steps is Worthy of Triple  Speed-Up for Diffusion Model Training","archived":false,"fork":false,"pushed_at":"2025-01-27T00:48:14.000Z","size":69099,"stargazers_count":178,"open_issues_count":5,"forks_count":7,"subscribers_count":8,"default_branch":"master","last_synced_at":"2025-05-07T20:17:55.317Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/NUS-HPC-AI-Lab.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE.txt","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-03-05T05:25:54.000Z","updated_at":"2025-04-08T13:55:27.000Z","dependencies_parsed_at":"2025-01-25T17:40:45.620Z","dependency_job_id":null,"html_url":"https://github.com/NUS-HPC-AI-Lab/SpeeD","commit_stats":null,"previous_names":["1zeryu/speedit","1zeryu/speed","kaiwang960112/speed","nus-hpc-ai-lab/speed"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/NUS-HPC-AI-Lab%2FSpeeD","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/NUS-HPC-AI-Lab%2FSpeeD/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/NUS-HPC-AI-Lab%2FSpeeD/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/NUS-HPC-AI-Lab%2FSpeeD/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/NUS-HPC-AI-Lab","download_url":"https://codeload.github.com/NUS-HPC-AI-Lab/SpeeD/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":252949239,"owners_count":21830154,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-10-25T11:02:21.497Z","updated_at":"2025-05-07T20:18:03.715Z","avatar_url":"https://github.com/NUS-HPC-AI-Lab.png","language":"Python","funding_links":[],"categories":["Accelerate"],"sub_categories":[],"readme":"\u003cp align=\"center\"\u003e\n    \u003cimg src=\"visuals/logo.png\" width=\"150\" style=\"margin-bottom: 0.2;\"/\u003e\n\n\n\u003cp\u003e\n\u003ch2 align=\"center\"\u003eA Closer Look at Time Steps is Worthy of Triple Speed-Up for Diffusion Model Training\n\u003ch5 align=\"center\"\u003e If you like SpeeD, please give us a star ⭐ on GitHub for the latest update.\n\u003c/h2\u003e\n\n### [Paper](https://arxiv.org/pdf/2405.17403) | [Project Page](https://bdemo.github.io/SpeeD/) | [Hugging Face]()\n\nThis repository contains the code and implementation details for the research paper titled \"A Closer Look at Time Steps is Worthy of Triple Speed-Up for Diffusion Model Training.\" In this paper, SpeeD, a novel speed-up method for diffusion model training, is introduced.\n\n## Authors\n\n- [Kai Wang](https://kaiwang960112.github.io/)\u003csup\u003e2\u003c/sup\u003e, Yukun Zhou\u003csup\u003e1,2\u003c/sup\u003e, [Mingjia Shi](https://www.samjs.online/)\u003csup\u003e2\u003c/sup\u003e, [Zekai Li](https://lizekai-richard.github.io/)\u003csup\u003e2\u003c/sup\u003e, [Zhihang Yuan](https://zhihang.cc/)\u003csup\u003e3\u003c/sup\u003e, [Yuzhang Shang](https://42shawn.github.io/)\u003csup\u003e4\u003c/sup\u003e, [Xiaojiang Peng*](https://pengxj.github.io/)\u003csup\u003e1\u003c/sup\u003e, [Hanwang Zhang](https://personal.ntu.edu.sg/hanwangzhang/)\u003csup\u003e5\u003c/sup\u003e, [Yang You](https://www.comp.nus.edu.sg/~youy/)\u003csup\u003e2\u003c/sup\u003e\n- \u003csup\u003e1\u003c/sup\u003e[Shenzhen Technology University](https://english.sztu.edu.cn/), \u003csup\u003e2\u003c/sup\u003e[National University of Singapore](https://nus.edu.sg/), \u003csup\u003e3\u003c/sup\u003e[Infinigence-AI](https://cloud.infini-ai.com/), \u003csup\u003e4\u003c/sup\u003e[Illinois Institute of Technology](https://www.iit.edu/), and \u003csup\u003e5\u003c/sup\u003e[Nanyang Technological University](https://www.ntu.edu.sg/) [Kai, Yukun, and Mingjia contribute equally to this work. We will update this repo asap.]\n\n## 😮 Highlights\n\nOur method, which is easily compatible, can accelerate the training of diffusion model.\n\n![comparision](visuals/accel.png)\n\n\n\n##  ✒️ Motivation\n\u003c!-- \nInspired by the uphill and downhill diffusion processes in physics. The following GIF illustrates the  commonalities  between image diffusion and electron diffusion. The left  figure of electric diffusion is simulated from  [PhET/diffusion](https://phet.colorado.edu/zh_CN/simulations/diffusion). The right figure is downloaded from [OpenAI website](https://images.openai.com/blob/b196df3a-6fea-4d86-87b2-f9bb50be64c7/leaf.png?trim=0,0,0,0\u0026width=2600).![comparision](visuals/consistency.gif)\n\nVisualization of different phases of reverse process and [uphill diffusion](https://en.wikipedia.org/wiki/Diffusion). For easy understanding, we assume that the direction of electronic velocity only has two cases: :arrow_left: and :arrow_right:.\n\n![motivation](visuals/motivation.png) --\u003e\n\nInspired by the following observation on time steps, we propose the re-sampling + re-weighting strategy as shown below.\n\nTo take a closer look at time steps, we find that the time steps could be divided into three areas: acceleration, decceleration and convergence areas. Samples of the corresponding time step in the convergence area are of limited benefit to training, while these time steps take up the most. Empirically, the training losses of these samples are quite low compare to the ones of the other two areas.\n![motivation](visuals/Findings.png)\n\nAsymmetric Sampling: Suppress the attendance of the time step in convergence areas.\n\nChange-Aware Weighting: The faster changing time steps in the diffusion process are given more weight.\n![method](visuals/Method.png)\n\n\u003c!-- ##  🔆 Method\n\nWe use the sampling and weighting strategy which are simple and easily compatible to achieve the acceleration. The following is the core code  [SpeeD/speed/diffusion/iddpm/speed.py](https://github.com/kaiwang960112/SpeeD/blob/master/speed/diffusion/iddpm/speed.py) ,\n\n```python\nclass SpeeDiffusion(SpacedDiffusion):\n    def __init__(self, faster, **kwargs):\n        super().__init__(**kwargs)\n        self.faster = faster\n        if faster:\n            grad = np.gradient(self.sqrt_one_minus_alphas_cumprod)\n\n            # set the meaningful steps in diffusion, which is more important in inference\n            self.meaningful_steps = np.argmax(grad \u003c 1e-4) + 1\n\n            # p2 weighting from: Perception Prioritized Training of Diffusion Models\n            self.p2_gamma = 1\n            self.p2_k = 1\n            self.snr = 1.0 / (1 - self.alphas_cumprod) - 1\n            sqrt_one_minus_alphas_bar = torch.from_numpy(self.sqrt_one_minus_alphas_cumprod)\n            # sample more meaningful step\n            p = torch.tanh(1e6 * (torch.gradient(sqrt_one_minus_alphas_bar)[0] - 1e-4)) + 1.5\n            self.p = F.normalize(p, p=1, dim=0)\n            self.weights = self._weights()\n        else:\n            self.meaningful_steps = self.num_timesteps\n\n    def _weights(self):\n        # process where all noise to noisy image with content has more weighting in training\n        # the weights act on the mse loss\n        weights =  1 / (self.p2_k + self.snr) ** self.p2_gamma\n        weights = weights\n        return weights\n\n    # get the weights and sampling t in training diffusion\n    def t_sample(self, n, device):\n        if self.faster:\n            t = torch.multinomial(self.p, n // 2 + 1, replacement=True).to(device)\n            # dual sampling, which can balance the step multiple task training\n            dual_t = torch.where(t \u003c self.meaningful_steps, self.meaningful_steps - t, t - self.meaningful_steps)\n            t = torch.cat([t, dual_t], dim=0)[:n]\n            weights = self.weights\n        else:\n            # if\n            t = torch.randint(0, self.num_timesteps, (n,), device=device)\n            weights = None\n\n        return t, weights\n```\n\n You can enable our acceleration module with **diffusion.faster=True**.\n\n```\n# config file\ndiffusion:\n    timestep_respacing: '250'\n    faster: true  #enabl module for training acceleration\n``` --\u003e\n\n\n\n## 🛠️ Requirements and Installation\n\nThis code base does not use hardware acceleration technology, experimental environment is not complicated.\n\nYou can create a new conda environment:\n\n```\nconda env create -f environment.yml\nconda activate speed\n```\n\nor install the necessary package by:\n\n```\npip install -r requirements.txt\n```\n\nIf necessary, we will provide more methods (e.g., docker) to facilitate the configuration of the experimental environment.\n\n## 🗝️ Tutorial\n\nWe provide a complete process for generating tasks including **training**, **inference** and **test**. The current code is only compatible with class-conditional image generation tasks. We will be compatible with more generation tasks about diffusion in the future.\n\nWe refactor the [facebookresearch/DiT](https://github.com/facebookresearch/DiT) code and loaded the configs using  [OmegaConf ](https://omegaconf.readthedocs.io/en/2.3_branch/). The configuration file loading rule is  recursive for easier argument modification. Simply put, the file in the latter path will override the previous setting of **base.yaml**.\n\nYou can modify the experiment setting by modifying the config file and the command line. More details about the reading of config are written in  [configs/README.md](https://github.com/kaiwang960112/SpeeD/blob/master/configs/README.md).\n\nFor each experiment, you must provide two arguments by command,\n\n```\n-c: config path;\n-p: phase including ['train', 'inference', 'sample'].\n```\n\n### Train \u0026 inference\n\n**Baseline**\n\nClass-conditional image generation task with 256x256 ImageNet dataset and DiT-XL/2 models.\n\n```bash\n# Training: training diffusion and saving checkpoints\ntorchrun --nproc_per_node=8 main.py -c configs/image/imagenet_256/base.yaml -p train\n# inference: generating samples for testing\ntorchrun --nproc_per_node=8 main.py -c configs/image/imagenet_256/base.yaml -p inference\n# sample: sample some images for visualization\npython main.py -c configs/image/imagenet_256/base.yaml -p sample\n```\n\n**Ablation**\n\nYou can modify the experiment setting by modifying the config file and the command line. More details about the configs are in [configs/README.md](https://github.com/kaiwang960112/SpeeD/blob/master/configs/README.md).\n\nFor example,  change the classifier-free guidance scale in sampling by command line:\n\n```\npython main.py -c configs/image/imagenet_256/base.yaml -p sample guidance_scale=1.5\n```\n\n### Test\n\nTest the generation tasks require the results of inference. The more details about testing in  [evaluations](https://github.com/kaiwang960112/SpeeD/tree/master/evaluations).\n\n## 🔒 License\n\nThe majority of this project is released under the Apache 2.0 license as found in the [LICENSE](https://github.com/PKU-YuanGroup/Video-LLaVA/blob/main/LICENSE) file.\n\n## ✏️Citation\n\n If you find our code useful in your research, please consider giving a star ⭐ and citation 📝.\n\n```\n@article{wang2024closer,\n      title={A Closer Look at Time Steps is Worthy of Triple Speed-Up for Diffusion Model Training}, \n      author={Kai Wang, Mingjia Shi, Yukun Zhou, Zekai Li, Zhihang Yuan, Yuzhang Shang, Xiaojiang Peng, Hanwang Zhang and Yang You},\n      year={2024},\n      journal={arXiv preprint arXiv:2405.17403},\n}\n```\n\n## 👍 Acknowledgement\n\nWe thank Tianyi Li, Yuchen Zhang, Yuxin Li, Zhaoyang Zeng, and Yanqing Liu for the comments on this work. Kai Wang (idea, writing, story, presentation), Yukun Zhou (implementation), and Mingjia Shi (theory, writing, presentation) contribute equally to this work. Xiaojiang Peng, Hanwang Zhang, and Yang You are equal advising. Xiaojiang Peng is the corresponding author.\n\nWe are grateful for the following exceptional work and generous contribution to open source.\n\n* [DiT](https://github.com/facebookresearch/DiT): Scalable Diffusion Models with Transformers.\n* [Open-Sora](https://github.com/hpcaitech/Open-Sora/tree/main) : Open-Sora: Democratizing Efficient Video Production for All\n* [OpenDiT](https://github.com/NUS-HPC-AI-Lab/OpenDiT): An acceleration for DiT training. We adopt valuable acceleration strategies for training progress from OpenDiT.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fnus-hpc-ai-lab%2Fspeed","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fnus-hpc-ai-lab%2Fspeed","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fnus-hpc-ai-lab%2Fspeed/lists"}