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!