Variational Inference Theory

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
// Model: z ~ N(0,1); x | z ~ N(z, sigma_x^2)
5
// Variational family: q(z|x) = N(mu, sigma^2)
6
// Optimize ELBO(mu, sigma) via reparameterization: z = mu + sigma * eps, eps ~ N(0,1)
7
8
struct RNG {
9
    mt19937 gen;
10
    normal_distribution<double> stdn{0.0, 1.0};
11
    RNG(unsigned seed = 42u) : gen(seed) {}
12
    double std_normal() { return stdn(gen); }
13
};
14
15
// Derivative of f(z) = log p(z) + log p(x|z) w.r.t z
16
// log p(z) = -0.5*(z^2 + c); d/dz = -z
17
// log p(x|z) = -0.5*((x - z)^2 / sigma_x^2 + c); d/dz = (x - z)/sigma_x^2
18
inline double df_dz(double z, double x, double sigma_x){
19
    return -z + (x - z)/(sigma_x * sigma_x);
20
}
21
22
// Compute a single-sample stochastic ELBO value for monitoring
23
inline double elbo_sample(double z, double mu, double sigma, double x, double sigma_x){
24
    // log p(z)
25
    double log_pz = -0.5*(z*z + log(2*M_PI));
26
    // log p(x|z)
27
    double log_px_z = -0.5*(((x - z)*(x - z))/(sigma_x*sigma_x) + log(2*M_PI*sigma_x*sigma_x));
28
    // log q(z|x)
29
    double log_q = -0.5*(((z - mu)*(z - mu))/(sigma*sigma) + 2.0*log(sigma) + log(2*M_PI));
30
    return log_pz + log_px_z - log_q;
31
}
32
33
int main(){
34
    // Observed datum
35
    const double x = 2.0;        // thermometer reading
36
    const double sigma_x = 0.5;  // known noise std
37
38
    // Variational parameters (initialize)
39
    double mu = 0.0;      // mean
40
    double rho = -0.5;    // log sigma (to keep sigma > 0)
41
42
    // Optimization hyperparameters
43
    const int K = 8;           // MC samples per iteration
44
    const int iters = 2000;    // iterations
45
    const double lr = 0.02;    // learning rate
46
47
    RNG rng(123);
48
49
    for(int t=1; t<=iters; ++t){
50
        double g_mu = 0.0;   // gradient w.r.t mu
51
        double g_rho = 0.0;  // gradient w.r.t rho
52
        double elbo_est = 0.0; // for monitoring
53
54
        double sigma = exp(rho);
55
        for(int k=0; k<K; ++k){
56
            double eps = rng.std_normal();
57
            double z = mu + sigma * eps; // reparameterization
58
59
            // Pathwise gradients for Gaussian q
60
            // dL/dmu = f'(z)
61
            double fp = df_dz(z, x, sigma_x);
62
            double dL_dmu = fp;
63
            // dL/dsigma = f'(z)*eps + 1/sigma  => chain to rho: dL/drho = dL/dsigma * sigma
64
            double dL_dsigma = fp * eps + 1.0 / sigma;
65
            double dL_drho = dL_dsigma * sigma;
66
67
            g_mu  += dL_dmu;
68
            g_rho += dL_drho;
69
70
            elbo_est += elbo_sample(z, mu, sigma, x, sigma_x);
71
        }
72
        g_mu  /= K; g_rho /= K; elbo_est /= K;
73
74
        // Gradient ascent on ELBO
75
        mu  += lr * g_mu;
76
        rho += lr * g_rho;
77
78
        if(t % 200 == 0){
79
            cout << "iter " << t
80
                 << "\tELBO~ " << fixed << setprecision(4) << elbo_est
81
                 << "\tmu= " << mu
82
                 << "\tsigma= " << exp(rho) << "\n";
83
        }
84
    }
85
86
    cout << "Final: mu= " << mu << ", sigma= " << exp(rho) << "\n";
87
    return 0;
88
}
89

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
// Target posterior (up to normalization): p(z) ∝ exp(-0.5 * (z - m)^T Λ (z - m))
5
// We approximate with mean-field q(z) = Π_i N(m_i_hat, s2_i), where s2_i = 1 / Λ_ii
6
// Coordinate update (Gaussian case):
7
//   m_i_hat ← m_i - (1/Λ_ii) * Σ_{j≠i} Λ_ij * m_j_hat
8
// Repeat until convergence.
9
10
struct MeanFieldGaussian {
11
    vector<double> m_true;     // target mean m
12
    vector<vector<double>> Lambda; // precision matrix Λ (symmetric PD)
13
    vector<double> m_hat;      // variational means
14
    vector<double> s2;         // variational variances (fixed): 1 / Λ_ii
15
16
    MeanFieldGaussian(const vector<double>& m, const vector<vector<double>>& L)
17
        : m_true(m), Lambda(L) {
18
        int n = (int)m.size();
19
        m_hat.assign(n, 0.0);
20
        s2.resize(n);
21
        for(int i=0;i<n;++i){
22
            s2[i] = 1.0 / Lambda[i][i];
23
        }
24
    }
25
26
    void run(int max_iters = 1000, double tol = 1e-9){
27
        int n = (int)m_hat.size();
28
        for(int it=1; it<=max_iters; ++it){
29
            double max_delta = 0.0;
30
            for(int i=0; i<n; ++i){
31
                double off = 0.0;
32
                for(int j=0; j<n; ++j){
33
                    if(j==i) continue;
34
                    off += Lambda[i][j] * m_hat[j];
35
                }
36
                double new_mi = m_true[i] - off / Lambda[i][i];
37
                max_delta = max(max_delta, fabs(new_mi - m_hat[i]));
38
                m_hat[i] = new_mi;
39
            }
40
            if(max_delta < tol){
41
                // converged
42
                break;
43
            }
44
        }
45
    }
46
47
    void print_state() const{
48
        int n = (int)m_hat.size();
49
        cout << fixed << setprecision(6);
50
        cout << "Variational means (m_hat): ";
51
        for(int i=0;i<n;++i) cout << m_hat[i] << (i+1==n?"\n":" ");
52
        cout << "Variational variances (s2): ";
53
        for(int i=0;i<n;++i) cout << s2[i] << (i+1==n?"\n":" ");
54
    }
55
};
56
57
int main(){
58
    // 2D example with correlation
59
    vector<double> m = {1.0, -2.0};
60
    // Precision Λ for covariance Σ = [[1, 0.8],[0.8, 1]] (PD); Λ = Σ^{-1}
61
    double rho = 0.8;
62
    double denom = 1.0 - rho*rho;
63
    vector<vector<double>> Lambda = {
64
        {  1.0/denom, -rho/denom },
65
        { -rho/denom,  1.0/denom }
66
    };
67
68
    MeanFieldGaussian mf(m, Lambda);
69
    mf.run(10000, 1e-12);
70
    mf.print_state();
71
72
    return 0;
73
}
74