Skip to content

Implement Belief propagation as a new inference backend #97

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 15 commits into from
May 12, 2025
Merged

Conversation

GiggleLiu
Copy link
Member

@GiggleLiu GiggleLiu commented Apr 28, 2025

New features

This PR implements a simple Belief propagation algorithm for computing the marginal probabilities. This method is exact for tree-like tensor network. It can be used to gauge tensor networks (for data compression), even if it is not exact. The following APIs are added:

  • BeliefPropgation: create a belief propagation model instance, which can be derived from the UAI model.
  • belief_propagate: perform message passing until convergence, returns a BPState object and extra info.
  • random_tensor_train_uai and random_matrix_product_uai: two UAI models for testing.

Example

using TensorInference, Test
using OMEinsum

@testset "belief propagation" begin
    n = 5
    chi = 3
    mps_uai = TensorInference.random_tensor_train_uai(Float64, n, chi)
    bp = BeliefPropgation(mps_uai)
    @test TensorInference.initial_state(bp) isa TensorInference.BPState
    state, info = belief_propagate(bp)
    @test info.converged
    @test info.iterations < 10
    contraction_res = TensorInference.contraction_results(state)
    tnet = TensorNetworkModel(mps_uai)
    expected_result = probability(tnet)[]
    @test all(r -> isapprox(r, expected_result), contraction_res)
    mars = marginals(state)
    mars_tnet = marginals(tnet)
    for v in 1:TensorInference.num_variables(bp)
        @test mars[[v]]  mars_tnet[[v]]
    end
end

Some API changes:

  • Changed the mars keyword of TensorNetworkModel argument to unity_tensors_labels to be more accurate.
  • Changed the vars field of TensorNetworkModel to nvars. It seems not nessesary to name the variables.
  • Removed a lot of unused interfaces for TensorNetworkModel.
  • Move the autodiff related feastures to OMEinsum.

@GiggleLiu GiggleLiu changed the title [WIP] Implement Belief propagation as a new inference backend Implement Belief propagation as a new inference backend Apr 29, 2025
Copy link

codecov bot commented Apr 29, 2025

Codecov Report

Attention: Patch coverage is 97.45763% with 3 lines in your changes missing coverage. Please review.

Project coverage is 85.15%. Comparing base (c6cd869) to head (2c41d54).
Report is 22 commits behind head on main.

Files with missing lines Patch % Lines
src/Core.jl 85.71% 1 Missing ⚠️
src/belief.jl 98.63% 1 Missing ⚠️
src/mar.jl 75.00% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main      #97      +/-   ##
==========================================
+ Coverage   84.81%   85.15%   +0.33%     
==========================================
  Files          10       10              
  Lines         540      687     +147     
==========================================
+ Hits          458      585     +127     
- Misses         82      102      +20     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@GiggleLiu GiggleLiu requested review from mroavi and ArrogantGao April 29, 2025 02:46
@GiggleLiu
Copy link
Member Author

This PR is ready for review. @mroavi , @ArrogantGao

@ArrogantGao
Copy link
Contributor

ArrogantGao commented Apr 29, 2025

Can this code be used to evaluate the joint distribution of multi-variables?
I also write something similar here: https://github.com/ArrogantGao/GenericMessagePassing.jl/blob/main/src/mp.jl, its results are also tested on a few uai problems.

@GiggleLiu
Copy link
Member Author

Can this code be used to evaluate the joint distribution of multi-variables? I also write something similar here: https://github.com/ArrogantGao/GenericMessagePassing.jl/blob/main/src/mp.jl, its results are also tested on a few uai problems.

No. Is that doable in BP?

@ArrogantGao
Copy link
Contributor

Can this code be used to evaluate the joint distribution of multi-variables? I also write something similar here: https://github.com/ArrogantGao/GenericMessagePassing.jl/blob/main/src/mp.jl, its results are also tested on a few uai problems.

No. Is that doable in BP?

I don't think that is possible. Then marginals of single varibles can be calculated by pure forward process, why the function cost_and_gradient is used?

@GiggleLiu
Copy link
Member Author

Can this code be used to evaluate the joint distribution of multi-variables? I also write something similar here: https://github.com/ArrogantGao/GenericMessagePassing.jl/blob/main/src/mp.jl, its results are also tested on a few uai problems.

No. Is that doable in BP?

I don't think that is possible. Then marginals of single varibles can be calculated by pure forward process, why the function cost_and_gradient is used?

Hard to explain, maybe we can schedule a discussion. In short, I use this trick to compute the "message" sent from tensor to multiple hyperedges simultaneouly.

@mroavi
Copy link
Collaborator

mroavi commented Apr 29, 2025

I could not run the test suit in Julia 1.11.5. I get the following error:

