Regression with mixture on coefficients

[max-de-rooij]

Update 1: I got inference to run, but my outcome does not make sense to me.

First of all, I really appreciate the examples and the extensive documentation of this package. I am working on analysis of latent profiles and I want to try out this package as a possible alternative to the flexmix R package. The problem is as follows: I have data over time of a number of individuals belonging to a specific class. For my first attempt I wanted to combine the gaussian mixtures example with the regression example and specify my regression coefficient $a$ as a mixture of Gaussians.

I tried the following model specification:

@model function mixture_regression(x, y)

    s ~ Beta(2.0, 2.0)
    b ~ Normal(mean=0.0, variance=100.0)
    vs ~ InverseGamma(1.0, 1.0)

    a[1] ~ Normal(mean = 0.0, variance = 1.0)
    w[1] ~ Gamma(shape = 0.01, rate = 0.01)

    a[2] ~ Normal(mean = 20.0, variance = 1.0)
    w[2] ~ Gamma(shape = 0.01, rate = 0.01)

    
    for i in 1:size(y,2)
        z[i] ~ Bernoulli(s)
        ca[i] ~ NormalMixture(switch = z[i], m = a, p = w)
        for j in eachindex(x)
            y[j,i] ~ Normal(mean = ca[i] * x[j] + b, variance = vs)
        end
    end
end

With an inference:

init_mixture = @initialization begin
    μ(b) = NormalMeanVariance(0.0, 100.0)
    
    q(a) = [vague(NormalMeanVariance), vague(NormalMeanVariance)]
    q(w) = [vague(GammaShapeRate), vague(GammaShapeRate)]
    q(s) = vague(Beta)
    q(vs) = vague(InverseGamma)
    q(z) = vague(Bernoulli)
end

results_mixture = infer(
    model           = mixture_regression(), 
    data            = (y = hcat(y_pop...), x=x_pop[1]), 
    initialization  = init_mixture, 
    returnvars      = KeepLast(),
    options         = (limit_stack_depth = 100,),
    iterations      = 250,
    constraints     = MeanField(),
    free_energy     = false
)

However when I check my posteriors of the selection variable $z$, it seems to assign every measurement to the same class. Any idea if something may be wrong there?

results_mixture.posteriors[:z]

Gives:

20-element Vector{Distributions.Categorical{Float64, Vector{Float64}}}:
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])
 Distributions.Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.999999999999929, 7.097988834875855e-14])

[bvdmitri]

Hey! Thanks for using RxInfer, an interesting model you have, I checked the original comment from you before editing, to me it seems like you wanted to do smth, like

@model function mixture_regression(x, y)

    s ~ Beta(1.0, 1.0)

    a[1] ~ Normal(mean = 0.0, variance = 1e3)
    w[1] ~ Gamma(shape = 0.01, rate = 0.01)

    a[2] ~ Normal(mean = 2.0, variance = 1e3)
    w[2] ~ Gamma(shape = 0.01, rate = 0.01)

    b ~ Normal(mean=0.0, variance=100.0)

    local m

    for i in eachindex(x)
        z[i] ~ Bernoulli(s)

        # Hand-written equivalent of 
        # the map operation, with `m` defined above
        for k in 1:length(a)
            m[i, k] := a[k] * x[i] + b
        end
        
        y[i] ~ NormalMixture(switch = z[i], m = m[i, :], p = w)
    end
    
end

GraphPPL.jl language isn't smart enough to make this transformation, so I needed to write the map operation explicitly, I also tested on some dummy data

x = randn(10)
y = randn(10);

with the following initialization

init_mixture = @initialization begin
    μ(b) = NormalMeanVariance(0.0, 100.0)
    q(a) = [vague(NormalMeanVariance), vague(NormalMeanVariance)]
    q(w) = [vague(GammaShapeRate), vague(GammaShapeRate)]
    q(s) = vague(Beta)
    q(z) = vague(Bernoulli)
end

and got these results

results_mixture = infer(
    model           = mixture_regression(), 
    data            = (y = y, x = x), 
    initialization  = init_mixture, 
    returnvars      = KeepLast(),
    options         = (limit_stack_depth = 100,),
    iterations      = 250,
    constraints     = MeanField(),
    free_energy     = false
)

10-element Vector{Categorical{Float64, Vector{Float64}}}:
 Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.000343460670220945, 0.999656539329779])
 Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.28039086281305386, 0.7196091371869461])
 Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.1393305283219813, 0.8606694716780188])
 Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[3.1157790403287214e-13, 0.9999999999996884])
 Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.15862856598765293, 0.841371434012347])
 Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.9999999586423757, 4.135762428551097e-8])
 Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[1.0619854384719683e-7, 0.9999998938014562])
 Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.2745391906103953, 0.7254608093896048])
 Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[2.9299715689295617e-5, 0.9999707002843107])
 Categorical{Float64, Vector{Float64}}(support=Base.OneTo(2), p=[0.2491203914041699, 0.75087960859583])

