Multivariate linear regression with gaussian mixture pool

[max-de-rooij]

So my problem is similar to what I wanted to achieve yesterday, but I decided to give myself a hard time with a more complicated example. Fortunately, this is all still dummy data so I can share everything (and perhaps this can be a cool example in RxInferExamples.jl if this works).

Anyway, on to the fun stuff. I have a simulated dataset of a modest ~1200 individuals with 2 variables changing over time:

Don't mind the training/test data labels, for now.

I want to do something similar to the partial pooling example, but instead of drawing the regression parameters for each individual from a single distribution for the population, the population distribution for the regression parameters is a Gaussian mixture of $n$ groups. In this case, I selected $n=3$, which matches the simulated data.

I decided to do a similar model setup as with the partial pooling example:

@model function latent_class_lm(individuals, time, parameters, data)

    # fetch information
    n_classes = parameters.n_classes
    prior_mean_a = parameters.prior_mean_a
    prior_var_a = parameters.prior_var_a
    prior_mean_b = parameters.prior_mean_b
    prior_var_b = parameters.prior_var_b
    prior_s = parameters.prior_s
    prior_w = parameters.prior_w
    dim = length(prior_mean_a[1])

    n_data_points = length(individuals)
    n_individuals = length(unique(individuals))

    # global selection variable
    s ~ Dirichlet(probvec(prior_s))

    local as
    local ws
    local bs

    # class level priors
    for i in 1:n_classes
        as[i] ~ MvNormal(mean = prior_mean_a[i], covariance = diageye(dim)*prior_var_a[i])
        bs[i] ~ MvNormal(mean = prior_mean_b[i], covariance = diageye(dim)*prior_var_b[i])
        ws[i] ~ Wishart(dim + 2, prior_w[i] * diageye(dim))
    end

    # patient-level parameters
    local α
    local β
    local z
    for i in 1:n_individuals
        z[i] ~ RxInfer.Categorical(s)
        α[i] ~ NormalMixture(switch = z[i], m = as, p = ws)
        β[i] ~ NormalMixture(switch = z[i], m = bs, p = ws)
    end

    σ ~ Gamma(shape = 1.75, scale = 45.0)

    local estimation

    for i in 1:n_data_points
        estimation[i] := α[individuals[i]] + β[individuals[i]] * time[i]
        data[i] ~ MvNormal(mean = estimation[i], covariance = diageye(dim) * σ)
    end

end

Right now, I use the same prior for the precision of both regression coefficients, which is not completely intentional. I then proceed to initialisation and inference:

n_classes = 3
priors = LatentClassParameters(
    n_classes,
    [zeros(2) for i in 1:n_classes],
    [10.0 for i in 1:n_classes],
    [zeros(2) for i in 1:n_classes],
    [10.0 for i in 1:n_classes],
    [1e2 for i in 1:n_classes],
    Dirichlet(ones(n_classes))
)

init_mixture = @initialization begin
    μ(as) = [MvNormalMeanCovariance(zeros(2), 10.0*diageye(2)) for i in 1:priors.n_classes]
    q(bs) = [MvNormalMeanCovariance(zeros(2), 10.0*diageye(2)) for i in 1:priors.n_classes]
    q(ws) = [Wishart(2 + 2, 1e2*diageye(2)) for i in 1:priors.n_classes]
    q(σ) = vague(Gamma)
    q(s) = vague(Dirichlet, priors.n_classes)
    q(z) = vague(RxInfer.Categorical, priors.n_classes)
end

results_mixture = infer(
    model           = latent_class_lm(individuals=individuals, time=timepoints, parameters=priors,), 
    data            = (data = train_data,), 
    initialization  = init_mixture, 
    returnvars      = KeepLast(),
    options         = (limit_stack_depth = 100,),
    iterations      = 30,
    constraints     = MeanField(),
)

But unfortunately, I get the following error message:

Variables [ α, as, σ, s, estimation, bs, z, β, ws ] have not been updated after an update event. 

Would you have any improvements in model specification that enable me to do this?

[bvdmitri]

My guess would be is that you also need to initialize messages on α and β, e.g.

init_mixture = @initialization begin
    μ(as) = [MvNormalMeanCovariance(zeros(2), 10.0*diageye(2)) for i in 1:priors.n_classes]
    q(bs) = [MvNormalMeanCovariance(zeros(2), 10.0*diageye(2)) for i in 1:priors.n_classes]
    q(ws) = [Wishart(2 + 2, 1e2*diageye(2)) for i in 1:priors.n_classes]
    q(σ) = vague(Gamma)
    q(s) = vague(Dirichlet, priors.n_classes)
    q(z) = vague(RxInfer.Categorical, priors.n_classes)
    
    μ(α) = ...
    μ(β) = ...