(TensorInference) pkg> resolve
ERROR: Unsatisfiable requirements detected for package Statistics [10745b16]:
 Statistics [10745b16] log:
 ├─possible versions are: 1.11.0 - 1.11.1 or uninstalled
 └─restricted to versions 1.10.0 by an explicit requirement — no versions left

It looks like the new dependency ProblemReductions.jl only allows Julia 1.10. Because I am using Julia 1.11.5, the environment cannot be resolved, and I cannot run the tests.

Could you please check if the [compat] section of ProblemReductions.jl can be updated to also allow Julia 1.11?

Also, I noticed that our Project.toml file in TensorInference.jl has julia = "1.9". Should we update it as well to allow Julia 1.10 and 1.11?

@ArrogantGao
Copy link
Contributor

I could not run the test suit in Julia 1.11.5. I get the following error:

(TensorInference) pkg> resolve
ERROR: Unsatisfiable requirements detected for package Statistics [10745b16]:
 Statistics [10745b16] log:
 ├─possible versions are: 1.11.0 - 1.11.1 or uninstalled
 └─restricted to versions 1.10.0 by an explicit requirement — no versions left

It looks like the new dependency ProblemReductions.jl only allows Julia 1.10. Because I am using Julia 1.11.5, the environment cannot be resolved, and I cannot run the tests.

Could you please check if the [compat] section of ProblemReductions.jl can be updated to also allow Julia 1.11?

Also, I noticed that our Project.toml file in TensorInference.jl has julia = "1.9". Should we update it as well to allow Julia 1.10 and 1.11?

That is strange, I am also using julia 1.11.5 and I confirm that the tests are runable.

@GiggleLiu
Copy link
Member Author

Updates:

  1. Normalize the message, and add a damping factor (as pointed out by @ArrogantGao ). Add a test for the loopy network.
  2. Add the UAI dataset test. Got the same precision at the one in the one in Xuanzhao's implementation.
  3. Bumped the Julia version to 1.10 or larger.

@mroavi
Copy link
Collaborator

mroavi commented May 1, 2025

I managed to run the test suite. All tests passed.

@GiggleLiu I have a question: In a graph/network with loops (and relatively high complexity), Belief Propagation should be faster at calculating the marginals of all variables compared to an exact approach. I was wondering if you had a chance to check or observe this?

@GiggleLiu
Copy link
Member Author

I managed to run the test suite. All tests passed.

@GiggleLiu I have a question: In a graph/network with loops (and relatively high complexity), Belief Propagation should be faster at calculating the marginals of all variables compared to an exact approach. I was wondering if you had a chance to check or observe this?

I have not done a serious benchmark. But it should be true since its computational cost is only number of tensors * number of iterations to converge. There is no exponential barrier caused by the tree width. I observe the following fact:

  • In the loop test (cycle graph), the longer the loop, the more accurate the result it, with a circle size >10, the result is also accurate to machine precision.

A typical number of iterations to converge is 10^2 in the tests. However, @ArrogantGao told me that in some corner cases, it may fail to converge. BTW, @ArrogantGao, I borrowed the default damping factor from your code, do you have any insight about this number?

@ArrogantGao
Copy link
Contributor

ArrogantGao commented May 2, 2025

I managed to run the test suite. All tests passed.
@GiggleLiu I have a question: In a graph/network with loops (and relatively high complexity), Belief Propagation should be faster at calculating the marginals of all variables compared to an exact approach. I was wondering if you had a chance to check or observe this?

I have not done a serious benchmark. But it should be true since its computational cost is only number of tensors * number of iterations to converge. There is no exponential barrier caused by the tree width. I observe the following fact:

  • In the loop test (cycle graph), the longer the loop, the more accurate the result it, with a circle size >10, the result is also accurate to machine precision.

A typical number of iterations to converge is 10^2 in the tests. However, @ArrogantGao told me that in some corner cases, it may fail to converge. BTW, @ArrogantGao, I borrowed the default damping factor from your code, do you have any insight about this number?

I borrowed the dumping factor from https://github.com/stecrotti/BeliefPropagation.jl.

An example of the corner case is the Antiferromagnetic spin-glass with uniform J and h on regular graphs, BP fail to converge or converge to wrong results in that case.

@ArrogantGao
Copy link
Contributor

I managed to run the test suite. All tests passed.

@GiggleLiu I have a question: In a graph/network with loops (and relatively high complexity), Belief Propagation should be faster at calculating the marginals of all variables compared to an exact approach. I was wondering if you had a chance to check or observe this?

The run time of BP is quite complex, in general if it converge its cost is linear to the problem size.

@GiggleLiu GiggleLiu merged commit 3290f9e into main May 12, 2025
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants