Approximate Message Passing on General Factor Graphs using Shallow Neural Networks

Factor graphs offer an efficient framework for probabilistic inference through message passing, with the added benefit of uncertainty quantification, which is crucial in safety-critical applications. However, their applicability is limited by the need to analytically solve update equations for factors, which are problem-specific and may involve intractable integrals. We propose to approximate the message update equations of individual factors with shallow neural networks, which we train on data generated by sampling from the respective factor equations, to capture complex factor relationships while maintaining computational tractability.