From Noise to Geometry: Training a Diffusion Model to Draw Polygons

Polydiff is a diffusion model that learns a distribution over 2D polygons and generates new ones starting from Gaussian noise.

I picked polygons for two reasons. They make a small sandbox for learning how generative models work, without the overhead of images or huge datasets. And generated polygons have uses of their own: synthetic geometric data feeds simulation, CAD-style workflows and shape priors, and later I want to condition on structure I can control.

Project repository: github.com/Jerenemy/polydiff

What I’m building

I’m training a DDPM-style model (polydiff) on synthetic near-regular polygons, hexagons for now. The workflow:

1. Generate a dataset of polygons with controlled deformation.
2. Train a diffusion denoiser to predict noise at random timesteps.
3. Sample new polygons by denoising step by step from random noise.
4. Score the samples with the same regularity metric used to generate the data.

Why polygons

Everything is small enough to inspect, so I can debug failures and iterate quickly on questions like these:

• How well does a simple MLP denoiser recover geometric structure?
• How does quality change with the number of diffusion steps or the architecture?
• Which inductive biases matter most before moving to heavier models such as GNNs or equivariant networks?

Technical details

Data generation

Training data comes from polydiff.data.gen_polygons.

• Each sample is an n-gon with smooth radial and angular jitter.
• Polygons are canonicalized by centering, RMS-scale normalization and CCW ordering.
• Self-intersecting samples can optionally be rejected.
• A continuous regularity score is computed from edge, angle and radius variation:
• higher score = more regular polygon
• lower score = more distorted polygon

The current training dataset (data/raw/hexagons.npz) has:

• 10,000 polygons
• n=6 vertices per polygon
• a stored regularity score for each sample

Model and diffusion process

The denoiser in polydiff/models/diffusion.py is a timestep-conditioned MLP:

• input: flattened polygon coordinates (2 * n_vertices) + sinusoidal timestep embedding
• output: predicted Gaussian noise (epsilon)

Diffusion config used in training:

• n_steps = 1000
• linear beta schedule from 1e-4 to 2e-2
• objective: MSE between true noise and predicted noise (standard DDPM epsilon prediction)

Training config

From configs/train_diffusion.yaml:

• batch size: 128
• epochs: 100
• learning rate: 2e-4
• hidden dim: 256
• time embedding dim: 64
• MLP layers: 3

Checkpoints are saved every 500 steps, with a final checkpoint at:

• models/diffusion_final.pt

With 10,000 samples and batch size 128, the run takes about 7,900 optimizer steps.

Preliminary results

I evaluated the current sample file (data/processed/samples.npz, 64 hexagons generated by diffusion_final.pt) against the training data, using the project’s regularity score and self-intersection check.

None of the 64 generated polygons intersects itself. Their regularity is well below the training data, though (mean 0.47 against 0.79), so the model has a long way to go before it matches the geometry it was trained on. For a first baseline that’s a reasonable place to start.

Next

• Evaluate larger generated sets than 64 for tighter statistics.
• Compare fewer sampling steps (sampling.n_steps) against quality and speed.
• Move from the flat MLP denoiser to graph-aware and eventually E(2)-equivariant architectures, to see which inductive biases improve the geometry most.

Reproducibility commands

# Train
python3 -m polydiff.training.train --config configs/train_diffusion.yaml

# Sample
python3 -m polydiff.sampling.sample --config configs/sample_diffusion.yaml

# Visualize generated polygons and compute scores
python3 -m polydiff.data.plot_polygons data/processed/samples.npz --num 16 --compute_score