Stochastic Differential Equations for Generation

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
// Simulate the 1D Variance Preserving (VP) SDE:
5
// dx = f(x,t) dt + g(t) dW, with f(x,t) = -0.5 * beta(t) * x, g(t) = sqrt(beta(t)).
6
// Beta schedule: linear between beta_min and beta_max.
7
8
struct VPSDE {
9
    double beta_min, beta_max; // schedule parameters
10
    explicit VPSDE(double bmin, double bmax) : beta_min(bmin), beta_max(bmax) {}
11
    inline double beta(double t) const { return beta_min + (beta_max - beta_min) * t; }
12
    inline double drift(double x, double t) const { return -0.5 * beta(t) * x; }
13
    inline double diffusion(double t) const { return sqrt(max(0.0, beta(t))); }
14
};
15
16
int main() {
17
    ios::sync_with_stdio(false);
18
    cin.tie(nullptr);
19
20
    // Simulation parameters
21
    const int M = 20000;    // number of independent paths
22
    const int N = 1000;     // number of time steps
23
    const double T = 1.0;   // final time
24
    const double dt = T / N;
25
26
    // VP SDE parameters (typical diffusion-model values)
27
    VPSDE sde(0.1, 20.0); // beta_min, beta_max
28
29
    // RNG setup
30
    random_device rd;
31
    mt19937_64 gen(rd());
32
    normal_distribution<double> normal01(0.0, 1.0);
33
34
    // Initialize x0 ~ N(0,1) to mimic standardized data
35
    vector<double> x(M);
36
    for (int i = 0; i < M; ++i) x[i] = normal01(gen);
37
38
    // Euler–Maruyama integration forward from t=0 to t=T
39
    double t = 0.0;
40
    for (int k = 0; k < N; ++k) {
41
        double g = sde.diffusion(t);
42
        double sqdt = sqrt(dt);
43
        for (int i = 0; i < M; ++i) {
44
            double drift = sde.drift(x[i], t);
45
            double noise = g * sqdt * normal01(gen);
46
            x[i] += drift * dt + noise; // EM step
47
        }
48
        t += dt;
49
    }
50
51
    // Empirical mean and variance at time T
52
    double mean = 0.0, var = 0.0;
53
    for (double xi : x) mean += xi;
54
    mean /= M;
55
    for (double xi : x) var += (xi - mean) * (xi - mean);
56
    var /= (M - 1);
57
58
    cout.setf(std::ios::fixed); cout << setprecision(6);
59
    cout << "Empirical mean at T: " << mean << "\n";
60
    cout << "Empirical variance at T: " << var << "\n";
61
62
    // For VP SDE with Var[x0]=1, the variance remains close to 1 theoretically.
63
    // We can compute alpha(t) = exp(-0.5 * \int_0^T beta(s) ds).
64
    auto integral_beta = [&](double Tfin){
65
        // For linear beta(t) = bmin + (bmax-bmin) t: integral = bmin T + 0.5 (bmax-bmin) T^2
66
        return sde.beta_min * Tfin + 0.5 * (sde.beta_max - sde.beta_min) * Tfin * Tfin;
67
    };
68
    double alpha = exp(-0.5 * integral_beta(T));
69
    double var_theory = (1 - alpha * alpha) + (alpha * alpha) * 1.0; // Var[x0]=1
70
    cout << "Theoretical variance (VP, Var[x0]=1): " << var_theory << "\n";
71
72
    return 0;
73
}
74

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
// Demonstrate reverse-time sampling without a neural net by choosing a driftless forward SDE
5
// with analytically Gaussian marginals. Forward SDE: dx = g(t) dW, f=0.
6
// Then Var[x_t] = Var[x_0] + \int_0^t g(s)^2 ds. Choose g(t) = sqrt(s), constant in t
7
// (here, s is just a scalar), so Var[x_t] = Var[x_0] + s t.
8
// We'll set Var[x_0] = 1 and s = 9, so Var[x_1] = 10. The score is ∂_x log p_t(x) = -x / Var[x_t].
9
// Reverse-time SDE: dx = [f - g(t)^2 * score] dt + g(t) d\bar W  (integrated from t=T down to 0).
10
// With f=0 and g^2 = s, the reverse drift is  (s / Var[x_t]) * x.
11
12
struct DriftlessGaussian {
13
    double var0; // Var at t=0
14
    double s;    // diffusion power: g^2 = s (constant)
15
    DriftlessGaussian(double var0_, double s_) : var0(var0_), s(s_) {}
16
    inline double g(double /*t*/) const { return sqrt(max(0.0, s)); }
17
    inline double var_t(double t) const { return var0 + s * t; }
18
    inline double score(double x, double t) const { return -x / var_t(t); } // ∂_x log p_t(x)
19
};
20
21
int main() {
22
    ios::sync_with_stdio(false);
23
    cin.tie(nullptr);
24
25
    // Model: Var[x_0] = 1, s=9 -> Var[x_1] = 10
26
    DriftlessGaussian model(1.0, 9.0);
27
28
    // Reverse-time integration settings
29
    const int M = 20000;    // number of samples to generate
30
    const int N = 1000;     // steps
31
    const double T = 1.0;   // start from t=T down to 0
32
    const double dt_pos = T / N; // positive step magnitude
33
34
    // RNG
35
    random_device rd; mt19937_64 gen(rd());
36
    normal_distribution<double> normal01(0.0, 1.0);
37
38
    // Initialize from the forward terminal distribution: x_T ~ N(0, Var[x_T]) with Var[x_T]=10
39
    double varT = model.var_t(T);
40
    normal_distribution<double> normal0_varT(0.0, sqrt(varT));
41
    vector<double> x(M);
42
    for (int i = 0; i < M; ++i) x[i] = normal0_varT(gen);
43
44
    // Integrate the reverse SDE from t=T to t=0 with negative dt
45
    double t = T;
46
    for (int k = 0; k < N; ++k) {
47
        double dt = -dt_pos;                 // negative because we go backward in time
48
        double g = model.g(t);
49
        double sqdt = sqrt(dt_pos);          // magnitude for Brownian increment
50
        for (int i = 0; i < M; ++i) {
51
            // reverse drift: -g^2 * score = s * x / Var[x_t]
52
            double rev_drift = (model.s / model.var_t(t)) * x[i];
53
            double noise = g * sqdt * normal01(gen); // uses d\bar W with variance |dt|
54
            x[i] += rev_drift * dt + noise;         // Euler–Maruyama backward step
55
        }
56
        t += dt; // decrease time
57
    }
58
59
    // Empirical mean/variance of generated samples (target: Var[x_0] ~ 1)
60
    double mean = 0.0, var = 0.0;
61
    for (double xi : x) mean += xi; mean /= M;
62
    for (double xi : x) var += (xi - mean) * (xi - mean); var /= (M - 1);
63
64
    cout.setf(std::ios::fixed); cout << setprecision(6);
65
    cout << "Generated mean at t=0: " << mean << "\n";
66
    cout << "Generated variance at t=0 (should be ~1): " << var << "\n";
67
68
    return 0;
69
}
70