Bayesian and causal inference by foundation models

August 29, 2024 — August 29, 2024

approximation
Bayes
generative
language
machine learning
meta learning
Monte Carlo
neural nets
NLP
optimization
probabilistic algorithms
probability
statistics
stringology
time series

As set functions, transformers look a lot like ‘generalized inference machines’. Are they? Can we make them do ‘proper’ inference, in some formal sense?

This is a scrapbook of interesting approaches; Bayesian inference over LLM outputs, understanding in-context learning as Bayesian conditioning, and so on.

Probably connected: LLM explanation via Sparse Autoencoders, causal inference in foundation models, and so on.

Figure 1

First one:

Alireza Makhzani introduces Zhao et al. (2024):

Many capability and safety techniques of LLMs—such as RLHF, automated red-teaming, prompt engineering, and infilling—can be viewed from a probabilistic inference perspective, specifically as sampling from an unnormalized target distribution defined by a given reward or potential function. Building on this perspective, we propose to use twisted Sequential Monte Carlo (SMC) as a principled probabilistic inference framework to approach these problems. Twisted SMC is a variant of SMC with additional twist functions that predict the future value of the potential at each timestep, enabling the inference to focus on promising partial sequences. We show the effectiveness of twisted SMC for sampling rare, undesirable outputs from a pretrained model (useful for harmlessness training and automated red-teaming), generating reviews with varied sentiment, and performing infilling tasks.

Our paper offers much more! We propose a novel twist learning method inspired by energy-based models; we connect the twisted SMC literature with soft RL; we propose novel bidirectional SMC bounds on log partition functions as a method for evaluating inference in LLMs; and finally we provide probabilistic perspectives for many more controlled generation methods in LLMs.

More methods in the references.

1 References

Geiger, Ibeling, Zur, et al. 2024. “Causal Abstraction: A Theoretical Foundation for Mechanistic Interpretability.”
Gloeckler, Deistler, Weilbach, et al. 2024. “All-in-One Simulation-Based Inference.”
Huh, Cheung, Wang, et al. 2024. “The Platonic Representation Hypothesis.”
Kinney, and Lombrozo. n.d. “Building Compressed Causal Models of the World.” Cognitive Psychology.
Korbak, Perez, and Buckley. 2022. “RL with KL Penalties Is Better Viewed as Bayesian Inference.”
Melnychuk, Frauen, and Feuerriegel. 2022. “Causal Transformer for Estimating Counterfactual Outcomes.” In Proceedings of the 39th International Conference on Machine Learning.
Müller, Hollmann, Arango, et al. 2022. “Transformers Can Do Bayesian Inference.”
Nichani, Damian, and Lee. 2024. “How Transformers Learn Causal Structure with Gradient Descent.”
Saengkyongam, Rosenfeld, Ravikumar, et al. 2024. “Identifying Representations for Intervention Extrapolation.”
Scetbon, Jennings, Hilmkil, et al. 2024. “FiP: A Fixed-Point Approach for Causal Generative Modeling.”
von Kügelgen, Besserve, Wendong, et al. 2023. “Nonparametric Identifiability of Causal Representations from Unknown Interventions.” In Advances in Neural Information Processing Systems.
Wang, Xu, Tong, et al. 2021. “InferBERT: A Transformer-Based Causal Inference Framework for Enhancing Pharmacovigilance.” Frontiers in Artificial Intelligence.
Ye, Yang, Siah, et al. 2024. “Pre-Training and in-Context Learning IS Bayesian Inference a La De Finetti.”
Zhao, Brekelmans, Makhzani, et al. 2024. “Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo.” In Proceedings of the 41st International Conference on Machine Learning.