discrete-distribution-network 0.1.1

Discrete Distribution Network

Discrete Distribution Network (wip)

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

$ discrete-distribution-network

Usage

import torch
from discrete_distribution_network.ddn 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)

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}, 
}