ガウス分布によるベイズ学習

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

import jax
import jax.numpy as jnp
from genjax import gen, normal, ChoiceMap
import jax.random as random

# Known parameters
DATA_STD = 2.0   # observation noise: N(mu, 4) means std 2

@gen
def generative_model():
    """Full model: prior on the mean, then one observation given that mean."""
    # Prior belief about the mean
    mu = normal(500.0, 5.0) @ "mu"           # Prior: N(500, 25)
    # Generate the observation from N(mu, 4)
    weight = normal(mu, DATA_STD) @ "weight_0"
    return mu

# Observe one bento: 497g. We CONDITION on it with a ChoiceMap (the modern way to
# pin an addressed random choice to an observed value), then importance-sample.
observed = ChoiceMap.d({"weight_0": 497.0})

# Run importance sampling to approximate the posterior over mu
key = random.PRNGKey(42)
num_samples = 10000

# model.importance(key, constraints, args) returns (trace, log_weight): the trace is
# a sample with weight_0 pinned to 497, and log_weight scores how well its mu explains
# the data. A weighted average of mu over many samples approximates the posterior.
mus = []
log_weights = []
for _ in range(num_samples):
    key, subkey = random.split(key)
    trace, log_w = generative_model.importance(subkey, observed, ())
    mus.append(trace.get_choices()["mu"])
    log_weights.append(log_w)

mus = jnp.array(mus)
log_weights = jnp.array(log_weights)

# Normalize the importance weights (log-space, like the GenJAX tutorial), then take
# the weighted mean and std of mu.
w = jnp.exp(log_weights - log_weights.max())
w = w / w.sum()
post_mean = jnp.sum(w * mus)
post_std = jnp.sqrt(jnp.sum(w * (mus - post_mean) ** 2))

print(f"Posterior mean: {post_mean:.2f}g")
print(f"Posterior std dev: {post_std:.2f}g")
print(f"Theoretical posterior mean: 497.4g")
print(f"Theoretical posterior std dev: 1.86g")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

import jax.numpy as jnp

def gaussian_gaussian_update(prior_mu, prior_var, data, data_var):
    """
    Analytical Bayesian update for Gaussian-Gaussian conjugate prior

    Args:
        prior_mu: Prior mean
        prior_var: Prior variance
        data: List of observations
        data_var: Known data variance

    Returns:
        posterior_mu, posterior_var
    """
    n = len(data)
    sample_mean = jnp.mean(jnp.array(data))

    # Precision-weighted update
    prior_precision = 1.0 / prior_var
    data_precision = n / data_var

    posterior_precision = prior_precision + data_precision
    posterior_var = 1.0 / posterior_precision

    posterior_mu = posterior_var * (prior_precision * prior_mu +
                                     data_precision * sample_mean)

    return posterior_mu, posterior_var

# Apply to our example
prior_mu, prior_var = 500.0, 25.0
data = [497.0]
data_var = 4.0

post_mu, post_var = gaussian_gaussian_update(prior_mu, prior_var, data, data_var)
post_std = jnp.sqrt(post_var)

print(f"After 1 observation:")
print(f"  Posterior mean: {post_mu:.2f}g")
print(f"  Posterior std dev: {post_std:.2f}g")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

# Additional observations
import jax.numpy as jnp

all_data = [497.0, 498.5, 496.0, 499.0, 497.5, 498.0, 496.5, 497.0, 498.5, 497.5]

# Start with prior
mu, var = 500.0, 25.0

print(f"Prior: N({mu:.2f}, {var:.2f})")
print(f"  Mean: {mu:.2f}g, Std dev: {jnp.sqrt(var):.2f}g\n")

# Update with each observation sequentially
for i, obs in enumerate(all_data, 1):
    mu, var = gaussian_gaussian_update(mu, var, [obs], data_var)
    std = jnp.sqrt(var)
    print(f"After observation {i} (x={obs}g):")
    print(f"  Posterior: N({mu:.2f}, {var:.2f})")
    print(f"  Mean: {mu:.2f}g, Std dev: {std:.2f}g")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

import matplotlib.pyplot as plt
from scipy.stats import norm as scipy_norm

# Prior
import jax.numpy as jnp

x_range = jnp.linspace(490, 505, 1000)
prior_pdf = scipy_norm.pdf(x_range, 500, 5)

# After 1, 5, and 10 observations
results = []
mu, var = 500.0, 25.0
for i, obs in enumerate(all_data):
    mu, var = gaussian_gaussian_update(mu, var, [obs], data_var)
    if i + 1 in [1, 5, 10]:
        results.append((i + 1, mu, jnp.sqrt(var)))

# Plot

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

# Posterior from before
import jax.numpy as jnp

post_mu = 497.65
post_var = 0.37

# Predictive distribution
pred_mu = post_mu  # Same mean
pred_var = data_var + post_var  # 4.0 + 0.37 = 4.37
pred_std = jnp.sqrt(pred_var)

print(f"Posterior for μ: N({post_mu:.2f}, {post_var:.2f})")
print(f"Predictive for next X: N({pred_mu:.2f}, {pred_var:.2f})")
print(f"  Predictive std dev: {pred_std:.2f}g")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

import jax
import jax.numpy as jnp
from genjax import gen, normal
import jax.random as random

@gen
def posterior_predictive(post_mu, post_var, data_var):
    """
    Sample from posterior predictive distribution
    """
    # First, sample a μ from the posterior
    mu = normal(post_mu, jnp.sqrt(post_var)) @ "mu"

    # Then, sample a new observation given that μ
    x_new = normal(mu, jnp.sqrt(data_var)) @ "x_new"

    return x_new

# Posterior from before (after observing 497g)
post_mu = 497.65
post_var = 0.37
data_var = 4.0

# Theoretical posterior-predictive: same mean, variance = posterior var + data var
pred_mu = post_mu
pred_std = jnp.sqrt(post_var + data_var)

# Simulate 10,000 predictions
key = random.PRNGKey(42)
predictions = []

for _ in range(10000):
    key, subkey = random.split(key)
    # model.simulate(key, args) runs the model once; args is the tuple of arguments.
    trace = posterior_predictive.simulate(subkey, (post_mu, post_var, data_var))
    predictions.append(trace.get_retval())

predictions = jnp.array(predictions)

print(f"Simulated predictive mean: {jnp.mean(predictions):.2f}g")
print(f"Simulated predictive std: {jnp.std(predictions):.2f}g")
print(f"Theoretical predictive mean: {pred_mu:.2f}g")
print(f"Theoretical predictive std: {pred_std:.2f}g")