Tree-AMP: Compositional Inference with Tree Approximate Message Passing

Antoine Baker, Florent Krzakala, Benjamin Aubin, Lenka Zdeborová; 24(57):1−89, 2023.

Abstract

The python package Tree-AMP is introduced as a tool for compositional inference in high-dimensional tree-structured models. This package offers a comprehensive framework for studying various approximate message passing algorithms that have been previously developed for different machine learning tasks, including generalized linear models, inference in multi-layer networks, matrix factorization, and reconstruction using non-separable penalties. In certain models, the algorithm’s asymptotic performance can be theoretically forecasted through state evolution, while the free entropy formalism can be used to estimate the measurements entropy. The package’s design is modular, allowing users to combine different factor modules to tackle complex inference tasks. The user only needs to specify the factor graph of the model, as the inference algorithm, state evolution, and entropy estimation are fully automated.

[abs]

[pdf][bib]
      
[code]