end

if that doesn't help, then perhaps also initialize posteriors on α and β, e.g.

init_mixture = @initialization begin
    μ(as) = [MvNormalMeanCovariance(zeros(2), 10.0*diageye(2)) for i in 1:priors.n_classes]
    q(bs) = [MvNormalMeanCovariance(zeros(2), 10.0*diageye(2)) for i in 1:priors.n_classes]
    q(ws) = [Wishart(2 + 2, 1e2*diageye(2)) for i in 1:priors.n_classes]
    q(σ) = vague(Gamma)
    q(s) = vague(Dirichlet, priors.n_classes)
    q(z) = vague(RxInfer.Categorical, priors.n_classes)
    
    μ(α) = ...
    μ(β) = ...
    q(α) = ...
    q(β) = ...
end

[max-de-rooij]

That seems to work, but now I get this error that I cannot really trace:

DomainError with [NaN, NaN, NaN]:
Categorical: vector p is not a probability vector

Any ideas?

[bvdmitri]

Hm, not immediately, no, my guess would be that the initialization with vague and zeros makes the computations in message passing rules somewhat unstable for this model.. I can allocate some time this week to debug your example it if you post the update model you used with initialization and maybe some dummy data we can try to run the example with.

We are also planning to refactor the mixture nodes, but that will probably take several months to complete, hopefully it will be easier to debug afterwards.

[max-de-rooij]

That would be awesome. I've got the model specs right here:

# read data
data = CSV.File("data/selected_individuals.csv") |> DataFrame

individuals = data.individual
timepoints = data.time
train_data = [[data.var1[i], data.var2[i]] for i in 1:size(data,1)]

struct LatentClassParameters
    n_classes::Int
    prior_mean_a
    prior_var_a
    prior_mean_b
    prior_var_b
    prior_w
    prior_s
end

@model function latent_class_lm(individuals, time, parameters, data)

    # fetch information
    n_classes = parameters.n_classes
    prior_mean_a = parameters.prior_mean_a
    prior_var_a = parameters.prior_var_a
    prior_mean_b = parameters.prior_mean_b
    prior_var_b = parameters.prior_var_b
    prior_s = parameters.prior_s
    prior_w = parameters.prior_w
    dim = length(prior_mean_a[1])

    n_data_points = length(individuals)
    n_individuals = length(unique(individuals))

    # global selection variable
    s ~ Dirichlet(probvec(prior_s))

    local as
    local ws
    local bs

    # class level priors
    for i in 1:n_classes
        as[i] ~ MvNormal(mean = prior_mean_a[i], covariance = diageye(dim)*prior_var_a[i])
        bs[i] ~ MvNormal(mean = prior_mean_b[i], covariance = diageye(dim)*prior_var_b[i])
        ws[i] ~ Wishart(dim + 2, prior_w[i] * diageye(dim))
    end

    # patient-level parameters
    local α
    local β
    local z
    for i in 1:n_individuals
        z[i] ~ RxInfer.Categorical(s)
        α[i] ~ NormalMixture(switch = z[i], m = as, p = ws)
        β[i] ~ NormalMixture(switch = z[i], m = bs, p = ws)
    end

    σ ~ Gamma(shape = 1.75, scale = 45.0)

    local estimation

    for i in 1:n_data_points
        estimation[i] := α[individuals[i]] + β[individuals[i]] * time[i]
        data[i] ~ MvNormalMeanScalePrecision(estimation[i], σ)
    end

end


n_classes = 3
priors = LatentClassParameters(
    n_classes,
    [zeros(2) for i in 1:n_classes],
    [10.0 for i in 1:n_classes],
    [zeros(2) for i in 1:n_classes],
    [10.0 for i in 1:n_classes],
    [1e2 for i in 1:n_classes],
    Dirichlet(ones(n_classes))
)

