Technical Summary: RAMP – Recognition Parametrisation by Amortised Message Passing

1. Problem Statement

Unsupervised learning aims to uncover latent factors that explain dependencies among observations. Probabilistic graphical models (PGMs) are well-suited for this task, representing uncertainty and non-deterministic links mathematically. However, a fundamental trade-off exists between expressivity and tractability.

• Tractability: Exact inference in PGMs typically requires restrictive distributional choices (e.g., linear-Gaussian) to allow for exact belief propagation.
• Expressivity: Nonlinear, complex dependencies often render the necessary integrals intractable. Existing approximation methods (e.g., variational inference, Monte Carlo) often scale poorly with model size or fail to provide a unified framework that integrates deep learning's approximation capabilities with probabilistic coherence.

The core challenge is to develop a method that learns flexible, nonlinear latent structures on complex, high-dimensional data while maintaining a principled probabilistic framework with a well-defined likelihood.

2. Methodology: RAMP

The authors propose RAMP (Recognition Parametrisation by Amortised Message Passing), a method that implicitly defines latent structure by learning a flexible, nonlinear, amortised message-passing framework. RAMP builds upon the Recognition-Parametrised Model (RPM) paradigm and adapts it to tree-structured graphical models.

2.1 Core Concepts

• Recognition-Parametrisation: Instead of parametrising the generative process $p(x|z)$, RAMP parametrises the recognition factors $f(z|x)$ that map observations to posterior beliefs. The joint distribution is constructed as:
  $$p(Z, X) = p(Z) \prod_p \frac{f_p(Z|x_p)}{F_p(Z)} p_{emp}(x_p)$$
  where $F_p(Z)$ is a normalising constant ensuring the model is properly defined.
• Amortised Message Passing: In standard belief propagation (BP) on trees, messages are computed via exact integration. RAMP replaces these integrals with learned, amortised transformations.
  • For a latent node $z_k$ and its neighbor $z_j$, the belief update is defined not by integration but by a functional $G_{j \to k}$:
    $$f^k_j(z_k | X^k_j) = G_{j \to k} \left( f^j_0(z_j) \prod_{l \in \partial j \setminus k} \frac{f^j_l(z_j | X^j_l)}{f^j_0(z_j)} \right)$$
  • This allows the network to learn the transformation of beliefs directly, bypassing the need for explicit generative factors or tractable integrals.

2.2 Model Architecture

• Tree Structure: RAMP operates on undirected, tree-structured graphical models where latent variables $Z$ and observed variables $X$ form a tree. Observations are constrained to leaf nodes.
• Nodewise RPMs: The model is trained by simultaneously optimising multiple "nodewise" RPMs, one for each latent variable $z_k$. Each nodewise model respects the conditional independence partition induced by $z_k$ on the observations.
• Joint Optimisation: The objective function is the sum of free energies (variational lower bounds) across all nodewise models:
  $$F = \frac{1}{K} \sum_k F_k$$
  This joint optimisation ensures that the amortised message-passing network captures the global tree structure and conditional independencies, rather than just local subtree dependencies.
• Exponential Family Parametrisation: To ensure tractability of the learning objective, beliefs are represented as members of an exponential family (e.g., Gaussians). The message passing is implemented as a transformation of natural parameters via neural networks ($g_{j \to k}$), where the product of distributions corresponds to the sum of natural parameters.

2.3 Theoretical Guarantees

The paper establishes that if the data is compatible with a tree-structured latent model satisfying specific conditional independence assumptions (Assumption 1), and the amortised network optimises the joint objective, then:

1. Each nodewise free energy attains its optimal value.
2. The learned beliefs $f^k_j$ are sufficient statistics for the observations, effectively recovering the true latent structure.
3. The method recovers the correct conditional independence structure among latent variables (Lemma 2).

3. Key Contributions

1. Novel Integration of Deep Learning and PGMs: RAMP offers a unified approach where neural networks are used to amortise the message-passing operator itself, rather than just approximating the posterior or the generative model.
2. Likelihood-Based Learning in Nonlinear Models: Unlike many deep latent variable models that rely on variational lower bounds with loose approximations, RAMP maintains a proper likelihood function derived from the RPM framework, even with nonlinear message passing.
3. Robustness to Misspecification: The method demonstrates robustness to structural misspecification (learning correct latents even if the assumed tree topology is slightly perturbed) and distributional misspecification (handling nonlinear data generating processes where true posteriors are non-Gaussian).
4. Bidirectional Semantics: In time-series applications, RAMP naturally implements bidirectional inference (forward-backward) with probabilistic semantics, outperforming standard contrastive methods and structured VAEs in recovering dynamical states.

4. Experimental Results

The authors evaluated RAMP on three distinct tasks:

• Hierarchical Linear/Nonlinear Trees:

  • Linear-Gaussian: RAMP recovered the exact solution with high fidelity, matching the performance of exact inference.
  • Structural Misspecification: When trained on perturbed tree structures, RAMP maintained high correlation with ground-truth latents for nodes where the local partition remained correct, demonstrating that learning signals are local to the conditional independence structure.
  • Nonlinear Data: On data generated by nonlinear transformations, RAMP outperformed Black-Box Variational Inference (BBVI) and supervised Gaussian approximations, successfully learning representations of the latent variables despite the true posterior being non-Gaussian.

• Nonlinear Dynamical Systems (Pendulum):

  • Trained on images of a swinging pendulum, RAMP successfully inferred the 2D phase space (angle and angular velocity) within a 2D latent space.
  • It outperformed Contrastive Predictive Coding (CPC) and Structured VAEs (SVAE), which failed to reliably recover angular velocity (a state variable not directly observable in a single frame) when constrained to the true dimensionality.

• Pose Estimation:

  • Applied to the Leeds Sports Pose (LSP) dataset, RAMP inferred body joint configurations from local image patches without global location information.
  • The inferred latent variables encoded the orientation of limbs connecting joints. The model successfully reconstructed pose sketches and handled occluded joints (e.g., a missing ankle) by leveraging the learned tree-structured prior.

5. Significance and Claims

The paper claims that RAMP represents a significant step toward merging adaptive neural network components into a coherent probabilistic framework. Its significance lies in:

• Identifiability: By leveraging the conditional independence structure of tree graphs, RAMP ensures that latent variables are discoverable based on the specific partitions of observations they induce.
• Coherence: Unlike generic deep learning approaches where internal representations lack probabilistic interpretation, RAMP's inference procedure directly defines the model, ensuring that the learned representations have well-defined likelihoods and probabilistic semantics.
• Scalability and Flexibility: By amortising message passing, RAMP scales to complex, high-dimensional data and nonlinear dependencies where exact inference is impossible, yet avoids the pitfalls of unstructured variational approximations.

The authors note that while the current work focuses on tree-structured models to ensure identifiability and inference coherence, the empirical success of belief propagation on loopy graphs suggests that adapting RAMP to loopy structures is a viable direction for future work.