discrete-distribution-network

Discrete Distribution Network

Exploration into Discrete Distribution Network, by Lei Yang out of Beijing

Besides the split-and-prune, may also throw in an option for crossover (mixing of top 2 nodes to replace the pruned)

Install

$ pip install discrete-distribution-network

Usage

import torch
from discrete_distribution_network import DDN

ddn = DDN(
    dim = 32,
    image_size = 256
)

images = torch.randn(2, 3, 256, 256)

loss = ddn(images)
loss.backward()

# after much training

sampled = ddn.sample(batch_size = 1)

assert sampled.shape == (1, 3, 256, 256)

The proposed GuidedSampler in the paper

import torch
from discrete_distribution_network import GuidedSampler

sampler = GuidedSampler(
    dim = 16,              # feature dimension
    dim_query = 3,         # the query image dimension
    codebook_size = 10,    # the codebook size K in the paper, which is K separate projections of the features, which is then measured distance wise to the query image guide
)

features = torch.randn(20, 16, 32, 32)
query_image = torch.randn(20, 3, 32, 32)

out, codes, commit_loss = sampler(features, query_image)

# (20, 3, 32, 32), (20,), ()

assert torch.allclose(sampler.forward_for_codes(features, codes), out, atol = 1e-5)

# after optimizer step, this needs to be called
# there is also a helper function by the same name that can take in your module and will invoke all of the guided samplers

sampler.split_and_prune_()

Oxford flowers

Install uv, which will probably become the default in the near future

$ pip install uv

Then

$ uv run train_oxford_flowers.py

Citations

@misc{yang2025discretedistributionnetworks,
    title   = {Discrete Distribution Networks},
    author  = {Lei Yang},
    year    = {2025},
    eprint  = {2401.00036},
    archivePrefix = {arXiv},
    primaryClass = {cs.CV},
    url     = {https://arxiv.org/abs/2401.00036},
}

Metadata

Release files for discrete-distribution-network 0.2.5

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for discrete-distribution-network 0.2.5
File	Size	Uploaded
discrete_distribution_network-0.2.5.tar.gz	13.5 kB	Aug 2, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for discrete-distribution-network 0.2.5
File	Interpreter	ABI	Platform	Reset
discrete_distribution_network-0.2.5-py3-none-any.whl 13.1 kB Aug 2, 2026	Python 3	none	any	Details

Total release size: 26.6 kB

Release files / discrete_distribution_network-0.2.5.tar.gz

Download URL	discrete_distribution_network-0.2.5.tar.gz
Size	13.5 kB
Tags	Source
SHA-256 checksum How to use checksums	`4676e21c8b194a4b55aae52a7fd5a984ee9e5b50fd1faea895c3561751b5dd16`
BLAKE2b-256 checksum How to use checksums	`f39cbc354835b9739a4ec1e54658d5b9cb238a728e7b9c3ca7f60036656744cb`
Upload date	Aug 2, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`

Release files / discrete_distribution_network-0.2.5-py3-none-any.whl

Download URL	discrete_distribution_network-0.2.5-py3-none-any.whl
Size	13.1 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`58d9fa020fe394d518a9e8b80634bff478d302858f5de57f1dc045e19069d4cc`
BLAKE2b-256 checksum How to use checksums	`81d9a5d64d49fdda4f5e26a6af1b26057119333644db363ab1695702f96355fb`
Upload date	Aug 2, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`