init_mixture = @initialization begin
    μ(as) = [MvNormalMeanCovariance(zeros(2), 10.0*diageye(2)) for i in 1:priors.n_classes]
    q(bs) = [MvNormalMeanCovariance(zeros(2), 10.0*diageye(2)) for i in 1:priors.n_classes]
    q(ws) = [Wishart(2 + 2, 1*diageye(2)) for i in 1:priors.n_classes]
    q(σ) = vague(Gamma)
    q(s) = vague(Dirichlet, priors.n_classes)
    q(z) = vague(RxInfer.Categorical, priors.n_classes)

    μ(α) = MvNormalMeanPrecision(zeros(2), 10.0*diageye(2))
    μ(β) = MvNormalMeanPrecision(zeros(2),  10.0*diageye(2))
    q(α) = MvNormalMeanPrecision(zeros(2),  10.0*diageye(2))
    q(β) = MvNormalMeanPrecision(zeros(2),  10.0*diageye(2))
end;

results_mixture = infer(
    model           = latent_class_lm(individuals=individuals, time=timepoints, parameters=priors,), 
    data            = (data = train_data,), 
    initialization  = init_mixture, 
    returnvars      = KeepLast(),
    options         = (limit_stack_depth = 100,),
    iterations      = 30,
    constraints     = MeanField(),
)

I've uploaded the data here, which is a smaller set of the data I'm currently using. In the end this should be a polynomial regression, but I guess I wanted to start with a linear regression first.
selected_individuals.csv

[bvdmitri]

Hey, I looked into it, and it does seem strange. However, I think I identified the main issue.

1. Index Out of Bounds Error

In the example you shared, the code didn’t run at all because the individuals array contained indices larger than the length of α, causing an out-of-bounds error when accessing α[individuals[i]]. I fixed this by mapping individuals to a valid range:

map_individuals = Dict()
counter = 1
foreach(individuals) do individual
    if !haskey(map_individuals, individual)
        map_individuals[individual] = counter
        global counter += 1
    end
end
fixed_individuals = map(i -> map_individuals[i], individuals)

I then used fixed_individuals instead of individuals in the rest of the code.

2. Issue with timepoints

The real problem was with time = timepoints, which contained 0 values. This caused a multiplication error, leading to infinite variance. I fixed it by shifting the time values:

time = timepoints .+ 1

This allowed the inference to run.

3. Poor Inference due to Initialization

Initially, all classes collapsed, and the inference results were poor. To fix this, I changed the initialization to be random:

n_classes = 6  # Arbitrary choice for fun
priors = LatentClassParameters(
    n_classes,
    [rand(2) for _ in 1:n_classes],
    [10.0 for _ in 1:n_classes],
    [rand(2) for _ in 1:n_classes],
    [10.0 for _ in 1:n_classes],
    [1e2 for _ in 1:n_classes],
    Dirichlet(rand(n_classes))
)

init_mixture = @initialization begin
    μ(as) = [MvNormalMeanCovariance(rand(2), 10.0 * diageye(2)) for _ in 1:priors.n_classes]
    q(bs) = [MvNormalMeanCovariance(rand(2), 10.0 * diageye(2)) for _ in 1:priors.n_classes]
    q(ws) = [Wishart(2 + 2, 1 * diageye(2)) for _ in 1:priors.n_classes]
    q(σ) = vague(Gamma)
    q(s) = Dirichlet(rand(priors.n_classes))
    q(z) = vague(RxInfer.Categorical, priors.n_classes)

    μ(α) = MvNormalMeanPrecision(rand(2), 10.0 * diageye(2))
    μ(β) = MvNormalMeanPrecision(rand(2), 10.0 * diageye(2))
    q(α) = MvNormalMeanPrecision(rand(2), 10.0 * diageye(2))
    q(β) = MvNormalMeanPrecision(rand(2), 10.0 * diageye(2))
end

This significantly improved the results. For reproducibility, I recommend using:

rng = StableRNG(some_seed)  # From the StableRNGs package
rand(rng, 2)

4. Bug in Free Energy Calculation

While checking the free energy for algorithm convergence, I encountered another issue. The following line:

estimation[i] := α[individuals[i]] + β[individuals[i]] * time[i]

needed to be changed to:

estimation[i] := α[individuals[i]] + time[i] * β[individuals[i]]

Silly bug in RxInfer, but the symmetric rules weren’t properly defined. After fixing this, I observed a nice convergence in the free energy.

Final Result

After these changes, the model now runs correctly and produces better inference results. Here’s the free energy convergence plot:

Let me know if you need any further clarifications!

Full code

using RxInfer, CSV, DataFrames, Plots, LinearAlgebra

# read data
data = CSV.File("selected_individuals.csv") |> DataFrame

