MedEvoNet

LLM-guided neural-architecture evolution for medical-image classification. An LLM proposes mutations to a Python file that defines a PyTorch model, every candidate is trained and scored, and the search keeps the architectures with the highest mean validation AUC. This repository provides two self-contained tracks:

• MedMNIST (top-level files) — evolve a CNN from scratch across five MedMNIST datasets.
• Histopathology (histo/) — evolve a multiple-instance-learning aggregator on frozen UNI whole-slide embeddings across TCGA-BRCA, TCGA-NSCLC and CPTAC-NSCLC. See histo/DATA.md to obtain the data.

The companion paper reports, for the MedMNIST track, a 3-layer CNN baseline (1.70M params, 0.884 mean test AUC) evolving into a 341k-parameter architecture at 0.904 mean test AUC (paired-bootstrap p=0.012); and, for the histopathology track, a mean test AUC of 0.955 across the three cohorts.

While this repository targets MedMNIST classification, the underlying principle is general-purpose: any optimization problem for which a scalar evaluation function can be defined over a candidate program is amenable to the same LLM-guided evolutionary loop. Swapping train.py for a different evaluator (and adjusting seed_program.py / the system prompt accordingly) is sufficient to retarget the pipeline to other architecture searches, algorithmic tasks, or any code-level optimization with a measurable objective.

---

Repository layout

MedEvoNet/
├── README.md
├── pyproject.toml          # uv-installable; pulls in the vendored openevolve
├── .env.example            # copy → .env, fill in your API key(s)
├── .gitignore              # ignores .env, runs/, caches
├── config.yaml             # evolution + LLM + MAP-Elites configuration
├── seed_program.py         # initial architecture (3-layer CNN)
├── dataloader.py           # MedMNIST datasets + RGB transform + RAM cache
├── train.py                # train_and_evaluate(program_path) — full eval pass
├── evaluator.py            # 6-line OpenEvolve evaluator that calls train.py
├── azure_llm.py            # Azure OpenAI wrapper (used iff AZURE_OPENAI_API_KEY set)
├── run_evolution.py        # entry point — `python run_evolution.py`
├── run_utils.py            # run-directory bookkeeping + artifacts env var
├── metrics_utils.py        # bootstrap AUC + per-iteration artifact persistence
├── histo/                  # histopathology track (MIL on frozen UNI features)
│   ├── DATA.md             # ← download WSIs + extract UNI features (start here)
│   ├── run_evolution.py    # `cd histo && python run_evolution.py meanpool|vanilla`
│   ├── seeds/              # vanilla_attention.py, mean_pooling.py
│   ├── dataloader.py       # configurable feature/label paths
│   ├── train.py, evaluator.py, config.yaml
│   └── labels/             # bundled train/val/test folds (case,slide,label,fold)
└── openevolve_lib/         # vendored OpenEvolve (Apache-2.0) with MedEvoNet patches
    ├── pyproject.toml
    └── openevolve/         # ← imported as `openevolve` from your scripts

Setup with uv

uv is required (pip works too — see further down).

# 1. Clone
git clone https://github.com/YOUR_ORG/MedEvoNet.git
cd MedEvoNet

# 2. Configure secrets — copy template, then fill in your API key(s)
cp .env.example .env
$EDITOR .env

# 3. Create the venv + install everything (incl. the patched openevolve)
uv sync

# 4. Smoke test: load the seed architecture and count its parameters
uv run python -c "from seed_program import create_model; \
    print(sum(p.numel() for p in create_model().parameters()), 'params')"
# → 1701383 params

The first MedMNIST dataset access will download ~2 GB of .npz files into ~/.medmnist/ (or $MEDMNIST_ROOT if set).

Running evolution

uv run python run_evolution.py

Output is written into a fresh runs/mnist/<timestamp>_<jobid>/ directory:

runs/mnist/20260512_134500_local/
├── run_info.json            # config snapshot, git sha, host, time
├── artifacts/               # per-iteration: <uuid>.json + <uuid>.npz
└── openevolve/              # openevolve's own checkpoints/best/logs
    ├── best/best_program.py
    ├── checkpoints/checkpoint_5, _10, ...
    └── logs/openevolve_*.log

A runs/mnist/latest symlink is kept up to date. The *.json artifact per iteration contains the program source, the per-epoch history per dataset, the best-val-epoch metrics, and bootstrap AUC stats. The companion *.npz holds the raw best-val-epoch probabilities + labels per (dataset, split) so you can re-bootstrap or run paired-bootstrap analyses without retraining.

Configuration

config.yaml controls:

• max_iterations, checkpoint_interval, max_code_length
• LLM provider (llm.api_base, llm.primary_model, llm.temperature)
• MAP-Elites archive feature axes (features: — params × efficiency)
• the system prompt that instructs the LLM what kinds of architectural mutations to attempt

To switch from Azure OpenAI to standard OpenAI just leave AZURE_OPENAI_API_KEY empty in .env and set OPENAI_API_KEY. The runner detects which provider to use at startup.

Single-program evaluation

To evaluate one candidate program end-to-end (10 epochs/dataset, full metrics + artifacts):

uv run python train.py seed_program.py 10

The OpenEvolve loop calls evaluator.evaluate(program_path), which is a 1-line wrapper around the same function.

Histopathology track

The histo/ directory evolves the MIL aggregator on frozen UNI patch embeddings, reusing the same evolution loop, provider selection and artifact format as the MedMNIST track. Because whole-slide images are large, features are pre-extracted, so start with histo/DATA.md — it walks through downloading the TCGA/CPTAC slides and extracting UNI features. Then:

export MEDEVONET_HISTO_FEATURES=/path/to/features   # see histo/DATA.md
cd histo
python run_evolution.py meanpool    # mean-pooling seed (shared baseline)
python run_evolution.py vanilla     # vanilla self-attention seed

The bundled histo/labels/*.csv fix the exact train/validation/test folds (794/95/234 for TCGA-BRCA, 739/156/157 for TCGA-NSCLC, 1505/316/325 for CPTAC-NSCLC), so the split is reproducible even from a partial download.

Installation without uv

python -m venv .venv && source .venv/bin/activate
pip install -e openevolve_lib
pip install -e .

Notes

• The MedMNIST images are loaded at 224×224, with grayscale datasets channel-replicated to RGB so a single architecture can be evaluated across all five tasks.
• For each dataset the model's final nn.Linear is auto-swapped to match the class count (2 / 4 / 7 / 14); seed programs just need to end in a Linear layer.
• Selection signal is pure mean validation AUC. Parameter count is not used in the fitness score itself — only as a behavioural axis of the MAP-Elites archive, so smaller models can occupy different cells but cannot replace a higher-AUC model in the same cell.

Citation

See the companion paper for methodology and results. The OpenEvolve codebase is from codelion/openevolve (Apache-2.0); local patches to controller.py, database.py and process_parallel.py are bundled in openevolve_lib/.

License

Apache-2.0. Bundled OpenEvolve sources remain under their upstream Apache-2.0 license.