LLM-Assisted Rule Derivation for RxInfer.jl

[albertpod]

RxInfer.jl has a performance bottleneck that's been bugging us for a while:
When we have analytical update rules → Lightning fast inference
When we don't have those rules → We fall back to slower methods
Specifically, for models without pre-derived update rules, we currently rely on projections method.

The fallback work, but it comes with significant speed penalties that limit scalability. Essentially, we're leaving performance on the table for entire classes of probabilistic models simply because the math hasn't been worked out yet.

A Wild Idea
Inspired by AlphaEvolve paper, and recent discussions with @murphyk we can attempt using LLM agents for mathematical derivation. The key: we have verification methods to ensure correctness, so it's not just "trust the AI."

The Vision

1. LLM generates candidate update rules
2. Our system verifies mathematical correctness
3. RxInfer gains fast inference for new model classes

Timeline & Collaboration
Starting mid/end-September. Looking for collaborators interested in:

1. Mathematical AI and LLMs
2. Probabilistic programming
3. Julia development
4. Pushing boundaries in automated math

Drop a comment with your background! Let's make probabilistic programming faster for everyone.

[ofSingularMind]

Super cool :)

[bvdmitri]

I like the idea, but I would like to see how The Vision # 2 can be implemented in practice. I actually want this to verify our existing rules, not only for the new ones. If we had this verifier we can practically remove all manual @test_rule and just automatically test all rules we have

[albertpod]

That's why there's a discussion.
I think two things are important here:

