Tree-AMP: Compositional Inference with Tree Approximate Message Passing

TL;DR

Tree-AMP is a Python package that enables compositional inference in high-dimensional tree-structured models, unifying various approximate message passing algorithms with automated inference, state evolution, and entropy estimation.

Contribution

It introduces a modular, automated framework for compositional inference in high-dimensional tree-structured models, integrating multiple algorithms and theoretical predictions.

Findings

The package accurately predicts asymptotic performance via state evolution.

It estimates measurement entropy using free entropy formalism.

The modular design simplifies complex inference tasks.

Abstract

We introduce Tree-AMP, standing for Tree Approximate Message Passing, a python package for compositional inference in high-dimensional tree-structured models. The package provides a unifying framework to study several approximate message passing algorithms previously derived for a variety of machine learning tasks such as generalized linear models, inference in multi-layer networks, matrix factorization, and reconstruction using non-separable penalties. For some models, the asymptotic performance of the algorithm can be theoretically predicted by the state evolution, and the measurements entropy estimated by the free entropy formalism. The implementation is modular by design: each module, which implements a factor, can be composed at will with other modules to solve complex inference tasks. The user only needs to declare the factor graph of the model: the inference algorithm, state evolution and entropy estimation are fully automated.