RasmahRasmah

Generative Design

println("Julia ", VERSION)
Julia 1.12.6

Overview

Generative design trains a model to sample new designs from a learned distribution, rather than optimize a single one. Given a dataset of designs, a generative model captures their manifold and then produces new, valid members of it. Rasmah implements the main families from the literature: the variational autoencoder (VAE), the denoising diffusion model (DDPM/DDIM) with classifier-free guidance, and the score-based SDE (VE/VP).

using Rasmah
using Random
using Statistics

A latent field manifold (VAE)

A VAE compresses each design to a low-dimensional latent code and reconstructs it. We build a dataset of smooth fields a·sin(2πt) + b·cos(2πt) and learn a 4-dimensional latent space:

rng = MersenneTwister(8)
d = 16
t = collect(range(0.0, 1.0; length = d))
N = 200
X = zeros(d, N)
for i in 1:N
    a, b = randn(rng), randn(rng)
    X[:, i] = a .* sin.(2π .* t) .+ b .* cos.(2π .* t)
end

lat = 4
enc = MLP([d, 32, 2 * lat]); dec = MLP([lat, 32, d])
vae = VAE(enc, dec, lat)
θv = vae_fit(vae, [xavier_init(rng, enc.sizes); xavier_init(rng, dec.sizes)], X;
             epochs = 500, lr = 1e-3, seed = 8)
1496-element Vector{Float64}:
  0.04751476105106996
 -0.2113136787376715
  0.0632178481169035
  0.10803879172804234
  0.24612956291786253
  0.13892647180040918
 -0.2825224674422241
 -0.09847107776188
  0.08290501017299898
  0.0735076419414228
  ⋮
 -0.019741391287818547
 -0.004715092366760297
  0.0026543578526247976
 -0.011076730872440835
 -0.011584439671856548
  0.07338273382112644
  0.05200570808915593
  0.06966931813405064
  0.043330339491668056

The reconstruction error measures how well the manifold captures the data; a latent interpolation between two designs should stay on the manifold:

re = mean(sum(abs2, vae_reconstruct(vae, θv, X[:, i]) .- X[:, i]) / d for i in 1:N)
println("mean reconstruction MSE = ", round(re; digits = 4))
z1 = vae_encode(vae, θv, X[:, 1]); z2 = vae_encode(vae, θv, X[:, 2])
println("interpolation finite = ", all(isfinite, vae_decode(vae, θv, 0.5 .* z1 .+ 0.5 .* z2)))
mean reconstruction MSE = 0.0905
interpolation finite = true

Diffusion (DDPM) over the designs

A diffusion model noisifies designs over a schedule and learns to reverse it, sampling new designs from pure noise. We fit a DDPM on a 1-D Gaussian design family and sample:

d1, N1 = 1, 500
X1 = reshape(1.0 .+ 0.5 .* randn(rng, d1 * N1), d1, N1)
ddpm = DDPM(MLP([2, 32, 32, 1]), 50)
θd = ddpm_fit(ddpm, xavier_init(rng, [2, 32, 32, 1]), X1; epochs = 600, lr = 1e-3, seed = 9)
samples = [ddpm_sample(ddpm, θd, rng)[1] for _ in 1:200]
println("DDPM samples: mean = ", round(mean(samples); digits = 3), ", std = ", round(std(samples); digits = 3),
        "  (target 1.0 / 0.5)")
DDPM samples: mean = 0.72, std = 0.544  (target 1.0 / 0.5)

Score-based SDE (VE)

The score-based SDE generalizes diffusion to a continuous process: learn the score ∇ₓ log pₜ(x) and sample by reversing the SDE. VESDE/ve_fit/ve_sample do this for the variance-exploding process:

ve = VESDE(σmin = 0.01, σmax = 5.0)
net_s = MLP([2, 32, 32, 1])
θs = ve_fit(ve, net_s, xavier_init(rng, net_s.sizes), X1; epochs = 600, lr = 1e-3, seed = 31)
s2 = [ve_sample(ve, net_s, θs, rng; d = d1, steps = 200, langevin_steps = 1)[1] for _ in 1:200]
println("VE-SDE samples: mean = ", round(mean(s2); digits = 3), ", std = ", round(std(s2); digits = 3))
VE-SDE samples: mean = 1.176, std = 1.112

Next steps