{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/large-scale-wasserstein-gradient-flows","title":"Large-Scale Wasserstein Gradient Flows","arxiv_id":"2106.00736","date":"2021-06-01","proceeding":"NeurIPS 2021 12","authors":["Petr Mokrov","Alexander Korotin","Lingxiao Li","Aude Genevay","Justin Solomon","Evgeny Burnaev"],"abstract":"Wasserstein gradient flows provide a powerful means of understanding and solving many diffusion equations. Specifically, Fokker-Planck equations, which model the diffusion of probability measures, can be understood as gradient descent over entropy functionals in Wasserstein space. This equivalence, introduced by Jordan, Kinderlehrer and Otto, inspired the so-called JKO scheme to approximate these diffusion processes via an implicit discretization of the gradient flow in Wasserstein space. Solving the optimization problem associated to each JKO step, however, presents serious computational challenges. We introduce a scalable method to approximate Wasserstein gradient flows, targeted to machine learning applications. Our approach relies on input-convex neural networks (ICNNs) to discretize the JKO steps, which can be optimized by stochastic gradient descent. Unlike previous work, our method does not require domain discretization or particle simulation. As a result, we can sample from the measure at each time step of the diffusion and compute its probability density. We demonstrate our algorithm's performance by computing diffusions following the Fokker-Planck equation and apply it to unnormalized density sampling as well as nonlinear filtering.","url_abs":"https://arxiv.org/abs/2106.00736v2","url_pdf":"https://arxiv.org/pdf/2106.00736v2.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"large-scale-wasserstein-gradient-flows","repo_url":"https://github.com/PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"pytorch","reach":null},{"paper_slug":"large-scale-wasserstein-gradient-flows","repo_url":"https://github.com/iamalexkorotin/Wasserstein2Barycenters","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok","spdx":"MIT"}},{"paper_slug":"large-scale-wasserstein-gradient-flows","repo_url":"https://github.com/iamalexkorotin/Wasserstein2Benchmark","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"pytorch","reach":null}],"tasks":[],"methods":[{"method_slug":"diffusion","method_name":"Diffusion"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2106.00736","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2106.00736"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows","reach":null},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/iamalexkorotin/Wasserstein2Barycenters","reach":{"status":"ok","spdx":"MIT"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/iamalexkorotin/Wasserstein2Benchmark","reach":null}],"summary":{"ran_draft_wrong":3},"by_repo_kind":{"official":{"samples":3,"ran":3,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"c506fa4e8b8fccec","entry":"extract_exp_ns","repo":"PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows","repo_kind":"official","path":"script.py","file_url":"https://github.com/PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows/blob/HEAD/script.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":"deterministic","behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"c506fa4e8b8fccec"}},{"code_sha256_prefix":"a11f707d494e7d87","entry":"extract_method","repo":"PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows","repo_kind":"official","path":"script.py","file_url":"https://github.com/PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows/blob/HEAD/script.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":"deterministic","behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"a11f707d494e7d87"}},{"code_sha256_prefix":"c528ab9d6897d3fb","entry":"extract_verbose","repo":"PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows","repo_kind":"official","path":"script.py","file_url":"https://github.com/PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows/blob/HEAD/script.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":"deterministic","behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"c528ab9d6897d3fb"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}