Error during inference with Negative Binomial model: “Variables [p, r] have not been updated after an update event”

[mrazomej]

Dear RxInfer team,

I’m trying to implement a simple Negative Binomial model using RxInfer.jl, but I’m encountering an error when running inference. Despite reviewing the initialization section in the documentation, I’m unsure how to proceed.

Code to reproduce the issue:

using RxInfer, Random, Distributions

# Register the Negative Binomial node
@node Distributions.NegativeBinomial Stochastic [out, r, p]

# Define the model for a single gene using RxInfer
@model function simple_nb(data, γ_α, γ_β, β_α, β_β)
    # Set prior for r
    r ~ GammaShapeRate(γ_α, γ_β)
    # Set prior for p
    p ~ Beta(β_α, β_β)
    # Negative Binomial likelihood
    data .~ NegativeBinomial(r, p)
end

# Generate synthetic dataset
Random.seed!(42)

n_samples = 100

# Ground truth parameters for p
β_α = 1.0
β_β = 1.0
p_true = rand(Beta(β_α, β_β))

# Ground truth parameters for r
γ_α = 5.0
γ_β = 0.1
r_true = rand(GammaShapeRate(γ_α, γ_β))

# Simulate the data
data = rand(NegativeBinomial(r_true, p_true), n_samples)

# Run inference
result = infer(
    model = simple_nb(γ_α=γ_α, γ_β=γ_β, β_α=β_α, β_β=β_β),
    data = (data = data,),
)

Error message:

ERROR: Variables [p, r] have not been updated after an update event.
Therefore, make sure to initialize all required marginals and messages.
See `initialization` keyword argument for the inference function.
See the official documentation for detailed information regarding the initialization.

What I’ve tried:

• Reviewed the initialization section in the documentation.
• Attempted to initialize variables using the initialization keyword in the infer function but wasn’t sure how to apply it in this context.

Question:

Could you please help me understand why this error occurs and how to initialize the variables p and r for this model properly? Any guidance on correctly setting up the inference for the Negative Binomial model would be greatly appreciated.

[wouterwln]

Hi, thank you for trying out RxInfer! The problem with your model is indeed with the initialization. I'll try to explain the problem:

You have two variables r and p that are used in multiple NegativeBinomial nodes, so if we draw the factor graph according to what we do here, we see that there is a loop in the graph (from r to a NegativeBinomial node to p to a different NegativeBinomial node back to r). This is not a problem, it just means we have to initialize something :)

Now, if we look at the documentation for defining a custom node, we can what we should do. here we see the update equations for sum-product message passing, where here we see the update rules for variational message passing. There is a subtle difference between these two update equations: The sum-product message passing update rule only depends on $\mu(x)$ (the message from an adjacent node), whereas the variational message passing update rule depends on $q(x)$ (the posterior marginal, which is the product of two messages). Now which one you will have to initialize depends on the constraints that you impose on your variational posterior, as this determines whether we will use variational message passing or sum-product message passing.

Since you don't impose any additional constraints on the variational posterior, RxInfer assumes you can do sum-product message passing, which would mean that we will have to initialize $\mu(r)$ and $\mu(p)$, so let's go ahead and do that:

initialization = @initialization begin
    μ(r) = GammaShapeRate(1.0, 1.0)
    μ(p) = Beta(1.0, 1.0)
end

Now if you try to run inference with this initialization, you'll run into the next error telling you that there's no rules available for the combination of your node and the incoming messages. This is because you defined your node yourself and RxInfer does not know how to handle the newly created NegativeBinomial node. Luckily there's several ways around this. The first would be the most correct way to go about this, but it requires some work on your part. here we have a tutorial detailing how you can derive rules for custom nodes yourself. This involves taking an integral over the product of incoming messages with the pdf of the Negative Binomial distribution. We have a builtin workaround for this, however. The result will be both slower and less accurate than deriving custom rules, but to get your inference up and running it should be fine. We will go through the route we have detailed here and use ExponentialFamilyProjection to derive an approximate posterior. Futhermore, we will use this method to derive rules automatically from the pdf of the NegativeBinomial distribution. This gives you the following (adjusted) snippet:

using RxInfer, Random, Distributions, ExponentialFamilyProjection

# Register the Negative Binomial node
@node Distributions.NegativeBinomial Stochastic [out, r, p]

# Define the model for a single gene using RxInfer
@model function simple_nb(data, γ_α, γ_β, β_α, β_β)
    # Set prior for r
    r ~ GammaShapeRate(γ_α, γ_β)
    # Set prior for p
    p ~ Beta(β_α, β_β)
    # Negative Binomial likelihood
    data .~ NegativeBinomial(r, p)
