{"id":30228930,"url":"https://github.com/compvis/tread","last_synced_at":"2025-08-14T20:07:05.649Z","repository":{"id":308545187,"uuid":"913984404","full_name":"CompVis/tread","owner":"CompVis","description":null,"archived":false,"fork":false,"pushed_at":"2025-08-06T13:06:16.000Z","size":6386,"stargazers_count":58,"open_issues_count":3,"forks_count":4,"subscribers_count":7,"default_branch":"master","last_synced_at":"2025-08-06T15:10:18.420Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/CompVis.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,"zenodo":null}},"created_at":"2025-01-08T18:15:46.000Z","updated_at":"2025-08-06T13:06:20.000Z","dependencies_parsed_at":"2025-08-06T15:20:50.079Z","dependency_job_id":null,"html_url":"https://github.com/CompVis/tread","commit_stats":null,"previous_names":["compvis/tread"],"tags_count":null,"template":false,"template_full_name":null,"purl":"pkg:github/CompVis/tread","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/CompVis%2Ftread","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/CompVis%2Ftread/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/CompVis%2Ftread/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/CompVis%2Ftread/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/CompVis","download_url":"https://codeload.github.com/CompVis/tread/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/CompVis%2Ftread/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":270475448,"owners_count":24590079,"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","status":"online","status_checked_at":"2025-08-14T02:00:10.309Z","response_time":75,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"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":"2025-08-14T20:07:01.603Z","updated_at":"2025-08-14T20:07:05.641Z","avatar_url":"https://github.com/CompVis.png","language":"Python","readme":"\u003ch2 align=\"center\"\u003e👟TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training\u003c/h2\u003e\n\u003cdiv align=\"center\"\u003e \n  \u003ca href=\"https://x.com/felix_m_krause\" target=\"_blank\"\u003eFelix Krause\u003c/a\u003e · \n  \u003ca href=\"\" target=\"_blank\"\u003eTimy Phan\u003c/a\u003e · \n  \u003ca href=\"\" target=\"_blank\"\u003eMing Gui\u003c/a\u003e · \n  \u003ca href=\"https://stefan-baumann.eu/\" target=\"_blank\"\u003eStefan Baumann\u003c/a\u003e · \n  \u003ca href=\"https://taohu.me\" target=\"_blank\"\u003eVincent Tao Hu\u003c/a\u003e · \n  \u003ca href=\"https://ommer-lab.com/people/ommer/\" target=\"_blank\"\u003eBjörn Ommer\u003c/a\u003e\n\u003c/div\u003e\n\u003cp align=\"center\"\u003e \n  \u003cb\u003eCompVis Group @ LMU Munich\u003c/b\u003e \u003cbr/\u003e\n\u003c/p\u003e\n\n[![Paper](https://img.shields.io/badge/arXiv-PDF-b31b1b)](https://arxiv.org/pdf/2501.04765)\n[![Project Page](https://img.shields.io/badge/Project-Page-blue)](https://compvis.github.io/tread/)\n\nThis repository contains the official implementation of the paper \"TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training\".\n\nWe propose TREAD, a new method to increase the efficiency of diffusion training by improving upon iteration speed and performance at the same time. For this, we use uni-directional token transportation to modulate the information flow in the network.\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"./docs/static/images/teaser.png\" alt=\"teaser\" style=\"width:50%;\"\u003e\n\u003c/div\u003e\n\n## 🚀 Training\n\nIn order to train a diffusion model, we offer a minimalistic training script in `train.py`. In its simplest form it can be started using:\n\n```python\naccelerate launch train.py model=tread\n```\n\nor\n\n```python\naccelerate launch train.py model=dit\n```\n\nwith `configs/config.yaml` having all the relevant information and settings for the actual training run. Please adjust this as needed before training.\n`Note:` We expect precomputed latents in this version.\nUnder `model` one can decide between `dit` and `tread` which are the preconfigured versions here with the former being the standard dit and the latter being supported by TREAD. How these changes are implemented can be seen in `dit.py` and `routing_module.py`.\n\nIn our paper, we show that TREAD can also work on other architectures. In practice, one needs to be more careful with the routing process in order to adhere to the characteristics of the specific architecture as some have a spatial bias (RWKV, Mamba, etc.). For simplicity, we only provide code for the Transformer architecture as it is the most widely used while being robust and easy to work with.\n\n## 🖼️ Sampling\n\nFor most experiments we use the [EDM](https://github.com/NVlabs/edm) training and sampling to stay consistent with prior art, and the FID calculation is done via the [ADM](https://github.com/openai/guided-diffusion) evaluation suite. We provide a `fid.py` to evaluate our models during training using the same reference batches as ADM.\n\n## 💥 Guiding TREAD\n\nTREAD works great during _training_! How about _inference_? \\\nIt turns out TREAD can be applied during guided inference as well to gain additional performance and reduce FLOPS at the same time! \\\nInstead of dropping the class label (CFG), we can guide with a selection rate delta. Since TREAD's selection rate (0.5) generalizes to other rates, this can be tuned in inference-time only.\n\nWe demonstrate this in `rf.py` which contains minimal flow matching code for training and sampling:\n\n`sample`: normal sampling\\\n`sample_tread`: TREAD sampling 🔥\n\n## 🎓 Citation\n\nIf you use this codebase or otherwise found our work valuable, please cite our paper:\n\n```bibtex\n@article{krause2025tread,\n  title={TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training},\n  author={Krause, Felix and Phan, Timy and Gui, Ming and Baumann, Stefan Andreas and Hu, Vincent Tao and Ommer, Bj{\\\"o}rn},\n  journal={arXiv preprint arXiv:2501.04765},\n  year={2025}\n}\n```\n\n## Acknowledgements\n\nThanks to the open source codebases such as [DiT](https://github.com/facebookresearch/DiT), [MaskDiT](https://github.com/Anima-Lab/MaskDiT), [ADM](https://github.com/openai/guided-diffusion), and [EDM](https://github.com/NVlabs/edm). Our codebase is built on them.\n","funding_links":[],"categories":[],"sub_categories":[],"project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcompvis%2Ftread","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fcompvis%2Ftread","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcompvis%2Ftread/lists"}