LLaMPPL: SMC Steering for LLMs

• LLaMPPL is a probabilistic programming framework that models constrained text generation as a Feynman-Kac process using SMC methods.
• It integrates a custom SMC inference engine and shared caching with LLaMA-family transformers for efficient, diversity-preserving decoding.
• Empirical results show that LLaMPPL achieves sampling-correct, constraint-respecting completions with runtime close to traditional beam search.

LLaMPPL is a probabilistic programming framework for controlling LLMs at inference time using sequential Monte Carlo (SMC) methods, with an explicit focus on enforcing syntactic and semantic constraints during text generation. Developed to address the limitations of prompt-based steering, LLaMPPL enables users to represent complex constrained or structured generation tasks as probabilistic programs and then sample strings from the exact posterior induced by those programs rather than resorting to heuristic decoding mechanisms. The framework provides a lightweight Python API for rapid prototyping of such inference-time constraints, with particular emphasis on integration with LLaMA-family Transformers. The core methodology and implementation are described in detail in "Sequential Monte Carlo Steering of LLMs using Probabilistic Programs" (Lew et al., 2023).

1. Design Goals and Architecture

LLaMPPL is designed to let users "write down" a broad class of constrained text-generation tasks as concise probabilistic programs, which are then executed using an SMC-based algorithm that guarantees sampling-correct completions under the model's posterior. Unlike ad-hoc constraint satisfaction or beam-search masking strategies, LLaMPPL encodes generation requirements into a Feynman-Kac model composed of a Markov sequence model (typically autoregressive token selection) with non-negative potentials enforcing constraints.

The architecture consists of:

• A minimal Python probabilistic programming library: All generative tasks are encoded by subclassing a Model object and implementing a step() method. This method interacts with five primitives:
  • sample(dist [, proposal])
  • condition(boolean)
  • observe(dist, value)
  • transformer(context)
  • finish()
• Underlying Feynman-Kac transformer model: Each probabilistic program defines an initial state, a sequence of Markov kernels $M_t$, and potentials $G_t$.
• Custom SMC inference engine: Operates across a particle population, sharing a single CachedTransformer to minimize expensive repeated model invocations, and utilizing a "without-replacement" stratified resampling strategy to maintain particle diversity.
• Efficient transformer integration: Special support for the LLaMA transformer family, using prefix tries and shared key/value caches across particles to ensure competitive runtime relative to standard beam search.

2. Probabilistic Program Representation for LLM Steering

Every LLaMPPL program defines:

• State: The current partial output string $s$, typically initialized in __init__().
• Actions: The step() method may at each token:
  • sample(D): Draw a token from distribution $D$ (e.g., transformer output).
  • sample(D, proposal=Q): Draw from proposal $Q$ for importance sampling, with $G_t$ weighted by $D(v)/Q(v)$.
  • observe(D, v): Condition on the next token being $v$, factoring in $D(v)$.
  • condition(bool): Zero out particle weight if boolean is false.
  • transformer(ctx): Return the next-token categorical LM distribution for context $ctx$.
  • finish(): Append EOS and terminate.
• Markov process: The chain $G_t$0 generates token sequences, with transitions governed by $G_t$1 and weighted by $G_t$2.
• Posterior sampling objective: The posterior over completed strings is proportional to $G_t$3, enabling explicit definition of desired constraints or structure.

3. Mathematical Framework: Feynman-Kac Models and SMC Steering

3.1 Feynman-Kac Sequence Models

Let $G_t$4 be the LM vocabulary, $G_t$5 the set of all finite sequences, EOS $G_t$6, and $G_t$7 the set of EOS-terminated strings. The generation process is a stochastic chain with:

• Initial state $G_t$8 (typically the prompt or empty string).
• At each time step $G_t$9, a Markov kernel $s$0 selects new tokens.
• Non-negative potential functions $s$1 encode arbitrary constraints.

The unnormalized posterior is: $s$2 with normalization

