PiGollum

Principle-Guided Bayesian Optimisation with GP as Experiment Agent

Fusing PiFlow principle-aware reasoning with Gollum LLM-featurised BO for principled scientific discovery.

---

Overview

PiGollum combines two complementary ideas:

After every oracle evaluation, an LLM extracts a generalizable principle from the observed (sequence → outcome) pair. These principles accumulate in a buffer and re-rank the BO acquisition scores at the next iteration — gradually shifting the search from statistical optimality toward mechanistically sound regions of sequence space.

A post-acquisition 3-agent pipeline (Planner → Hypothesis → Scorer) further refines candidate selection after the GP scores are computed, ensuring the final pick is both statistically strong and scientifically coherent.

---

Architecture

                        ┌──────────────────────────────────────────────────┐
  WARM-START (once)     │  LLM → 5 broad principles                       │
                        │  Evidence retrieval → LLM refines each principle │
                        └──────────────────────────────────────────────────┘

  PER BO ITERATION
  ─────────────────────────────────────────────────────────────────────────
  1. Train DeepGP on (train_x, train_y)
  2. GP predicts means + stds for all candidates   ← Experiment Agent

  3. Inner PiFlow loop  (n_inner_steps × per iteration)
     ┌─ Scorer  → action: explore / refine / validate
     ├─ Planner → select candidates + build guidance
     ├─ LLM     → predictive hypothesis per candidate   ← Hypothesis Agent
     ├─ GP      → validate hypothesis (posterior)       ← Experiment Agent
     └─ LLM     → extract principle from (hyp, GP pred) ← Principle Agent

  4. Score all candidates
     final = (1−α) × GP_score  +  α × principle_alignment

  5. Post-acquisition 3-agent pipeline
     ┌─ Agent 1 · Planner    → re-rank principles given new evidence
     ├─ Agent 2 · Hypothesis → directional hypothesis for next candidate
     └─ Agent 3 · Scorer     → re-rank top-k by hypothesis alignment

  6. Greedy select → Oracle evaluation → extract oracle principle
  ─────────────────────────────────────────────────────────────────────────

---

Key Features

• GP as Experiment Agent — the GP posterior validates LLM hypotheses in the inner loop, replacing expensive wet-lab experiments with in-silico validation.
• Principle buffer — discovered principles are stored with embeddings, rewards, and provenance (broad / refined / gp / oracle).
• PiFlow Min-Max action selection — exploration-exploitation scoring drives explore / refine / validate decisions per iteration.
• Post-acquisition 3-agent pipeline — Planner, Hypothesis, and Scorer agents inject scientific reasoning between the acquisition function and greedy selection.
• Three LLM backends — local HuggingFace, any OpenAI-compatible API, or Google Gemini (no extra SDK required).
• W&B integration — logs BO metrics, per-principle scores, and full agent outputs at every iteration.

---

Installation

Prerequisites

• CUDA-capable GPU (A6000 or equivalent recommended for local LLM inference)
• Conda / Mamba
• The gollum repository cloned as a sibling directory (~/gollum/)

Create the environment

# Clone the repo
git clone <repo-url> pigollum
cd pigollum

# Create and activate conda environment
conda env create -f requirements.yaml
conda activate pigollum

Data

Place the biocat dataset under ~/gollum/data/biocat/:

~/gollum/
└── data/
    └── biocat/
        ├── enzymes_sequence_yield_ee_processed_ddg_yeo-johnson.csv   # 200 labelled
        └── enzymes_sequence_10k.csv                                   # 10k candidates

---

Quick Start

Local HuggingFace model (default)

conda activate pigollum
python run_experiment.py \
  --config configs/biocat_pigollum.yaml \
  --data_root ../gollum \
  --output_dir results/biocat_pigollum \
  --n_iters 10

OpenAI-compatible API (GPT-4o-mini, Together, Groq, Ollama, …)

export PIGOLLUM_LLM_API_KEY=sk-...
export PIGOLLUM_LLM_MODEL=gpt-4o-mini          # or llama-3-70b, etc.
# export PIGOLLUM_LLM_BASE_URL=http://localhost:11434/v1  # for local servers

python run_experiment.py --config configs/biocat_pigollum.yaml

Google Gemini

export GEMINI_API_KEY=AIza...
# model defaults to gemini-2.0-flash; override with GEMINI_MODEL=gemini-1.5-pro

Then set backend: gemini in configs/biocat_pigollum.yaml.

With W&B tracking

python run_experiment.py \
  --config configs/biocat_pigollum.yaml \
  --wandb_project my-enzyme-discovery \
  --wandb_run_name biocat-run-01

---

Configuration

All settings live under the pigollum: key in the YAML config.

Core BO settings

Post-acquisition 3-agent pipeline

LLM backend

W&B logging

---

Outputs

All outputs are written to --output_dir (default: results/biocat_pigollum/).

W&B logged metrics (per iteration)

---

Project Structure

pigollum/
├── configs/
│   ├── biocat_pigollum.yaml      # Main experiment config
│   ├── bh_pigollum.yaml          # Placeholder config
│   └── flip2_pigollum.yaml       # Placeholder config
│
├── src/pigollum/
│   ├── bo/
│   │   └── pi_optimizer.py       # PiGollumOptimizer: GP training, inner loop,
│   │                             #   post-acquisition 3-agent pipeline, selection
│   ├── principle/
│   │   ├── buffer.py             # Principle dataclass + in-memory buffer
│   │   ├── extractor.py          # LLM backends + all agent prompts and methods
│   │   ├── scorer.py             # Exploration-exploitation scoring; candidate ranking
│   │   ├── planner.py            # Inner loop orchestration; guidance synthesis
│   │   └── journal.py            # Iteration snapshots; evolution report generation
│   └── utils/
│       ├── llm_client.py         # OpenAI-compatible + Gemini client builders
│       └── sequence_utils.py     # Amino acid composition → text description
│
├── run_experiment.py             # Experiment entry point
├── setup.py                      # Package metadata
└── requirements.yaml             # Conda environment

---

LLM Agent Prompts

The three post-acquisition agents follow the structured format from PiFlow Appendix Q:

---

Acknowledgements

• PiFlow — Principle-Aware Scientific Discovery with Multi-Agent Collaboration (Pu et al., 2026)
• Gollum — LLM-guided Bayesian Optimisation for molecular design
• BoTorch — Bayesian Optimisation in PyTorch
• ESM-2 — Evolutionary Scale Modeling for protein sequences