end

# Generate synthetic dataset
Random.seed!(42)

n_samples = 100

# Ground truth parameters for p
β_α = 1.0
β_β = 1.0
p_true = rand(Beta(β_α, β_β))

# Ground truth parameters for r
γ_α = 5.0
γ_β = 0.1
r_true = rand(GammaShapeRate(γ_α, γ_β))

# Simulate the data
data = rand(NegativeBinomial(r_true, p_true), n_samples)

initialization = @initialization begin
    μ(r) = GammaShapeRate(1.0, 1.0)
    μ(p) = Beta(1.0, 1.0)
end

constraints = @constraints begin
    q(r) :: ProjectedTo(Gamma)
    q(p) :: ProjectedTo(Beta)
    μ(r) :: ProjectedTo(Gamma)
    μ(p) :: ProjectedTo(Beta)
end

# Run inference
result = infer(
    model = simple_nb(γ_α=γ_α, γ_β=γ_β, β_α=β_α, β_β=β_β),
    data = (data = data,),
    initialization = initialization,
    constraints = constraints,
    options = (
        rulefallback = NodeFunctionRuleFallback(),
    ),
    iterations=10
)

This gives you an inference result. You can improve the inference by deriving rules yourself. Hopefully you can get started with RxInfer like this!

EDIT: After messing around a bit more with your model, I have found a way to improve the inference result. For this, I imposed a Mean Field constraint on r and p. This means that we will try to approximate independent marginal posterior distributions and RxInfer will try to use variational message passing, but it also means that we don't have to approximate the sum-product messages from our node. This means that we have to initialize $q(r)$ and $q(p)$ instead of $\mu(r)$ and $\mu(p)$ and impose a constraint. You can realize this with the following code:

using RxInfer, Random, Distributions, ExponentialFamilyProjection

# Register the Negative Binomial node
@node Distributions.NegativeBinomial Stochastic [out, r, p]

# Define the model for a single gene using RxInfer
@model function simple_nb(data, γ_α, γ_β, β_α, β_β)
    # Set prior for r
    r ~ GammaShapeRate(γ_α, γ_β)
    # Set prior for p
    p ~ Beta(β_α, β_β)
    # Negative Binomial likelihood
    data .~ NegativeBinomial(r, p)
end

# Generate synthetic dataset
Random.seed!(42)

n_samples = 100

# Ground truth parameters for p
β_α = 1.0
β_β = 1.0
p_true = rand(Beta(β_α, β_β))

# Ground truth parameters for r
γ_α = 5.0
γ_β = 0.1
r_true = rand(GammaShapeRate(γ_α, γ_β))

# Simulate the data
data = rand(NegativeBinomial(r_true, p_true), n_samples)

initialization = @initialization begin
    q(r) = GammaShapeRate(1.0, 1.0)
    q(p) = Beta(1.0, 1.0)
end

constraints = @constraints begin
    q(r) :: ProjectedTo(Gamma)
    q(p) :: ProjectedTo(Beta)
    q(r, p) = MeanField()
end

# Run inference
result = infer(
    model = simple_nb(γ_α=γ_α, γ_β=γ_β, β_α=β_α, β_β=β_β),
    data = (data = data,),
    initialization = initialization,
    constraints = constraints,
    options = (
        rulefallback = NodeFunctionRuleFallback(),
    ),
    iterations = 10
)

The inference result I get with this code is a lot better than when we approximate the message first and the posterior marginal afterwards, so hopefully this method is a bit more useful.
Cheers!

[mrazomej]

Thank you so much for the quick response, @wouterwln. Unfortunately, I tried both suggestions, and the resulting posteriors entirely miss the true parameter. The quick approximation seems not accurate enough to address this problem.

I will try to derive the proper rules for the NegativeBinomial node, but I guess that if you haven't already implemented them, there might not be closed-form solutions for the required integrals. I really want RxInfer to work for the actual application I'm after, but I might be hitting a wall here.

If you don't mind, I'll bug you a little more if I make some progress or get stuck again.

Thanks again for your help. I am really excited about what you guys are doing, and I hope to get enough understanding to make it work for my goals.

[wouterwln]

No worries, let us know if there's anything we can help with. It's nice to hear that you are exited about RxInfer and hopefully we can make something work! I don't know if anyone in our lab has ever derived proper rules for the Negative Binomial node. If the sum-product route does not work, maybe the VMP rules are defined. In any case, we can help you with implementational details. Hope this helps!