$s$3

3.2 Sequential Monte Carlo (SMC) Algorithm

LLaMPPL's SMC implementation proceeds as follows:

1. Expansion: Non-EOS particles are cloned $s$4 times, each extended using the proposal kernel $s$5.
2. Weight Update: Each clone's weight is updated via

$s$6

with $s$7 the current active particle/cloning count.

1. Normalization: All particle weights normalized to sum to 1.
2. Without-replacement Resampling: Compute minimal $s$8 to ensure $s$9; select particles deterministically or stochastically, with unbiased estimator for $D$0.
3. Termination: Repeat until all particles have produced EOS-terminated strings.

This "without-replacement" mechanism preserves sample diversity better than multinomial SMC and yields more faithful posterior inference.

4. LLaMPPL Python API and Example Use Cases

The LLaMPPL API centers on subclassing Model and defining the step() method. Selected task examples:

A. Word Length Constraint

$Q$1

B. Infilling between Fragments

$Q$2

C. Prompt Intersection (Product-of-Experts)

$Q$3

These high-level abstractions make explicit a range of constraints through direct probabilistic modeling.

5. Integration with LLaMA-family Transformers

LLaMPPL supports efficient inference-time control of Meta's LLaMA family by introducing a CachedTransformer object:

• Prefix Trie: Tracks all prefixes seen by any particle, storing for each:
  • Token sequence
  • Cached logits from last expansion
  • Full key/value (KV) caches from all transformer layers
• Cache Traversal: Logits for any context $D$1 can be retrieved in $D$2 time if present, or else by running the transformer only on the difference from the longest shared prefix (“new suffix mode”), after which logits and KV activations are cached.
• Efficiency: All particles share this cache, thus wall-clock time matches that of beam search with width $D$3 (number of particles), not $D$4 independent LM runs. Empirically, SMC decoding with $D$5 particles and $D$6 clones is only $D$7 as costly as beam search with beam size $D$8, but yields a diverse set of solutions faithful to imposed constraints (Lew et al., 2023).

6. Experimental Results and Performance Evaluation

Empirical findings using LLaMPPL with LLaMA-7B demonstrate:

• Prompt Intersection: Comparing two Feynman-Kac model formulations for generating sentences likely under two prompts, Model B (using product-of-logits proposals) achieves higher average unbiased log-normalizers $D$9 and qualitatively better prompt intersections than Model A, across a range of particle counts $Q$0.
• Sample Diversity and Cost: SMC steering with LLaMPPL produces diverse samples fulfilling hard constraints, avoiding the mode collapse associated with greedy or beam search. Runtime remains within a small constant factor of standard beam search, due to shared data structures and KV cache optimization.
• Constraint Satisfaction: With properly specified probabilistic programs, SMC steering yields "sampling-correct" constrained completions, validating the approach's rigor and extensibility (Lew et al., 2023).

7. Implications, Limitations, and Extensions

LLaMPPL illustrates that the probabilistic programming paradigm, paired with optimized SMC inference, can deliver principled, practical solutions for imposing arbitrary constraints on LLM outputs without retraining or functionally intrusive search hacks.

• Scalability and Generality: The approach generalizes to any constraint expressible as a potential in a Feynman-Kac model, ranging from syntactic templates to semantic requirements and multi-prompt intersections.
• Limitations: The method's performance and memory costs scale with particle number and per-step expansion, constraining its use for extremely long or high-dimensional generation tasks. All constraints must be tractably encoded as Markov kernels and local potentials.
• Extension Potential: The shared cache and inference machinery could enable extension to other transformer models, or hybrid approaches incorporating human feedback or symbolic reasoning.

In conclusion, LLaMPPL formalizes constrained language generation as probabilistic inference and makes SMC-based steering practically viable for LLaMA-family LLMs, producing faithful, diverse, and constraint-respecting outputs at competitive inference cost (Lew et al., 2023).