EDIT: During more experiments I noticed that the inference actually sometimes indeed returns all ones & zeros... similar to what you observed. If I run it several times depending on x=rand(10) it sometimes clamps to one category and sometimes converges to a less extreme answers. Perhaps it is data related or priors, could you test the model that I used for your real dataset? I guess better priors (and initialization) on a and b parameters might be helpful here.

[max-de-rooij]

Thank @bvdmitri for the reply. Indeed this resembles my goal here and this gets rid of the double priors on the regression coefficient above. Instead of $y$ and $x$ being vectors, they are matrices, however, but I got the model adjusted based on your suggestion:

@model function mixture_regression(x, y)

    s ~ Beta(5.0, 5.0)
    b ~ Normal(mean=20.0, variance=10.0)

    a[1] ~ Normal(mean = 0.0, variance = 1.0)
    w[1] ~ Gamma(shape = 0.1, rate = 0.1)

    a[2] ~ Normal(mean = 10.0, variance = 1.0)
    w[2] ~ Gamma(shape = 0.1, rate = 0.1)

    local m
    for i in 1:size(y,2)
        z[i] ~ Bernoulli(s)
        for j in 1:size(x, 1)
            for k in eachindex(a)
                m[i,j,k] := a[k] * x[j,i] + b
            end
            y[j,i] ~ NormalMixture(switch = z[i], m = m[i,j,:], p = w)
        end
    end
end

I managed to get a nice example with some simulated data:

function generate_data(a, b, v, nr_samples; rng=StableRNG(1234))
    x = float.(collect(1:nr_samples))
    y = a .* x .+ b .+ randn(rng, nr_samples) .* sqrt(v)
    return x, y
end;

function generate_univariate_data(nr_samples; rng = StableRNG(1234))

    # data generating parameters
    class        = [1/2, 1/2]
    mean1, mean2 = 0.5, 5.5
    variance    = 0.8

    # generate data
    z = rand(rng, RxInfer.Categorical(class), nr_samples)
    y = zeros(nr_samples)
    for k in 1:nr_samples
        y[k] = max(0, rand(rng, NormalMeanVariance(z[k] == 1 ? mean1 : mean2, variance)))
    end

    return y, z

end;

rng = StableRNG(270523)
a_univariate, classes = generate_univariate_data(20, rng=rng)

x_pop = []
y_pop = []

for a in a_univariate
    x_data, y_data = generate_data(a, 25.0, 50.0, 8, rng=rng)

    push!(x_pop, x_data)
    push!(y_pop, y_data)
end

@model function mixture_regression(x, y)

    s ~ Beta(5.0, 5.0)
    b ~ Normal(mean=20.0, variance=10.0)

    a[1] ~ Normal(mean = 0.0, variance = 10.0)
    w[1] ~ Gamma(shape = 0.1, rate = 0.1)

    a[2] ~ Normal(mean = 2.5, variance = 10.0)
    w[2] ~ Gamma(shape = 0.1, rate = 0.1)

    local m
    for i in 1:size(y,2)
        z[i] ~ Bernoulli(s)
        for j in 1:size(x, 1)
            for k in eachindex(a)
                m[i,j,k] := a[k] * x[j,i] + b
            end
            y[j,i] ~ NormalMixture(switch = z[i], m = m[i,j,:], p = w)
        end
    end
end

init_mixture = @initialization begin
    μ(b) = NormalMeanVariance(20.0, 100.0)
    q(a) = [vague(NormalMeanVariance), vague(NormalMeanVariance)]
    q(w) = [vague(GammaShapeRate), vague(GammaShapeRate)]
    q(s) = vague(Beta)
    q(z) = vague(Bernoulli)
end

results_mixture = infer(
    model           = mixture_regression(), 
    data            = (y = hcat(y_pop...), x=hcat(x_pop...)), 
    initialization  = init_mixture, 
    returnvars      = KeepLast(),
    options         = (limit_stack_depth = 100,),
    iterations      = 30,
    constraints     = MeanField(),
    free_energy     = true
)

Resulting in a nice classification based on regression slope:

Next step will be to transfer this to a generalised linear model to do latent class trajectory analysis (see https://diabetesjournals.org/care/article/41/8/1740/36359/Pathophysiological-Characteristics-Underlying for an example of that technique).

Glad to see this seems to be working at least and keep up the good work on this package! I hope this'll be a real game changer for machine learning applications.

[bvdmitri]

Cool! Don't hesitate to reach to us if you have any troubles running the generalized model!