[mrazomej]

It has been a fascinating journey to understand enough of RxInfer's mathematical background. Last time, I couldn't understand @wouterwln's suggestions, but after many hours of work, I think I know now—at least the basics of message passing and variational message passing.

After this long journey, I came back to my original problem. I was able to derive the variational message passing rule for the newly-defined negative binomial on the :p interface as

@node Distributions.NegativeBinomial Stochastic [out, r, p]

@rule NegativeBinomial(:p, Marginalisation) (
    q_out::PointMass, q_r::Any
) = begin
    # Get mean of the counts
    x = mean(q_out)

    # Get mean of q(r)
    r = mean(q_r)

    # Returns Beta distribution with parameters (1 + r, 1 + x)
    return Beta(one(r) + r, one(x) + x)
end

However, regardless of how much I tried, I think that the message on the r interface does not have a closed-form solution. The problem comes from the fact that the parameter is an argument in the $\Gamma$ function since the negative binomial PMF is of the form

$$ \text{NegBinom}(u \mid r, p) = \frac{\Gamma(u + r)}{u!\Gamma(r)} (1 - p)^u p^r. \tag{1} $$

I got excited about the idea behind ExponentialFamilyProjection.jl, especially after watching the recent JuliaCon talk. So I decided to basically try again what @wouterwln with the :p rule defined to see if that improved the inference. My model takes the very simple form

@model function negbinom_model(
    counts,
    r_prior_shape,
    r_prior_rate,
    p_prior_shape,
    p_prior_rate,
)
    # Prior for the Gamma shape parameter
    r ~ GammaShapeRate(r_prior_shape, r_prior_rate)

    # Prior for the Beta shape parameter
    p ~ Beta(p_prior_shape, p_prior_rate)

    # Loop over each count
    for i in eachindex(counts)
        counts[i] ~ NegativeBinomial(r, p)
    end # for
end # negbinom_model

I then generate synthetic data to test it as

Random.seed!(42)
# Define number of samples
n_samples = 1_000

# Define ground truth parameters for each r
r_α = 5.0
r_β = 0.1
r_true = rand(RxInfer.GammaShapeRate(r_α, r_β))

# Define prior parameters
p_α = 1.0
p_β = 1.0
p_true = rand(RxInfer.Beta(p_α, p_β))


# Simulate the data
counts = rand(NegativeBinomial(r_true, p_true), n_samples)

# Store data in dictionary
data = Dict(:counts => counts)

I then run the inference with the following initialization and constraints

# Initialize variational distributions
initialization = @initialization begin
    q(r) = vague(Gamma)
    q(p) = vague(Beta)
end

# Define constraints
constraints = @constraints begin
    q(r)::ProjectedTo(Gamma)
    q(r, p) = MeanField()
end

# Run inference
results = infer(
    model=negbinom_model(
        r_prior_shape=r_α,
        r_prior_rate=r_β,
        p_prior_shape=p_α,
        p_prior_rate=p_β,
    ),
    data=data,
    initialization=initialization,
    constraints=constraints,
    options=(
        rulefallback=NodeFunctionRuleFallback(),
    ),
)

But the inference looks terrible. It is not a subtle failure, but it completely misses the ground truth. Do you have any suggestions on whether this can be improved?

I tried changing the number of iterations when performing the projection, doing

ProjectedTo(Gamma; parameters = ProjectionParameters(niterations = 500))

or the tolerance

ProjectedTo(Gamma; parameters = ProjectionParameters(tolerance = 1e-2))

But nothing worked. Is it simply that RxInfer won't work with this distribution that does not belong to the exponential family when r is unknown?

[Nimrais]

Hi, thanks for this piece of work. It's an excellent example!

I would not say that inference is wrong; it looks like your model parameters are not identifiable from the observations.

Look, I took your code and modified it a bit:

using RxInfer, Random, Distributions, ExponentialFamilyProjection
using StatsPlots
using Manopt

# Register the Negative Binomial node
@node Distributions.NegativeBinomial Stochastic [out, r, p]

@model function negbinom_model(
    counts,
    r_prior_shape,
    r_prior_rate,
    p_prior_shape,
    p_prior_rate,
)
    # Prior for the Gamma shape parameter
    r ~ GammaShapeRate(r_prior_shape, r_prior_rate)

    # Prior for the Beta shape parameter
    p ~ Beta(p_prior_shape, p_prior_rate)

    for i in eachindex(counts)
        counts[i] ~ NegativeBinomial(r, p)
    end
end
# Generate synthetic dataset
Random.seed!(42)
# Define number of samples
n_samples = 1_000

