From 161352cf4504c2113ad148ce5b5917a5bed36c06 Mon Sep 17 00:00:00 2001 From: y-jan137 Date: Wed, 18 Mar 2026 11:44:05 +0300 Subject: Add continuous normalizing flow --- src/cnf_train.c | 128 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 src/cnf_train.c (limited to 'src/cnf_train.c') diff --git a/src/cnf_train.c b/src/cnf_train.c new file mode 100644 index 0000000..21b357f --- /dev/null +++ b/src/cnf_train.c @@ -0,0 +1,128 @@ +#include "cnf_train.h" +#include "cnf.h" +#include "adam.h" + +#include +#include +#include + +static const double TARGET_MU[4][2] = { + { 1.5, 0.0}, {-1.5, 0.0}, + { 0.0, 1.5}, { 0.0, -1.5} +}; +static const double TARGET_SIGMA2 = 0.2; +static const int TARGET_K = 4; + +static double log_p_target(const double *z) { + double lc[4], mx = -1e300; + for (int k = 0; k < TARGET_K; k++) { + double dx = z[0] - TARGET_MU[k][0], dy = z[1] - TARGET_MU[k][1]; + lc[k] = -0.5 * (dx * dx + dy * dy) / TARGET_SIGMA2; + if (lc[k] > mx) mx = lc[k]; + } + double s = 0.0; + for (int k = 0; k < TARGET_K; k++) s += exp(lc[k] - mx); + return mx + log(s / (double)TARGET_K) + - log(2.0 * M_PI * TARGET_SIGMA2); +} + +static void grad_log_p_target(const double *z, double *g) { + double lc[4], mx = -1e300; + for (int k = 0; k < TARGET_K; k++) { + double dx = z[0] - TARGET_MU[k][0], dy = z[1] - TARGET_MU[k][1]; + lc[k] = -0.5 * (dx * dx + dy * dy) / TARGET_SIGMA2; + if (lc[k] > mx) mx = lc[k]; + } + double w[4], wsum = 0.0; + for (int k = 0; k < TARGET_K; k++) { w[k] = exp(lc[k] - mx); wsum += w[k]; } + g[0] = g[1] = 0.0; + for (int k = 0; k < TARGET_K; k++) { + double wk = w[k] / wsum; + g[0] += wk * (-(z[0] - TARGET_MU[k][0]) / TARGET_SIGMA2); + g[1] += wk * (-(z[1] - TARGET_MU[k][1]) / TARGET_SIGMA2); + } +} + +/* ---- Training loop ---- */ +void cnf_train_demo(RNG *r) { + const int D = 2, H = 32; + const int ITERS = 200, BATCH = 16, LOG_EVERY = 20; + const double t0 = 0.0, t1 = 1.0; + const double atol = 1e-4, rtol = 1e-4; + const double LR = 1e-3; + + int nparams = dynmlp_nparams(D, H); + double *theta = vec_alloc(nparams); + CNF cnf; + cnf_init(&cnf, D, H, theta, r); + Adam adam = adam_init(nparams, LR, 0.9, 0.999, 1e-8); + + printf("\n--- CNF density matching (4-Gaussian target, D=2, H=%d) ---\n\n", H); + printf("%-6s %-14s %-10s %-10s\n", "Iter", "Loss", "FwdNFE", "BwdNFE"); + printf("--------------------------------------------------\n"); + + for (int iter = 1; iter <= ITERS; iter++) { + double *dL_dtheta = vec_zeros(nparams); + double loss = 0.0; + int total_nfe_fwd = 0, total_nfe_bwd = 0; + + for (int b = 0; b < BATCH; b++) { + double z0[2]; + for (int i = 0; i < D; i++) z0[i] = rng_normal(r); + + CNFSampleResult sr = cnf_sample(&cnf, theta, z0, t0, t1, atol, rtol); + double *z1 = sr.z1; + total_nfe_fwd += sr.nfe; + + double log_pb = -0.5 * (z0[0]*z0[0] + z0[1]*z0[1]) + - (double)D * 0.5 * log(2.0 * M_PI); + + double log_pm = log_pb + sr.delta_logp; + double log_pt = log_p_target(z1); + + * = log_pb + delta_logp - log_pt + * (KL(p_model || p_target) estimator) */ + loss += log_pm - log_pt; + + double dL_dz1[2]; + grad_log_p_target(z1, dL_dz1); + dL_dz1[0] = -dL_dz1[0]; + dL_dz1[1] = -dL_dz1[1]; + + CNFBackwardResult br = cnf_backward(&cnf, theta, z1, + 1.0, dL_dz1, + t0, t1, atol, rtol); + total_nfe_bwd += br.nfe; + vec_add_scaled(dL_dtheta, 1.0 / BATCH, br.dL_dtheta, nparams); + free(br.dL_dz0); free(br.dL_dtheta); + free(z1); + } + + adam_update(&adam, theta, dL_dtheta); + free(dL_dtheta); + + if (iter % LOG_EVERY == 0) { + printf("%-6d %-14.4f %-10d %-10d\n", + iter, loss / BATCH, + total_nfe_fwd / BATCH, total_nfe_bwd / BATCH); + fflush(stdout); + } + } + + printf("\n--------------------------------------------------\n"); + + const int EVAL_N = 200; + double eval_loss = 0.0; + for (int i = 0; i < EVAL_N; i++) { + double z0[2]; + for (int j = 0; j < D; j++) z0[j] = rng_normal(r); + CNFSampleResult sr = cnf_sample(&cnf, theta, z0, t0, t1, atol, rtol); + eval_loss += log_p_target(sr.z1); + free(sr.z1); + } + printf("Eval mean log p_target of samples: %.4f\n", eval_loss / EVAL_N); + printf("Total parameters: %d\n", nparams); + + adam_free(&adam); + free(theta); +} -- cgit v1.2.3