Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Zhao, Stephen, Brekelmans, Rob, Makhzani, Alireza, Grosse, Roger
Formato: Preprint
Publicado: 2024
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866911859418857472
author Zhao, Stephen
Brekelmans, Rob
Makhzani, Alireza
Grosse, Roger
author_facet Zhao, Stephen
Brekelmans, Rob
Makhzani, Alireza
Grosse, Roger
contents Numerous capability and safety techniques of Large Language Models (LLMs), including RLHF, automated red-teaming, prompt engineering, and infilling, can be cast as sampling from an unnormalized target distribution defined by a given reward or potential function over the full sequence. In this work, we leverage the rich toolkit of Sequential Monte Carlo (SMC) for these probabilistic inference problems. In particular, we use learned twist functions to estimate the expected future value of the potential at each timestep, which enables us to focus inference-time computation on promising partial sequences. We propose a novel contrastive method for learning the twist functions, and establish connections with the rich literature of soft reinforcement learning. As a complementary application of our twisted SMC framework, we present methods for evaluating the accuracy of language model inference techniques using novel bidirectional SMC bounds on the log partition function. These bounds can be used to estimate the KL divergence between the inference and target distributions in both directions. We apply our inference evaluation techniques to show that twisted SMC is effective for sampling undesirable outputs from a pretrained model (a useful component of harmlessness training and automated red-teaming), generating reviews with varied sentiment, and performing infilling tasks.
format Preprint
id arxiv_https___arxiv_org_abs_2404_17546
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo
Zhao, Stephen
Brekelmans, Rob
Makhzani, Alireza
Grosse, Roger
Machine Learning
Artificial Intelligence
Computation and Language
Numerous capability and safety techniques of Large Language Models (LLMs), including RLHF, automated red-teaming, prompt engineering, and infilling, can be cast as sampling from an unnormalized target distribution defined by a given reward or potential function over the full sequence. In this work, we leverage the rich toolkit of Sequential Monte Carlo (SMC) for these probabilistic inference problems. In particular, we use learned twist functions to estimate the expected future value of the potential at each timestep, which enables us to focus inference-time computation on promising partial sequences. We propose a novel contrastive method for learning the twist functions, and establish connections with the rich literature of soft reinforcement learning. As a complementary application of our twisted SMC framework, we present methods for evaluating the accuracy of language model inference techniques using novel bidirectional SMC bounds on the log partition function. These bounds can be used to estimate the KL divergence between the inference and target distributions in both directions. We apply our inference evaluation techniques to show that twisted SMC is effective for sampling undesirable outputs from a pretrained model (a useful component of harmlessness training and automated red-teaming), generating reviews with varied sentiment, and performing infilling tasks.
title Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2404.17546