# Define ground truth parameters for each r
r_α = 5.0
r_β = 0.1
r_true = rand(RxInfer.GammaShapeRate(r_α, r_β))

# Define prior parameters
p_α = 1.0
p_β = 1.0
p_true = rand(RxInfer.Beta(p_α, p_β))

# Simulate the data
counts = rand(NegativeBinomial(r_true, p_true), n_samples)
@show mean(counts), var(counts)

# Store data in dictionary
data = Dict(:counts => counts)

# Initialize variational distributions
initialization = @initialization begin
    q(r) = vague(Gamma)
    q(p) = vague(Beta)
end

# Define constraints
constraints = @constraints begin
    q(r)::ProjectedTo(Gamma, parameters = ExponentialFamilyProjection.ProjectionParameters(stepsize=ConstantLength(1)))
    q(p)::ProjectedTo(Beta)
    q(r, p) = MeanField()
end

# Run inference
results = infer(
    model=negbinom_model(
        r_prior_shape=r_α,
        r_prior_rate=r_β,
        p_prior_shape=p_α,
        p_prior_rate=p_β,
    ),
    data=data,
    initialization=initialization,
    constraints=constraints,
    options=(
        rulefallback=NodeFunctionRuleFallback(),
    ),
    showprogress = true,
    iterations=10,
)


mean_r = mean(results.posteriors[:r][end])
mean_p = mean(results.posteriors[:p][end])

mean_counts = rand(NegativeBinomial(mean_r, mean_p), n_samples)
@show mean(counts), var(counts)

begin
    histogram(counts, bins=40, alpha=0.5, label="True", normalize=:probability)
    histogram!(mean_counts, bins=40, alpha=0.5, label="Generated from the inferred parameters", normalize=:probability)
end

The histograms for the observations and generated out of the model:

The result that RxInfer provides is a posterior over-parameters generative model that can generate your observations. It could be not unique sometimes, which is what it looks like here. It's a shame that the prior contribution almost vanishes after the inference with many observations, but this behavior is a known problem with mean-field approximations. There is a known remedy for this problem that sometimes works (using a hierarchical prior). I don't have time this week to play with this model, but I could check it next week.

[bvdmitri]

Thanks for the detailed analysis, @mrazomej! I just wanted to add a bit more context to this discussion.

Whether the discrepancies in the inferred parameters are “large” or not depends a lot on the specific context and warrants further investigation. Here’s why:

Imagine you’re trying to infer a and b given their sum, a + b = 7. What are the “real” values for a and b? Both (a = 1, b = 6) and (a = 6, b = 1) are valid solutions. This ambiguity reflects an inherent identifiability problem that some problems naturally have. This can be resolved by assigning good priors for the parameters (both in the model specification and the initialization).

Another possible explanation is floating-point errors. For example, the NegativeBinomial distribution relies on the inverse gamma function, which approaches zero very quickly. Numerically, there’s little difference between 1/Γ(10) and 1/Γ(60)—modern computers essentially treat them as equivalent.

While I’m not entirely sure if both or any of these explanations apply here, the histogram shared by @Nimrais suggests that the inferred parameters perform well enough to recover the underlying distribution and generate the same dataset. That’s a promising result! If I would you I would also check how other methods perform e.g. NUTS or ADVI for this particular task on the exact same dataset.

[mrazomej]

Thank you very much for such a quick and insightful response!

I never thought the negative binomial had identifiability problems—quite shocking. But I agree that @Nimrais's histogram is quite revealing. I had the same thought as @bvdmitri, so I ran this very simple model with Turing.jl using both NUTS and ADVI. When sampling from the posterior with HMC, the posterior does overlap with the ground truth. However, I can see that even with MCMC, the model is not perfectly identifiable since, with 1000 data points in the dataset, the posterior still does not surround the ground truth perfectly.

But, as @Nimrais suggested, the problem seems to be related to the mean-field approximation because ADVI has a very similar issue to RxInfer, completely missing the ground truth—although to a lesser extent.

However, I can confirm the model's degeneracy, as when I generate posterior predictive checks with either approach, I am able to capture the data generative process.

I am curious to hear more about @Nimrais hierarchical prior suggestion!

Thanks again for your help. Learning about this completely different way of performing inference has been a fascinating experience. The software you guys are developing is a masterclass in open-source software.

[rzhli]

A silly question, can prior be a customed factor node? like r ~ Exponential(r_prior_β) or LogNormal(μ,σ), or Chisq(ν) ? I saw most examples use custom factor node for likelihood only