{"id":23534980,"url":"https://github.com/soran-ghaderi/torchebm","last_synced_at":"2026-06-01T00:01:02.903Z","repository":{"id":269636703,"uuid":"868083982","full_name":"soran-ghaderi/torchebm","owner":"soran-ghaderi","description":"🍓 Build and train energy-based and diffusion models in PyTorch ⚡.","archived":false,"fork":false,"pushed_at":"2026-05-31T11:56:18.000Z","size":89005,"stargazers_count":88,"open_issues_count":16,"forks_count":10,"subscribers_count":5,"default_branch":"master","last_synced_at":"2026-05-31T13:19:43.641Z","etag":null,"topics":["contrastive-divergence","cuda","diffusion-models","energy-based-model","equilibrium","equilibrium-matching","equilibrium-modeling","generative-ai","hamiltonian","hamiltonian-monte-carlo","langevin-dynamics","noise-contrastive-estimation","probabilistic-machine-learning","reasoning","sampling-methods","score-matching","variational-inference"],"latest_commit_sha":null,"homepage":"https://soran-ghaderi.github.io/torchebm/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/soran-ghaderi.png","metadata":{"files":{"readme":"README.md","changelog":"CHANGELOG.md","contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","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,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2024-10-05T12:41:38.000Z","updated_at":"2026-05-31T11:53:37.000Z","dependencies_parsed_at":null,"dependency_job_id":"1da5b616-a525-428c-b019-1e7c3b610937","html_url":"https://github.com/soran-ghaderi/torchebm","commit_stats":null,"previous_names":["soran-ghaderi/torchebm"],"tags_count":103,"template":false,"template_full_name":null,"purl":"pkg:github/soran-ghaderi/torchebm","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/soran-ghaderi%2Ftorchebm","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/soran-ghaderi%2Ftorchebm/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/soran-ghaderi%2Ftorchebm/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/soran-ghaderi%2Ftorchebm/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/soran-ghaderi","download_url":"https://codeload.github.com/soran-ghaderi/torchebm/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/soran-ghaderi%2Ftorchebm/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":33753925,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-05-31T02:00:06.040Z","response_time":95,"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":["contrastive-divergence","cuda","diffusion-models","energy-based-model","equilibrium","equilibrium-matching","equilibrium-modeling","generative-ai","hamiltonian","hamiltonian-monte-carlo","langevin-dynamics","noise-contrastive-estimation","probabilistic-machine-learning","reasoning","sampling-methods","score-matching","variational-inference"],"created_at":"2024-12-26T01:14:07.336Z","updated_at":"2026-06-01T00:01:02.887Z","avatar_url":"https://github.com/soran-ghaderi.png","language":"Python","funding_links":[],"categories":["Python","Open Source Libraries"],"sub_categories":[],"readme":"\u003cp align=\"center\"\u003e\n  \u003cimg src=\"docs/assets/images/logo_with_text.svg\" alt=\"TorchEBM Logo\" width=\"350\"\u003e\n\u003c/p\u003e\n\n\u003cp align=\"center\" style=\"margin-bottom: 20px;\"\u003e\n    \u003ca href=\"https://pypi.org/project/torchebm/\" target=\"_blank\" title=\"PyPI version\"\u003e\n        \u003cimg alt=\"PyPI\" src=\"https://img.shields.io/pypi/v/torchebm?style=flat-square\u0026color=blue\"\u003e\n    \u003c/a\u003e\n    \u003ca href=\"https://github.com/soran-ghaderi/torchebm/blob/master/LICENSE\" target=\"_blank\" title=\"License\"\u003e\n        \u003cimg alt=\"License\" src=\"https://img.shields.io/github/license/soran-ghaderi/torchebm?style=flat-square\u0026color=brightgreen\"\u003e\n    \u003c/a\u003e\n    \u003ca href=\"https://github.com/soran-ghaderi/torchebm\" target=\"_blank\" title=\"GitHub Repo Stars\"\u003e\n        \u003cimg alt=\"GitHub Stars\" src=\"https://img.shields.io/github/stars/soran-ghaderi/torchebm?style=social\"\u003e\n    \u003c/a\u003e\n    \u003ca href=\"https://deepwiki.com/soran-ghaderi/torchebm\"\u003e\u003cimg src=\"https://deepwiki.com/badge.svg\" alt=\"Ask DeepWiki\"\u003e\u003c/a\u003e\n    \u003c!-- Consider adding: build status, documentation status, code coverage --\u003e\n    \u003ca href=\"https://github.com/soran-ghaderi/torchebm/actions\" target=\"_blank\" title=\"Build Status\"\u003e\n      \u003cimg alt=\"Build Status\" src=\"https://img.shields.io/github/actions/workflow/status/soran-ghaderi/torchebm/tag-release.yml?branch=master\u0026style=flat-square\u0026label=build\"\u003e\n    \u003c/a\u003e\n    \u003c!-- Docs badge --\u003e\n    \u003ca href=\"https://github.com/soran-ghaderi/torchebm/actions\" target=\"_blank\" title=\"Documentation\"\u003e\n      \u003cimg alt=\"Docs\" src=\"https://img.shields.io/github/actions/workflow/status/soran-ghaderi/torchebm/docs_ci.yml?branch=master\u0026style=flat-square\u0026label=docs\"\u003e\n    \u003c/a\u003e\n\u003c!--     \u003ca href=\"https://pepy.tech/project/torchebm\" target=\"_blank\" title=\"Downloads\"\u003e\n        \u003cimg alt=\"Downloads\" src=\"https://img.shields.io/pypi/dm/torchebm?style=flat-square\"\u003e\n    \u003c/a\u003e --\u003e\n    \u003ca href=\"https://pepy.tech/project/torchebm\" target=\"_blank\" title=\"Downloads\"\u003e\n        \u003cimg alt=\"Downloads\" src=\"https://static.pepy.tech/badge/torchebm?style=flat-square\"\u003e\n    \u003c/a\u003e\n    \u003ca href=\"https://pypi.org/project/torchebm/\" target=\"_blank\" title=\"Python Versions\"\u003e\n        \u003cimg alt=\"Python Versions\" src=\"https://img.shields.io/pypi/pyversions/torchebm?style=flat-square\"\u003e\n    \u003c/a\u003e\n\u003c/p\u003e\n\n\n\u003cp align=\"center\"\u003e⚡ A PyTorch library for energy-based modeling, with support for flow and diffusion methods.\u003c/p\u003e\n\n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"docs/assets/animations/ebm_training_animation.gif\" alt=\"EBM Training Animation\"/\u003e\n\u003c/p\u003e\n\n## What is ∇ TorchEBM 🍓? \n\nEnergy-based models define distributions through a scalar energy function, where lower energy means higher probability. This is a very general formulation and many generative approaches, from MCMC sampling to score matching to flow-based generation, can be understood through this lens.\n\n**TorchEBM** is a PyTorch library that gives you composable tools for this entire spectrum. You can define energy landscapes, train models with various learning objectives, and sample via MCMC, optimization, or learned continuous-time dynamics (ODEs/SDEs). The library handles classical EBM training (contrastive divergence, score matching) as well as modern interpolant-based and equilibrium-based generation methods.\n\n📚 For the full documentation, please visit the [official website of TorchEBM 🍓](https://soran-ghaderi.github.io/torchebm/).\n\n## Features\n\n- **Energy models** with built-in analytical potentials and support for custom neural network energy functions\n- **MCMC and optimization-based samplers** for drawing samples from energy landscapes\n- **Flow and diffusion samplers** that generate via ODE/SDE integration of learned velocity or score fields\n- **Training objectives** including contrastive divergence variants, score matching variants, and equilibrium matching\n- **Interpolation schemes** for specifying noise-to-data paths in flow and diffusion models\n- **Numerical integrators** for SDE, ODE, and Hamiltonian dynamics\n- **Neural network architectures** ready for conditional generation\n- **Synthetic datasets** for rapid prototyping and benchmarking\n- **Hyperparameter schedulers** for step sizes, noise scales, and other training parameters\n- **CUDA acceleration** and mixed precision support\n\n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"docs/assets/animations/8gaussians_flow.gif\" alt=\"8 Gaussians Flow\" width=\"700\"/\u003e\n\u003c/p\u003e\n\u003ctable align=\"center\"\u003e\n  \u003ctr\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/e_functions/gaussian.png\" alt=\"Gaussian\" width=\"200\"/\u003e\u003c/td\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/e_functions/double_well.png\" alt=\"Double Well\" width=\"200\"/\u003e\u003c/td\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/e_functions/rastrigin.png\" alt=\"Rastrigin\" width=\"200\"/\u003e\u003c/td\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/e_functions/rosenbrock.png\" alt=\"Rosenbrock\" width=\"200\"/\u003e\u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd align=\"center\"\u003eGaussian\u003c/td\u003e\n    \u003ctd align=\"center\"\u003eDouble Well\u003c/td\u003e\n    \u003ctd align=\"center\"\u003eRastrigin\u003c/td\u003e\n    \u003ctd align=\"center\"\u003eRosenbrock\u003c/td\u003e\n  \u003c/tr\u003e\n\u003c/table\u003e\n\n\u003ctable align=\"center\"\u003e\n  \u003ctr\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/datasets/gaussian_mixture.png\" alt=\"Gaussian Mixture\" width=\"200\"/\u003e\u003c/td\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/datasets/two_moons.png\" alt=\"Two Moons\" width=\"200\"/\u003e\u003c/td\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/datasets/swiss_roll.png\" alt=\"Swiss Roll\" width=\"200\"/\u003e\u003c/td\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/datasets/checkerboard.png\" alt=\"Checkerboard\" width=\"200\"/\u003e\u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd align=\"center\"\u003eGaussian Mixture\u003c/td\u003e\n    \u003ctd align=\"center\"\u003eTwo Moons\u003c/td\u003e\n    \u003ctd align=\"center\"\u003eSwiss Roll\u003c/td\u003e\n    \u003ctd align=\"center\"\u003eCheckerboard\u003c/td\u003e\n  \u003c/tr\u003e\n\u003c/table\u003e\n\n\n## Installation\n\n```bash\npip install torchebm\n```\n\n#### Dependencies\n- [PyTorch](https://pytorch.org/) (with CUDA support for optimal performance)\n- Other dependencies are listed in [requirements.txt](requirements.txt)\n\n\n## Usage Examples\n\n### MCMC Sampling\n\n```python\nimport torch\nfrom torchebm.core import GaussianModel\nfrom torchebm.samplers import LangevinDynamics\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = GaussianModel(mean=torch.zeros(2), cov=torch.eye(2), device=device)\n\nsampler = LangevinDynamics(model=model, step_size=0.01, device=device)\nsamples = sampler.sample(x=torch.randn(500, 2, device=device), n_steps=100)\nprint(samples.shape)  # torch.Size([500, 2])\n```\n\n### Training with Contrastive Divergence\n\n```python\nimport torch\nfrom torchebm.core import BaseModel\nfrom torchebm.samplers import LangevinDynamics\nfrom torchebm.losses import ContrastiveDivergence\nfrom torchebm.datasets import GaussianMixtureDataset\nfrom torch.utils.data import DataLoader\n\nclass MLPEnergy(BaseModel):\n    def __init__(self, dim):\n        super().__init__()\n        self.net = torch.nn.Sequential(\n            torch.nn.Linear(dim, 64), torch.nn.SiLU(),\n            torch.nn.Linear(64, 64), torch.nn.SiLU(),\n            torch.nn.Linear(64, 1),\n        )\n    def forward(self, x):\n        return self.net(x).squeeze(-1)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = MLPEnergy(dim=2).to(device)\nsampler = LangevinDynamics(model=model, step_size=0.01, device=device)\ncd_loss = ContrastiveDivergence(model=model, sampler=sampler, k_steps=10)\n\ndata = GaussianMixtureDataset(n_samples=1000, n_components=4).get_data()\nloader = DataLoader(data, batch_size=64, shuffle=True)\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\nfor epoch in range(10):\n    for batch in loader:\n        optimizer.zero_grad()\n        loss, _ = cd_loss(batch.to(device))\n        loss.backward()\n        optimizer.step()\n```\n\n### Hamiltonian Monte Carlo\n\n```python\nimport torch\nfrom torchebm.core import GaussianModel\nfrom torchebm.samplers import HamiltonianMonteCarlo\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = GaussianModel(mean=torch.zeros(10), cov=torch.eye(10), device=device)\n\nhmc = HamiltonianMonteCarlo(model=model, step_size=0.1, n_leapfrog_steps=10, device=device)\nsamples = hmc.sample(dim=10, n_steps=500, n_samples=1000)\nprint(samples.shape)  # torch.Size([1000, 10])\n```\n\n## Library Structure\n\n```\ntorchebm/\n├── core/           # Base classes, energy models, schedulers, device management\n├── samplers/       # MCMC, optimization, and flow/diffusion samplers\n├── losses/         # Training objectives (CD, score matching, equilibrium matching)\n├── interpolants/   # Noise-to-data interpolation schemes\n├── integrators/    # Numerical integrators for SDE/ODE/Hamiltonian dynamics\n├── models/         # Neural network architectures\n├── datasets/       # Synthetic data generators\n├── utils/          # Visualization and training utilities\n└── cuda/           # CUDA-accelerated implementations\n```\n\n## Visualization Examples\n\n\u003ctable align=\"center\"\u003e\n  \u003ctr\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/sampling.jpg\" alt=\"Langevin Dynamics Sampling\" width=\"250\"/\u003e\u003c/td\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/trajectory.jpg\" alt=\"Langevin Dynamics Trajectory\" width=\"250\"/\u003e\u003c/td\u003e\n    \u003ctd\u003e\u003cimg src=\"docs/assets/images/parallel.jpg\" alt=\"Parallel Sampling\" width=\"250\"/\u003e\u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd align=\"center\"\u003eLangevin Dynamics Sampling\u003c/td\u003e\n    \u003ctd align=\"center\"\u003eLangevin Dynamics Trajectory\u003c/td\u003e\n    \u003ctd align=\"center\"\u003eParallel Sampling\u003c/td\u003e\n  \u003c/tr\u003e\n\u003c/table\u003e\n\n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"docs/assets/animations/circles_flow.gif\" alt=\"Flow Comparison\" width=\"700\"/\u003e\n  \u003cbr\u003e\n  \u003cem\u003eEquilibrium Matching: Linear, VP, and Cosine interpolants transforming noise into data.\u003c/em\u003e\n\u003c/p\u003e\n\nCheck out the `examples/` directory for sample scripts.\n\n## Contributing\n\nContributions are welcome! Step-by-step instructions for contributing to the project can be found on the [contributing.md](docs/developer_guide/contributing.md) page on the website.\n\nPlease check the issues page for current tasks or create a new issue to discuss proposed changes.\n\n## Show your Support for ∇ TorchEBM 🍓\n\nPlease ⭐️ this repository if ∇ TorchEBM helped you and spread the word.\n\nThank you! 🚀\n\n## Citation\n\nIf TorchEBM is useful in your research, please cite it:\n\n```bibtex\n@misc{torchebm_library_2025,\n  author       = {Ghaderi, Soran and Contributors},\n  title        = {{TorchEBM}: A PyTorch Library for Training Energy-Based Models},\n  year         = {2025},\n  url          = {https://github.com/soran-ghaderi/torchebm},\n}\n```\n\n## Changelog\n\nSee [CHANGELOG](CHANGELOG.md) for version history.\n\n## License\n\nMIT License. See [LICENSE](LICENSE) for details.\n\n## Research Collaboration\n\nIf you are interested in collaborating on research around energy-based, flow-based, or diffusion models, feel free to reach out. Contributions to TorchEBM 🍓 and discussions that push the field forward are always welcome.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsoran-ghaderi%2Ftorchebm","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsoran-ghaderi%2Ftorchebm","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsoran-ghaderi%2Ftorchebm/lists"}