individuals = data.individual
timepoints = data.time
train_data = [[data.var1[i], data.var2[i]] for i in 1:size(data,1)];

map_individuals = Dict()
counter = 1
foreach(individuals) do individual
    if !haskey(map_individuals, individual)
        map_individuals[individual] = counter
        global counter = counter + 1
    end
end
fixed_individuals = map(i -> map_individuals[i], individuals);

struct LatentClassParameters
    n_classes::Int
    prior_mean_a
    prior_var_a
    prior_mean_b
    prior_var_b
    prior_w
    prior_s
end

@model function latent_class_lm(individuals, time, parameters, data)

    # fetch information
    n_classes = parameters.n_classes
    prior_mean_a = parameters.prior_mean_a
    prior_var_a = parameters.prior_var_a
    prior_mean_b = parameters.prior_mean_b
    prior_var_b = parameters.prior_var_b
    prior_s = parameters.prior_s
    prior_w = parameters.prior_w
    dim = length(prior_mean_a[1])

    n_data_points = length(individuals)
    n_individuals = length(unique(individuals))

    # global selection variable
    s ~ Dirichlet(probvec(prior_s))

    local as
    local ws
    local bs

    # class level priors
    for i in 1:n_classes
        as[i] ~ MvNormal(mean = prior_mean_a[i], covariance = diageye(dim)*prior_var_a[i])
        bs[i] ~ MvNormal(mean = prior_mean_b[i], covariance = diageye(dim)*prior_var_b[i])
        ws[i] ~ Wishart(dim + 2, prior_w[i] * diageye(dim))
    end

    # patient-level parameters
    local α
    local β
    local z
    for i in 1:n_individuals
        z[i] ~ RxInfer.Categorical(s)
        α[i] ~ NormalMixture(switch = z[i], m = as, p = ws)
        β[i] ~ NormalMixture(switch = z[i], m = bs, p = ws)
    end

    σ ~ Gamma(shape = 1.75, scale = 45.0)

    local estimation

    for i in 1:n_data_points
        estimation[i] := α[individuals[i]] + time[i] * β[individuals[i]]
        data[i] ~ MvNormalMeanScalePrecision(estimation[i], σ)
    end

end


n_classes = 6
priors = LatentClassParameters(
    n_classes,
    [rand(2) for i in 1:n_classes],
    [10.0 for i in 1:n_classes],
    [rand(2) for i in 1:n_classes],
    [10.0 for i in 1:n_classes],
    [1e2 for i in 1:n_classes],
    Dirichlet(rand(n_classes))
)

init_mixture = @initialization begin
    μ(as) = [MvNormalMeanCovariance(rand(2), 10.0*diageye(2)) for i in 1:priors.n_classes]
    q(bs) = [MvNormalMeanCovariance(rand(2), 10.0*diageye(2)) for i in 1:priors.n_classes]
    q(ws) = [Wishart(2 + 2, 1*diageye(2)) for i in 1:priors.n_classes]
    q(σ) = vague(Gamma)
    q(s) = Dirichlet(rand(priors.n_classes))
    q(z) = vague(RxInfer.Categorical, priors.n_classes)

    μ(α) = MvNormalMeanPrecision(rand(2), 10.0*diageye(2))
    μ(β) = MvNormalMeanPrecision(rand(2),  10.0*diageye(2))
    q(α) = MvNormalMeanPrecision(rand(2),  10.0*diageye(2))
    q(β) = MvNormalMeanPrecision(rand(2),  10.0*diageye(2))
end;

results_mixture = infer(
    model           = latent_class_lm(individuals=fixed_individuals, time=timepoints .+ 1, parameters=priors,), 
    data            = (data = train_data,), 
    initialization  = init_mixture, 
    returnvars      = KeepLast(),
    options         = (limit_stack_depth = 100,),
    iterations      = 20,
    constraints     = MeanField(),
    catch_exception = true,
    free_energy     = true,
    # postprocess = NoopPostprocess()
    # addons          = (AddonMemory(), )
)

[bvdmitri]

Ah, another thing, this piece of the model most likely will not work out

covariance = diageye(dim) * σ

we don't have specialized message passing rules for this case, but we have a specialized factor node for this.
you should replace

data[i] ~ MvNormal(mean = estimation[i], covariance = diageye(dim) * σ)

with

data[i] ~ MvNormalMeanScalePrecision(estimation[i], σ) # equivalent to the line above and uses conjugate `Gamma` prior too