1. The update rule should be a valid Julia code (this is rather easy to check I'd say).
2. The result of the update corresponds to the result of MCMC (either custom or even via Turing.jl).

On 2, the way I see it (helicopter view), as a first step, the user defines a functional form of the node with interfaces and placeholders for rules and (perhaps in the first iteration) we could automatically generate a corresponding Turing.jl models (or ask user to provide ti, or maybe we can have a generic placeholder for MCMC (not sure here)) for that node.
Then the verification process becomes:

1. Run LLM generated message passing with the analytical rule on some test messages
2. Run MCMC sampling on the equivalent Turing models
3. Compare the resulting posterior distributions.

graph LR
    A[Node/Model Specification] --> B[LLM generated Message Passing Result]
    A --> C[MCMC Result]
    
    B --> D{Results Match?}
    C --> D
    
    D -->|Yes| E[✅ Verified]
    D -->|No| F[❌ Invalid]
    
    style E fill:
    style F fill:

Loading

[murphyk]

Can you compute ground truth update rule using numerical integration instead of MCMC? Then you can create a labeled dataset of the form (x=messages in, y=message out). Now you just have to solve a symbolic regression problem, ie. find f such that y=f(x). For this, you can use https://github.com/MilesCranmer/SymbolicRegression.jl. (One caveat is that it might be restricted to univariate outputs.) They have also combined SR with LLMs in this paper: A. Grayeli, A. Sehgal, O. Costilla-Reyes, M. Cranmer, and S. Chaudhuri, “Symbolic regression with a learned concept library,” in *NIPS*, Sep. 2024. Available: https://arxiv.org/abs/2409.09359

…

On Thu, Aug 28, 2025 at 3:44 AM Albert ***@***.***> wrote: That's why there's a discussion. I think two things are important here: 1. The update rule should be a valid Julia code (this is rather easy to check I'd say). 2. The result of the update corresponds to the result of MCMC (either custom or even via Turing.jl). On 2, the way I see it (helicopter view), in the first iteration, the user defines a functional form of the node with interfaces and placeholders for rules and (perhaps in the first iteration) we could automatically generate a corresponding Turing.jl models (or ask user to provide ti?) for that node. Then the verification process becomes: 1. Run LLM generated message passing with the analytical rule on some test messages 2. Run MCMC sampling on the equivalent Turing models 3. Compare the resulting posterior distributions. graph LR A[Node/Model Specification] --> B[LLM generated Message Passing Result] A --> C[MCMC Result] B --> D{Results Match?} C --> D D -->|Yes| E[✅ Verified] D -->|No| F[❌ Invalid] style E fill: style F fill: Loading — Reply to this email directly, view it on GitHub <#501 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABDK6EAE5BDVWF7GVNR2P4L3P3MADAVCNFSM6AAAAACFA23XCWVHI2DSMVQWIX3LMV43URDJONRXK43TNFXW4Q3PNVWWK3TUHMYTIMRUGQYDKNQ> . You are receiving this because you were mentioned.Message ID: ***@***.*** com>

[bvdmitri]

Can you compute ground truth update rule using numerical integration instead of MCMC?

For univariate distributions and simple nodes with one input and one output its probably a good idea to use numerical integration, but I'm afraid that it will not scale beyond univariate distributions and/or for more complex nodes. Numerical integration for multivariate Gaussians of size 10 is already a big problem as the complexity grows exponentially. Nodes also can have multiple inputs of different types/dimensions, would be smth like (x1 = messages in, x2 = messages in, ... xn = messages in, y = messages out). It will be similar to amortized inference where one can train a neural network on inputs and outputs.

@ismailsenoz wdyt

[Nimrais]

Instead of comparing message passing results to MCMC samples from the generative model, we should directly verify that the proposed rule satisfies the stationarity condition of the constrained Bethe free energy. This is a more principled approach that respects the variational constraints inherent in message passing. Because I do not know how to change MCMC behaviour by imposing constraints. Is it possible in anyway?

The correct way it seems to follow the fact that results of message passing should be stationary point of the Lagrangian, so the message passing in the end should give such marginals that $\int \frac{\delta L}{\delta q}(s) \phi(s) ds = 0$ for any test function $\phi$. And that is what we need to check, not the inference in a respective Turing model.

Probably selection of these test functions is also our leverage if we can select them smartly we can avoid the integration at all — think of bump functions: they will effectivly reduce these integrals to evaluation at points. However the points where we are centring bump functions probably should be sampled somehow.

[bvdmitri]

I think @Nimrais is right here. It would be interesting to see if we can apply symbolic methods as suggested by @murphyk there

[ismailsenoz]

In ideal scenario what we want is to verify that LLM generated rule to compute the following sum product message update (with abuse of notation)

$$\mu(x_j) = \sum_{\substack{x_k \\ k \in \mathcal{I}}} \int f(x_1, \dots, x_N) \prod_{\substack{i \neq j \\ i \in \mathcal{J}}} \mu (x_i) \mathrm{d}x_i \,,$$

where $\mathcal{I}$ denotes the indices for discrete random variables, $\mathcal{J}$ denotes the indices for continuous random variables, $f$ is the node function and $\mu$ are messages that are possibly unnormalized (normalization might not even be possible for some messages, i.e. sigmoids) is correct in some sense. Since the case for considering messages that are not normalizable is very challenging, we should limit ourselves to consider messages that are normalizable. This means verifier can work with incoming messages that have the following property

$$\int \mu(x_j) \mathrm{d}x_j < \infty\,.$$

In this restricted scenario, I do not see any other reliable/robust method other than MCMC or sparse-grid (Smolyak) numerical integrators. I think of the following recipe for the verifier:

1. Sum out all the discrete random variables and obtain a joint function after summing out all the discrete random variables, i.e.

$$\tilde{f}(\mathcal{J}) \triangleq \prod_{i \in \mathcal{J}}\mu(x_i) \sum_{x_k} f(x_1, \dots, x_N) \prod_{k \in \mathcal{I}}\mu(x_k)$$

(Yes I know this is going to be slow)
2. Use HMC to obtain samples from the target $\tilde{f}(\mathcal{J})$,i.e.
3. Use the list of samples for the random variable $x_j$ compute statistics and check that they match with the statistics of the result returned by LLM.

As for symbolic regression, we might try reducing the node function to be bivariate perhaps by defining (x = [messages in1, messages in2, ...], y = messages out) in an attempt to prevent (x1 = messages in, x2 = messages in, ... xn = messages in, y = messages out) but for that we might need more functionality than that is available in SymbolicRegression.jl . I dont know enough about SR to have an in depth discussion but could work.

Why would we even try changing the MCMC behavior?

Because I do not know how to change MCMC behaviour by imposing constraints. Is it possible in anyway?

I dont think for the message computation imposing stationarity condition is the right way. For marginal computations it certainly is but for message computations I dont think so. Or did you mean fixed point conditions @Nimrais ?

@murphyk If I understand you correctly, you propose to build a dictionary of results with sparse cubature methods and then using these dictionaries to create a symbolic regression problem? I think it is worth a shot. At least we can try it on a small scale 1d problem to test.

[murphyk]

Yes I propose you use SR on 1-2d problems where ground truth output is from integration. Then generalize to Nd by using induction, either manually or with LLM., where in latter case you provide the symbolic result for 1d as a hint , ie part of the prompt. Sent from Gmail Mobile

…

On Tue, Sep 2, 2025 at 2:35 AM Ismail Senoz ***@***.***> wrote: In ideal scenario what we want is to verify that LLM generated rule to compute the following sum product message update (with abuse of notation) $$\mu(x_j) = \sum_{\substack{x_k \\ k \in \mathcal{I}}} \int f(x_1, \dots, x_N) \prod_{\substack{i \neq j \\ i \in \mathcal{J}}} \mu (x_i) \mathrm{d}x_i \,,$$ where $\mathcal{I}$ denotes the indices for discrete random variables, $\mathcal{J}$ denotes the indices for continuous random variables, $f$ is the node function and $\mu$ are messages that are possibly unnormalized (normalization might not even be possible for some messages, i.e. sigmoids) is correct in some sense. Since the case for considering messages that are not normalizable is very challenging, we should limit ourselves to consider messages that are normalizable. This means verifier can work with incoming messages that have the following property $$\int \mu(x_j) \mathrm{d}x_j < \infty\,.$$ In this restricted scenario, I do not see any other reliable/robust method other than MCMC or sparse-grid (Smolyak) numerical integrators. I think of the following recipe for the verifier: 1. Sum out all the discrete random variables and obtain a joint function after summing out all the discrete random variables, i.e. $$\tilde{f}(\mathcal{J}) \triangleq \prod_{i \in \mathcal{J}}\mu(x_i) \sum_{x_k} f(x_1, \dots, x_N) \prod_{k \in \mathcal{I}}\mu(x_k)$$ (Yes I know this is going to be slow) 2. Use HMC to obtain samples from the target $\tilde{f}(\mathcal{J})$,i.e. 3. Use the list of samples for the random variable $x_j$ compute statistics and check that they match with the statistics of the result returned by LLM. As for symbolic regression, we might try reducing the node function to be bivariate perhaps by defining (x = [messages in1, messages in2, ...], y = messages out) in an attempt to prevent (x1 = messages in, x2 = messages in, ... xn = messages in, y = messages out) but for that we might need more functionality than that is available in SymbolicRegression.jl . I dont know enough about SR to have an in depth discussion but could work. Why would we even try changing the MCMC behavior? Because I do not know how to change MCMC behaviour by imposing constraints. Is it possible in anyway? I dont think for the message computation imposing stationarity condition is the right way. For marginal computations it certainly is but for message computations I dont think so. Or did you mean fixed point conditions @Nimrais <https://github.com/Nimrais> ? @murphyk <https://github.com/murphyk> If I understand you correctly, you propose to build a dictionary of results with sparse cubature methods and then using these dictionaries to create a symbolic regression problem? I think it is worth a shot. At least we can try it on a small scale 1d problem to test. — Reply to this email directly, view it on GitHub <#501 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABDK6EFCNIW23WYHRILT3TD3QVQF7AVCNFSM6AAAAACFA23XCWVHI2DSMVQWIX3LMV43URDJONRXK43TNFXW4Q3PNVWWK3TUHMYTIMRYGM3TMNI> . You are receiving this because you were mentioned.Message ID: ***@***.*** com>

[albertpod]

Thanks for the discussion! Just to clarify: I do care about message correctness, but I want to verify it through posterior correctness rather than direct mathematical verification.
My approach: LLM generates rules → compile in RxInfer → compare posteriors with Turing.jl ground truth → iterate with feedback until they match.
Working on a small demo in the next few days to show this in action. I welcome different verification approaches - excited to see what everyone comes up with! 🚀

[murphyk]

FWIW, tutorial on SR https://symbolicregression2025.github.io/

…

On Wed, Sep 3, 2025 at 10:45 AM Albert ***@***.***> wrote: Thanks for the discussion! Just to clarify: I do care about message correctness, but I want to verify it through posterior correctness rather than direct mathematical verification. My approach: LLM generates rules → compile in RxInfer → compare posteriors with Turing.jl ground truth → iterate with feedback until they match. Working on a small demo in the next few days to show this in action. I welcome different verification approaches - excited to see what everyone comes up with! 🚀 — Reply to this email directly, view it on GitHub <#501 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABDK6EECHNUXWFD4A5KJ3IL3Q4SKHAVCNFSM6AAAAACFA23XCWVHI2DSMVQWIX3LMV43URDJONRXK43TNFXW4Q3PNVWWK3TUHMYTIMRZHA4DQNY> . You are receiving this because you were mentioned.Message ID: ***@***.*** com>

[t3chw]

Hi @albertpod , I like the idea.
I would like to work on "Mathematical AI and LLMs" or "Probabilistic programming".
Background - Mathemathics and Computer science, (worked on bayesian, contributed to numpyro)

[bvdmitri]

Hey @t3chw , you don't have to ask the permission to work on this. Many people have contributed to RxInfer, its a community project. If you have some nice ideas for the mathematical foundations of LLMs and probabilistic programming please share!

[murphyk]

https://arxiv.org/abs/2302.06675 Vaguely related - LLM powered program search discovered muon optimizer Sent from Gmail Mobile

…

On Thu, Sep 4, 2025 at 9:49 AM Mohit ***@***.***> wrote: Hi @albertpod <https://github.com/albertpod> , I like the idea. I would like to work on Mathematical AI and LLMs or Probabilistic programming. Background - Mathemathics and Computer science, (worked on bayesian, contributed to numpyro) — Reply to this email directly, view it on GitHub <#501 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABDK6EG7C2B3VC3NTCU5XET3RA7PXAVCNFSM6AAAAACFA23XCWVHI2DSMVQWIX3LMV43URDJONRXK43TNFXW4Q3PNVWWK3TUHMYTIMZQHAYDCMQ> . You are receiving this because you were mentioned.Message ID: ***@***.*** com>

[bayesianempirimancer]

This is a great idea.

Also, I am not sure if this is helpful, but a few years back I tried to get similiar functionality by using regular old neural networks to estimate partition functions for non-named exponential family distributions (i.e. arbitrary sufficient statistics $T(x)$) by using MCMC samples. The idea was that if you just made every $\log q(x_i) + c = \left<\log p(x)\right>{q(x{\not i} ) } = T(x_i)\cdot \eta$ an exponential family distribution then messages consist of simply translating easy to compute expectations into natural parameters as one does when working with any p(x) cleverly built out of named exponential family distributions. This would effectively expand the set of distributions that are as easy to work with as named distributions by a huge amount since you are no longer stuck with a highly limited set of choices for $T(x)$ for which explicit update rules and expectation evaluations exist. You might still have to decompose certain $T_k(x_1,x_2)$ into something like $\sum_{ij} a_{ijk} \phi_i(x_1)\phi_j(x_2)$ but hardly seems like much of a penalty if $T_i(x) = \phi_i(x)$ is a sufficient statistic you can work with.

Anyway, I never got it to work well for the partition functions itself, which was problematic because I was interested in getting a good estimate of it for the purposes of doing mixture modeling. However, when I tried to just get it to produce the expectations of the sufficient statistics $\mu_T(\eta) = \left&lt; T(x) | \eta \right&gt;$ as a function of natural parameters, it did all right (with enough data relative to the dimensionality of $x$ and $T(x)$). There were two tricks that helped get decent results:

The first was giving the network as inputs a bunch of functions of $\eta$, i.e. $\eta$, $\log(|\eta|)$, $||\eta||$, $\eta/||\eta||$, $1/\eta$, ... These choices were not particularly well grounded, I was just trying to ensure that it would be easy for the network to find a good estimate of the maximum likelihood estimate of $x$ which is invariant to the magnitude of $\eta$. That was a lazy SVM inspired choice. The second trick was choosing a loss function that (soft of) respected the geometry of exponential family distributions. This was accomplished by using the KL divergence between two gaussian distribution: $N(\mu_T^{MCMC}(\eta), \Sigma_{TT}^{MCMC}(\eta))$ and $N(\mu_T^{net}(\eta),\nabla\mu_T^{net}(\eta))$. Here $\nabla$ is the gradient wrt to $\eta$ and $\Sigma_{TT}$ is the covariance of the sufficient statistic.

What I never got around to was trying to estimate the partition function, $A(\eta)$ up to an arbitrary constant using the same objective function, but with $\mu_T(\eta) = \nabla A^{net}(\eta)$. Not sure that would actually work, but might be worth a try. I will also add that i initially tried fitting $\mu_T(\eta)$ with invertible neural networks since, you know, its an invertible function. Indeed the realnvp paper was what inspired the project in the first place. Anyway, I could not get that to work very well but I was just using affine coupling layers and not the more sophisticated stuff the INN literature has come up with more recently.

Anyway, let me know if you start up a repo. I have wanted to clean up and modernize that work and would be happy to contribute it if only for the purposes of comparison to your LLM optimized approach.

[albertpod]

I think @ismailsenoz proposed the product of Laplaces:
#501 (reply in thread)

[bayesianempirimancer]

That not an exponential family distribution!!!! :) It looks like one but the mu in the abs means it has a parameters that can't be put in exponential family form..... Anyway, I will fake it with a mixture of two exponentials and see what happens.

[albertpod]

It should be under certain conditions. @ismailsenoz knows better

[ismailsenoz]

That not an exponential family distribution!!!! :) It looks like one but the mu in the abs means it has a parameters that can't be put in exponential family form..... Anyway, I will fake it with a mixture of two exponentials and see what happens.

It is in exponential family conditioned that the location parameters are fixed. That is why we have conditioner field in the ef distribution structures. Many ef distributions are exponential family when some parameters are conditioned to be fixed. Laplace distribution for example is only EF if it is location parameters are fixed!! So the product of two Laplace distributions whose location parameters are fixed is in exponential family. So the method you develop should recover the multiplicative form assuming the locations are fixed. I hope this helps

[bayesianempirimancer]

one laplace product distribution with two components separated by 2.

sampling time per natural parameter using blackjax mcmc was 90 seconds or 0.09 seconds per sample.

800 different $\eta$, $\mu_T$ pairs spanning eta values from 0.2 to 10 were used to train a simple MLP with some custom feature engineering. Total training time was 90 seconds. But more critically, evaluation of $mu_T(\eta)$ comes in at <0.3 ms. MSE is comparable to that expected by sampling.

This was a pretty easy one so I didnt bother to test the more sophisticated networks or clever solutions on it. sloppy poorly documented code can be found at: https://github.com/bayesianempirimancer/nnef-dists

[albertpod]

The PR is pending, not ideal, solo effort, so looking for feedback and collaborative improvement.

Notebook blob