Elicitation without Backpropagation

Official code for Elicitation without Backpropagation: Steering Model Behavior by Optimizing the Latent Posterior.

Given a Bayes-filtered transformer (BFT) π_θ, find a hard prompt z ∈ {0,1}^m maximizing a continuation-utility objective J(z) = E[U(Y) | z] — without backpropagating through the transformer.

Quick start — reproduce all tables and figures

pip install uv
git clone https://github.com/timaeus-research/elicitation.git
cd elicitation
uv sync

# Place the data archive elicitation_suff_stats_tilt_raw_export.tar.gz (or .zip)
# in the repo root (available from the authors on request), then:
./reproduce_all.sh --from data

This extracts the tarball, sets up checkpoints and results, and produces all paper tables and figures into figures/.

Requirements

• Python ≥ 3.10
• GPU with CUDA (for model evaluation and PMC generation)

pip install uv
uv sync

Pre-trained models

Two BFT checkpoints are included in checkpoints/:

Model	DGP	Architecture	Params	Path
k0_bos	Beta–Bernoulli (k=0)	1-layer, d=64	34K	checkpoints/k0_bos/
k1_learned_bos	Reinforced urn (k=1)	8-layer, d=256	11.4M	checkpoints/k1_learned_bos/

To retrain from scratch:

./reproduce_all.sh --stage train

Evaluation

Run a single experiment (PPT-RB with analytic prior, frequency match, k=0, m=6):

uv run python run_all_methods.py \
    --model k0_bos --utility freq_match --p_target 0.3 \
    --methods 3 --m_prompts 6 --seed 1 --init random \
    --pmc_source dirichlet --outdir results/k0_m6

Run GCG with top-K=1 gradient filtering:

uv run python run_all_methods.py \
    --model k0_bos --utility freq_match --p_target 0.3 \
    --methods 5 --gcg_top_k 1 --m_prompts 6 --seed 1 --init random \
    --outdir results/k0_m6

Methods

Five methods are compared across two DGPs (k=0 Beta–Bernoulli, k=1 Reinforced urn), two prompt lengths (m ∈ {6, 50}), and sweeps over utility parameterizations. Each cell is run with 10 random seeds.

Method	Description
GCG	Greedy Coordinate Gradient with top-K=1 gradient-filtered candidate selection
PPT-RB (A)	Posterior Prefix Tuning, Rao-Blackwellized, analytic prior samples
PPT-RB	Posterior Prefix Tuning, Rao-Blackwellized, PMC prior samples
PPT (A)	Posterior Prefix Tuning, analytic prior samples
PPT	Posterior Prefix Tuning, PMC prior samples

Reproducing paper tables and figures

From pre-computed results (recommended):

./reproduce_all.sh --stage figures

Individual scripts:

# Results tables (LaTeX, to stdout)
uv run python scripts/generate_results_tables.py \
    --results-dir results \
    --gcg-topk-json results/gcg_topk1_all.json \
    --j-cache-dir results/

# ESS tables (LaTeX, to stdout)
uv run python scripts/generate_ess_table.py

# PMC validation plots (replot from cached data, no GPU needed)
uv run python validate_pmc.py \
    --replot checkpoints/pmc/k0_bos/pmc_n0.pkl \
             checkpoints/pmc/k1_learned_bos/pmc_n0.pkl

Full reproduction

./reproduce_all.sh --from data       # Extract tarball → figures (recommended)
./reproduce_all.sh --from scratch    # Train → PMC → experiments → figures
./reproduce_all.sh --from models     # PMC → experiments → figures
./reproduce_all.sh --from pmc        # Experiments → figures

Individual stages:

./reproduce_all.sh --stage train        # Retrain BFTs from scratch
./reproduce_all.sh --stage pmc          # Generate PMC posteriors
./reproduce_all.sh --stage experiments  # Run all method × utility × seed combos
./reproduce_all.sh --stage figures      # Generate paper tables and figures

Repository structure

Core library
  k_markov_elicitation.py         BFT model, meta_train, PMC, PPT optimization
  gcg_baseline.py                 GCG baseline (top-K gradient-filtered coordinate search)

Experiment drivers
  run_all_methods.py              Single-experiment CLI (methods 3=PPT-RB, 4=PPT, 5=GCG)
  run_gcg_topk_test.py            GCG top-K=1 sweep, m=6
  run_gcg_topk_test_m50.py        GCG top-K=1 sweep, m=50

Table and figure generation
  scripts/generate_results_tables.py  Paper results tables (LaTeX)
  scripts/generate_ess_table.py       ESS/L diagnostic tables (LaTeX)
  build_j_cache.py                    Enumerate J(z) for all 2^m prompts (m=6)
  validate_pmc.py                     PMC posterior validation and plotting
  simple_bar_plots.py                 Per-configuration bar plots (alternative view)
  summarize_ranks.py                  Density-warped strip plots (alternative view)
  summarize_tables.py                 HTML summary tables

Infrastructure
  reproduce_all.sh                End-to-end reproduction script
  scripts/train_bft.py            BFT training (AdamW, cosine LR, bf16)
  scripts/run_tables.py           Batch dispatch across GPU instances
  scripts/add_topk_to_tarball.py  Append GCG results to data tarball

Data
  checkpoints/                    Pre-trained BFT weights (tracked in git)
  results/*.json                  Pre-computed GCG top-K and J-cache data (tracked in git)
  results/                        Experiment pkls (from tarball or reproduce_all.sh)

Notes on reproducibility

Rank tables assign tied prompts (J within 1e-5) the same rank, so they are stable across hardware. The J-caches used for ranking are exhaustive evaluations of J(z) for all 2^m prompts; cross-hardware float differences (~1e-6) exist but fall below the tie-breaking threshold. The --from data path uses pre-computed J-caches from the tarball; --from scratch regenerates them on your GPU.