TensorInference: A Julia package for tensor-based probabilistic inference

A major challenge in developing intelligent systems is the ability to reason under uncertainty, a challenge that appears in many real-world problems across various domains, including artificial intelligence, medical diagnosis, computer vision, computational biology, and natural language processing. Reasoning under uncertainty involves calculating the probabilities of relevant variables while taking into account any information that is acquired. This process, which can be thought of as drawing global insights from local observations, is known as probabilistic inference.

Probabilistic graphical models (PGMs) provide a unified framework to perform probabilistic inference.These models use graphs to represent the joint probability distribution of complex systems in a concise manner by exploiting the conditional independence between variables in the model.Additionally, they form the foundation for various algorithms that enable efficient probabilistic inference.
However, even with the representational aid of PGMs, performing probabilistic inference remains an intractable endeavor on many real-world models.The reason is that performing probabilistic inference involves complex combinatorial optimization problems in very high dimensional spaces.To tackle these challenges, more efficient and scalable inference algorithms are needed.
As an attempt to tackle the aforementioned challenges, we present TensorInference.jl,a Julia package for probabilistic inference that combines the representational capabilities of PGMs with the computational power of tensor networks.By harnessing the best of both worlds, TensorInference.jlaims to enhance the performance of probabilistic inference, thereby expanding the tractability spectrum of exact inference for more complex, real-world models.

Summary
Probabilistic inference entails the process of drawing conclusions from observed data through the axioms of probability theory.Inference algorithms fall into two broad categories: exact and approximate methods.The main challenge in applying exact inference to real-world problems is its NP-hard computational complexity tied to the model's treewidth, a metric of network connectivity.This has prompted a research shift to approximate methods like Markov chain Monte Carlo and variational inference.Prominent examples of packages that implement such algorithms include Stan (Carpenter et al., 2017), PyMC3 (Oriol et al., 2023), Turing.jl(Ge et al., 2018), and RxInfer.jl(Bagaev et al., 2023).However, while these methods offer superior scalability, they do not provide formal guarantees of accuracy -a challenge that is, in itself, NP-hard to address.Consequently, exact inference methods are gaining renewed interest for their promise of higher accuracy.probabilistic inference in discrete graphical models.Capitalizing on the recent advances in the field of tensor networks (Orús, 2014(Orús, , 2019;;Robeva & Seigal, 2019), TensorInference.jloffers high-performance solutions for prevalent inference problems.Specifically, it provides methods to: 1. calculate the partition function (also known as the probability of evidence).2. compute the marginal probability distribution over each variable given evidence.3. find the most likely assignment to all variables given evidence.4. find the most likely assignment to a set of query variables after marginalizing out the remaining variables.5. draw samples from the posterior distribution given evidence (Cheng et al., 2019;Han et al., 2018).
A tensor is a mathematical object that generalizes scalars, vectors, and matrices to higher dimensions.In essence, it is a multi-dimensional array of numbers, often used for representing complex data structures in physics, engineering, computer science, and data analytics.A tensor network consists of a set of tensors in which some or all indices are contracted according to a specific pattern (Jutho et al., 2023).The term contraction refers to the summation over all the possible values along one or more dimensions of a set of tensors.These networks excel at capturing the correlations of different states in complex systems.
The order in which tensor indices are contracted plays a pivotal role in computational efficiency.Different contraction sequences can produce the same mathematical outcome, but the computational costs can vary by orders of magnitude.Since tensor network methods frequently involve multiple contractions, optimizing the contraction order becomes crucial.
The use of a tensor network-based infrastructure (Jutho et al., 2023) offers several advantages when dealing with complex computational tasks.Firstly, it simplifies the process of computing gradients by employing differentiable programming (Liao et al., 2019), a critical operation for the aforementioned inference tasks.Secondly, it supports generic element types without a significant compromise on performance.This feature enables the solution of a variety of problems using the same tensor network contraction algorithm, simply by varying the element types used.This allowed us to seamlessly implement solutions for several of the inference tasks described above (Liu et al., 2021(Liu et al., , 2022)).Thirdly, it allows users to define a hyperoptimized contraction order, which is known to have a significant impact on the computational performance of contracting tensor networks (Gao et al., 2021;Markov & Shi, 2008;Pan & Zhang, 2022).TensorInference.jlprovides a predefined set of state-of-the-art contraction ordering methods, each identified by a specific name for ease of reference.These methods include a local search-based method, denoted as TreeSA (Kalachev et al., 2022), two methods based on min-cut algorithms, denoted as SABipartite and KaHyParBipartite (Gray & Kourtis, 2021); as well as a greedy algorithm, denoted as GreedyMethod.Lastly, TensorInference.jlleverages the cutting-edge developments commonly found in tensor network libraries, including a highly optimized set of BLAS routines (Blackford et al., 2002) and GPU technology.
TensorInference.jl succeeds JunctionTrees.jl (Roa-Villescas et al., 2022, 2023), a Julia package implementing the Junction Tree Algorithm (JTA) (Jensen et al., 1990;Lauritzen & Spiegelhalter, 1988).While the latter employs tensor-based technology to optimize the computation of individual sum-product messages within the JTA context, TensorInference.jltakes a different route.It adopts a holistic tensor network approach, which opens new doors for optimization opportunities and significantly reduces the algorithm's complexity compared to the JTA.Other prominent examples of exact inference packages for probabilistic inference include libDAI (Mooij, 2010), Merlin (Marinescu, 2022), and toulbar2 (Hurley et al., 2016).For a performance comparison of TensorInference.jlagainst these alternatives, please see the Performance evaluation section in the documentation of TensorInference.jl.

Usage example
The graph below corresponds to the ASIA network (Lauritzen & Spiegelhalter, 1988), a simple Bayesian network (Pearl, 1985) used extensively in educational settings.It describes the probabilistic relationships between different random variables which correspond to possible diseases, symptoms, risk factors and test results.The ASIA network: a simplified example of a Bayesian network from the context of medical diagnosis (Lauritzen & Spiegelhalter, 1988).
In the example, a patient has recently visited Asia and is now experiencing dyspnea.These conditions serve as the evidence for the observed variables ( and ).The doctor's task is to assess the likelihood of various diseases -tuberculosis, lung cancer, and bronchitis -which constitute the query variables in this scenario (, , and ).
We now demonstrate how to use TensorInference.jlfor conducting a variety of inference tasks on this toy example.Please note that as the API may evolve, we recommend checking the examples directory of the official TensorInference.jlrepository for the most up-to-date version of this example.
# Import the TensorInference package, which provides the functionality needed # for working with tensor networks and probabilistic graphical models. In Figure 1:The ASIA network: a simplified example of a Bayesian network from the context of medical diagnosis(Lauritzen & Spiegelhalter, 1988).