{"id":15652993,"url":"https://github.com/ssnl/poisson_quasimetric_embedding","last_synced_at":"2025-10-12T22:28:21.910Z","repository":{"id":41513702,"uuid":"462014852","full_name":"ssnl/poisson_quasimetric_embedding","owner":"ssnl","description":"Open source code for paper \"On the Learning and Learnability of Quasimetrics\".","archived":false,"fork":false,"pushed_at":"2022-11-28T22:14:55.000Z","size":137,"stargazers_count":32,"open_issues_count":0,"forks_count":1,"subscribers_count":2,"default_branch":"master","last_synced_at":"2025-04-30T20:31:47.241Z","etag":null,"topics":["deep-learning","machine-learning","pytorch"],"latest_commit_sha":null,"homepage":"https://www.tongzhouwang.info/quasimetric/","language":"C++","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"bsd-3-clause","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/ssnl.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2022-02-21T20:07:03.000Z","updated_at":"2024-11-22T06:53:09.000Z","dependencies_parsed_at":"2023-01-21T20:19:42.150Z","dependency_job_id":null,"html_url":"https://github.com/ssnl/poisson_quasimetric_embedding","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/ssnl/poisson_quasimetric_embedding","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ssnl%2Fpoisson_quasimetric_embedding","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ssnl%2Fpoisson_quasimetric_embedding/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ssnl%2Fpoisson_quasimetric_embedding/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ssnl%2Fpoisson_quasimetric_embedding/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/ssnl","download_url":"https://codeload.github.com/ssnl/poisson_quasimetric_embedding/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ssnl%2Fpoisson_quasimetric_embedding/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":279013266,"owners_count":26085250,"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-10-12T02:00:06.719Z","response_time":53,"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":["deep-learning","machine-learning","pytorch"],"created_at":"2024-10-03T12:44:26.307Z","updated_at":"2025-10-12T22:28:21.861Z","avatar_url":"https://github.com/ssnl.png","language":"C++","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Poisson Quasimetric Embedding (PQE)\n\n**[Tongzhou Wang](https://ssnl.github.io), [Phillip Isola](http://web.mit.edu/phillipi/)**\n\n---\n:fire::fire: **Checkout [`torchqmet`](https://github.com/quasimetric-learning/torch-quasimetric): our new PyTorch package of many SOTA quasimetric learning methods, including [Interval Quasimetric Embeddings (IQEs)](https://tongzhouwang.info/interval_quasimetric_embedding), a new quasimetric learning method that improves performances with a much simpler form**!\n\n---\nThis repository provides a PyTorch implementation of the Poisson Quasimetric Embedding (PQE) module for learning quasimetrics (i.e., asymmetrical distances).\n\nPQEs are proposed in **On the Learning and Learnability of Quasimetrics (ICLR 2022)**. It is the first method that has known guarantee to well approximate any quasimetric, and can be trained in gradient-based optimization (e.g., deep learning). Other common choices (e.g., unconstrained deep networks) provably fail. Empirically, PQEs show good results on learning quasimetrics of various sizes and structures, including an large-scale social graph and an offline Q-learning task.\n\n+ [arXiv](https://arxiv.org/abs/2206.15478)\n+ [ICLR 2022 (OpenReview)](https://openreview.net/forum?id=y0VvIg25yk)\n+ [Project Page](https://ssnl.github.io/quasimetric)\n\nRequirements:\n+ Python \u003e= 3.7\n+ PyTorch \u003e= 1.9.0\n+ NumPy \u003e= 1.19.2\n\n## Documentation\n\nA PQE represents a quasimetric (i.e., distances that can be asymmetrical) over an embedding space.\n\nThe main entry point is the ``pqe.PQE`` class that initializes a latent PQE quasimetric.  This module takes in two generic latent vectors ``u`` and ``v`` (both of a specific size) and output the estimated quasimetric distance from ``u`` to ``v``.\n\nCrucially, PQE can be trained to approximate any quasimetric distance. It does so by:\n- Parametrizes several distributions of quasipartitions (a specific kinds of quasimetric)\n    using Poisson Processes, where the latent vectors parametrize shapes (sets);\n- For each quasipartition distribution, compute the expected value;\n- Combine the expectations together together with non-negative weights (i.e., as a mixture\n    of them);\n\n```py\nclass PQE(nn.Module):\n    def __init__(\n        self,\n        num_quasipartition_mixtures: int,                   # Number of mixtures of quasipartition distributions\n                                                            #     to form the quasimetric estimate\n        num_poisson_processes_per_quasipartition: int = 4,  # Number of Poisson processes used for each\n                                                            #     quasipartition distributions\n        measure: str = 'lebesgue',                          # Mean measure of the Poisson processes;\n                                                            #     choices: 'lebesgue', 'gaussian'\n        shape: str = 'halfline',                            # Parametrization of shapes (sets) in Poisson\n                                                            #     process spaces by the input latents;\n                                                            #     choices: 'halfline', 'gaussian'\n        discounted: bool = False,                           # Whether this outputs distance or discounted\n                                                            #     distance (no need to specify the base\n                                                            #     because changing base is equivalent with\n                                                            #     applying a scale, which is already learned)\n    ):\n        ...\n```\n\n### **Measures and shape parameterizations**\n\nThe default values `measure='lebesgue', shape='halfline'` means that it defaults to the PQE-LH variant (Eqn. (10)),\nwhich is both straightforward to implement and works generally well in our experiments.\n\nThe choice `measure='lebesgue', shape='gaussian'` is prohibited as it leads to a very restricted symmetrical estimator.\n\nOther choices will internally use a utility extension at [`pqe.cdf_ops`](./pqe/cdf_ops) (included in this repo, see [here](./pqe/cdf_ops) for details and documentations). Notably, `measure='gaussian', shape='gaussian'`  gives the PQE-GG variant used in paper.\n\n## Getting Started and Examples\n\nTo start, we recommend adjusting `num_quasipartition_mixtures` according to the input latent\nspace dimensionality and keep the other arguments at default. The resulting PQE expects latent\nvectors of size ``4 * num_quasipartition_mixtures``, and is a PQE-LH variant.\n\n1. Quasimetric distances between two 64-dimensional latent vectors (with batch size 5)\n    ```py\n    \u003e\u003e\u003e pqe = PQE(num_quasipartition_mixtures=16,\n    ...           num_poisson_processes_per_quasipartition=4)\n    \u003e\u003e\u003e print(pqe)\n    \u003e\u003e\u003e u = torch.randn(5, 16 * 4)\n    \u003e\u003e\u003e v = torch.randn(5, 16 * 4)\n    \u003e\u003e\u003e print(pqe(u, v))  # non-negativity\n    tensor([0.2902, 0.3297, 0.2735, 0.3561, 0.2403], grad_fn=\u003cSqueezeBackward1\u003e)\n    \u003e\u003e\u003e print(pqe(v, u))  # asymmetrical\n    tensor([0.2554, 0.2326, 0.3263, 0.2562, 0.2975], grad_fn=\u003cSqueezeBackward1\u003e)\n    \u003e\u003e\u003e print(pqe(u, u))  # identity\n    tensor([0., 0., 0., 0., 0.], grad_fn=\u003cSqueezeBackward1\u003e)\n    \u003e\u003e\u003e t = torch.randn(5, 16 * 4)\n    \u003e\u003e\u003e print(pqe(u, v) + pqe(v, t) \u003e= pqe(u, t))  # triangle inequality\n    tensor([True, True, True, True, True])\n    ```\n\n2. **Discounted** quasimetric distances between two 64-dimensional latent vectors (with batch size 5)\n    ```py\n    \u003e\u003e\u003e discounted_pqe = PQE(num_quasipartition_mixtures=16,\n    ...                      num_poisson_processes_per_quasipartition=4,\n    ...                      discounted=True)\n    \u003e\u003e\u003e u = torch.randn(5, 16 * 4)\n    \u003e\u003e\u003e v = torch.randn(5, 16 * 4)\n    \u003e\u003e\u003e print(discounted_pqe(u, v))  # non-negativity (i.e., discounted distance \u003c= 1)\n    tensor([0.6986, 0.6614, 0.7698, 0.6864, 0.6138], grad_fn=\u003cProdBackward1\u003e)\n    \u003e\u003e\u003e print(discounted_pqe(v, u))  # asymmetrical\n    tensor([0.7258, 0.7232, 0.7233, 0.7440, 0.7511], grad_fn=\u003cProdBackward1\u003e)\n    \u003e\u003e\u003e print(discounted_pqe(u, u))  # identity (i.e., discounted distance = 1)\n    tensor([1., 1., 1., 1., 1.], grad_fn=\u003cProdBackward1\u003e)\n    \u003e\u003e\u003e t = torch.randn(5, 16 * 4)\n    \u003e\u003e\u003e print(discounted_pqe(u, v) * discounted_pqe(v, t) \u003c= discounted_pqe(u, t))  # triangle inequality\n    tensor([True, True, True, True, True])\n    ```\n\n3. Components of a PQE\n    ```py\n    \u003e\u003e\u003e pqe = PQE(num_quasipartition_mixtures=16,\n    ...           num_poisson_processes_per_quasipartition=4)\n    \u003e\u003e\u003e print(pqe)\n    PQE(\n        num_quasipartition_mixtures=16\n        num_poisson_processes_per_quasipartition=4\n        num_effective_parameters=16\n        discounted=False\n        (measure): LebesgueMeasure()\n        (shape): HalfLineShape()\n        (quasipartition_aggregator): DistanceAggregator(\n            (alpha_net): DeepLinearNet(\n                bias=True, non_negative=True\n                (mats): ParameterList(\n                    (0): Parameter containing: [torch.FloatTensor of size 1x64]\n                    (1): Parameter containing: [torch.FloatTensor of size 64x64]\n                    (2): Parameter containing: [torch.FloatTensor of size 64x64]\n                    (3): Parameter containing: [torch.FloatTensor of size 64x16]\n                )\n            )\n        )\n    )\n    ```\n    This PQE expects inputs from a (16*4)-dimensional latent space, where\n     - each 4 consecutive dimensions are used for a quasipartition distribution (`num_poisson_process_per_quasipartition=4`)\n     - a total of 16 such distributions mixed together (`num_quasipartition_mixtures=16`) to form the quasimetric\n     - the mixture weights are given by `pqe.quasipartition_aggregator.alpha_net` which parametrizes the 16 non-negative\n       numbers with deep linear networks (whose effective parameters are few).\n\n4. Training a PQE\n    ```py\n    # Training with MSE on discounted distances\n\n    # 128-dimensional paired data\n    x = ...      # shape [N, 128]\n    y = ...      # shape [N, 128]\n    d_x2y = ...  # shape [N], \u003e=0, could be +\\infty\n\n    # Simple 2 layer encoder from data space to latent space\n    encoder = nn.Sequential(\n        nn.Linear(128, 512),\n        nn.ReLU()\n        nn.Linear(512, 512),\n        nn.ReLU()\n        nn.Linear(512, 16 * 4),\n    )\n\n    # Define the latent PQE, use discounted version\n    discounted_pqe = PQE(num_quasipartition_mixtures=16,\n                         num_poisson_processes_per_quasipartition=4,\n                         discounted=True)\n    optim = torch.optim.Adam(\n        itertools.chain(encoder.parameters(), discounted_pqe.parameters()),  # both encoder and pqe have parameters\n        lr=1e-3,\n    )\n    while keep_training():\n        optim.zero_grad()\n        zx = encoder(x)\n        zy = encoder(y)\n        pred = discounted_pqe(zx, zy)\n        loss = F.mse_loss(0.9 ** d_x2y, pred)\n        loss.backward()\n        optim.step()\n    ```\n\n## Citation\n\nTongzhou Wang, Phillip Isola. \"On the Learning and Learnability of Quasimetrics\" International Conference on Learning Representations. 2022.\n\n```\n@inproceedings{wang2022learning,\n  title={On the Learning and Learnability of Quasimetrics},\n  author={Wang, Tongzhou and Isola, Phillip},\n  booktitle={International Conference on Learning Representations},\n  year={2022}\n}\n```\n\n## Questions\n\nFor questions about the code provided in this repository, please open an GitHub issue.\n\nFor questions about the paper, please contact Tongzhou Wang (`tongzhou _AT_ mit _DOT_ edu`).\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fssnl%2Fpoisson_quasimetric_embedding","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fssnl%2Fpoisson_quasimetric_embedding","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fssnl%2Fpoisson_quasimetric_embedding/lists"}