Skip to content

About

[NeurlPS 2026] Particle-Guided Trunk Repulsion for Diverse Protein Conformations

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Latest commit

 

History

16 Commits

Folders and files

Repository files navigation

ConforFlux

ConforFlux generates diverse protein conformations by applying particle guidance to a model's trunk embeddings. It supports Boltz-2, Protenix and OpenFold3-preview2, with guidance added to each backbone's inference command.

See the paper (NeurIPS 2026) for the method and experiments.

ConforFlux


Installation

Clone the repository and run the installation commands from its root directory.

git clone https://github.com/suzuki-2001/conforflux
cd conforflux

Choose the backbone you want to use. Keep a separate environment for each backbone.

Backbone Environment file Environment name
Boltz-2 envs/boltz.yml conforflux-boltz
Protenix envs/protenix.yml conforflux-protenix
OpenFold3 envs/openfold3.yml conforflux-openfold3

For example, create and activate the Boltz-2 environment with the following commands. You can use conda in place of mamba.

mamba env create -f envs/boltz.yml
conda activate conforflux-boltz

To use an existing environment, install the shared ConforFlux package and the backbone you need. For Boltz-2, run the following commands.

pip install -e .
pip install -e "./boltz[cuda]"

For Protenix or OpenFold3, replace the second command with pip install -e ./protenix or pip install -e ./openfold3.


Usage

Run the command for the backbone installed in your active environment.

Boltz-2

boltz predict input.yaml --out_dir ./out \
    --num_particles 5 --sigma 2.0 --alpha_s 0.02 --alpha_z 0.02 \
    --recycling_steps 3 --sampling_steps 200 --output_format pdb

Protenix

protenix pred -i input.json -o ./out \
    --num_particles 5 --sigma 2.0 --alpha_s 0.02 --alpha_z 0.02

OpenFold3

run_openfold predict --query_json input.json --output_dir ./out \
    --runner_yaml openfold3/of3_runner.yaml --seeds 101 \
    --num_particles 5 --sigma 2.0 --alpha_s 0.02 --alpha_z 0.02

The default guidance uses the paper's Ours (λ) settings. Add --noise_level for Ours (λ_t). Set --num_particles 0 to run the backbone without ConforFlux guidance.


Shared trunk

The default --trunk_mode particle keeps separate single and pair embeddings for each particle, as in the paper. All three backbones also support --trunk_mode shared, which uses one set of embeddings while retaining independent noisy coordinates. Gradients from all particles are summed before RMS normalization and a shared update.

boltz predict input.yaml --out_dir ./out --num_particles 5 --trunk_mode shared

This reduces embedding and gradient storage, but total memory still depends on the particle count. Boltz-2 and OpenFold3 also accept --gradient_checkpointing. Shared mode can change the generated conformations and was not evaluated in the paper.


Combining MSA sampling and masking

Both trunk modes can use the backbone's native MSA subsampling together with --msa_mask_fraction 0.4. Masking replaces 40% of protein MSA columns with unknown tokens in non-query rows before computing the profile. The query remains unchanged. The mask applies to the MSA itself rather than to attention weights.

For Boltz-2, the following example combines native subsampling, masking and shared guidance.

boltz predict input.yaml --out_dir ./out \
    --num_particles 5 --trunk_mode shared \
    --subsample_msa --num_subsampled_msa 256 \
    --msa_mask_fraction 0.4 --sigma 0.5 --seed 41

Try a smaller bandwidth such as --sigma 0.5 or 1.0 when combining these options, and check diversity and quality for your target. --alpha_s and --alpha_z control the update sizes. In the pipeline, use --sigmas 0.5,1.0 instead of --sigma.

The prediction seed determines the mask, shared by all particles in a round. Settings and column indices are saved under msa_mask/. Masking is disabled by default, and the unguided references for --plddt_filter always use unmasked MSAs.


Sampling pipeline

Use --target_samples to collect candidates across multiple guided rounds. The same pipeline options work with all three backbones.

boltz predict input.yaml --out_dir ./out --num_particles 5 --target_samples 50
  1. Generate --num_particles structures per round, using a new seed and cycling through --sigmas.
  2. Apply the geometry filter. Add --plddt_filter to compare confidence against unguided predictions generated first.
  3. Stop when enough candidates pass or --max_rounds is reached. By default, keep up to --target_samples passing structures with the highest mean pLDDT.

The pipeline uses --reg_weight 0.1 by default to pull the embeddings toward their original trunk values. This pullback and the pipeline were not used in the paper.

Structural representatives

Add --representatives to select conformational representatives from the passing candidates. Install Foldseek separately and put it on PATH, or set --foldseek_binary to its executable path.

boltz predict input.yaml --out_dir ./out \
    --num_particles 5 --trunk_mode shared \
    --msa_mask_fraction 0.4 --sigmas 0.5,1.0 \
    --target_samples 100 --representatives 10 --foldseek_threads 4

Here, --target_samples 100 sets the passing candidate target and --representatives 10 sets the final output count. All passing candidates are clustered using Foldseek global TM alignments and average linkage. The highest mean-pLDDT structure in each cluster is retained, and representatives are sorted by pLDDT. Extra candidates from the final round are included in clustering.

If fewer candidates pass, the output count is limited to the available candidates. The requested cluster count controls output size and does not estimate the number of biological states. Selection requires one protein chain per structure. All-pairs alignment cost grows quadratically with the candidate count.

Output files

Each input has its own output directory, named after the input.

Path Contents
rounds/ All generated candidates and per-round mask records when masking is enabled.
kept/ Selected individual structure files.
kept.cif Selected structures in pLDDT order as a multi-model mmCIF file.
samples.tsv Filter results, selection status and cluster IDs when clustering is enabled.
run.json Sampling settings and output counts.
unguided/ Reference predictions when --plddt_filter is enabled.
selection/ Foldseek scores, cluster membership in clusters.json, and commands, version and timing in selection.json.

Use a fresh output directory for a new Foldseek selection run.


Filter

The geometry filter follows ConforMix. A structure must satisfy all three conditions.

  • Consecutive Cα atoms are less than 4.5 Å apart.
  • Peptide C and N atoms are less than 2.0 Å apart.
  • Heavy atoms are more than 0.5 Å apart between residues separated by three or more positions within a chain.

The optional confidence filter uses ConforMix's 10-residue window criterion. Reference pLDDT is the minimum across unguided predictions at each residue. No window may lose more than 20 pLDDT points relative to this reference.

To filter existing predictions, run the following command.

python -m conforflux.filter ./guided --reference ./unguided --out filter.tsv --keep_dir ./kept

Without --reference, only geometry is checked. pLDDT is read from Cα B-factors. For Protenix predictions, use --need_atom_confidence true during generation. The sampling pipeline enables this automatically.


Options

Guidance

Flag Default Meaning
--num_particles 5 Number of coupled particles. Guidance requires at least 2. Use 0 for unguided inference. Replaces the backbone's sample-count flag.
--trunk_mode particle Separate embeddings per particle, or one set with shared.
--sigma 2.0 RBF bandwidth on Cα RMSD in Å for a single run. The paper sweeps 0.5 to 2.5.
--alpha_s / --alpha_z 0.02 RMS-normalised update size for the single and pair embeddings.
--update_interval 3 Apply a guidance update every K diffusion steps.
--noise_level off Scale updates by the current noise level. Can widen the ensemble at a higher clashscore.
--reg_weight 0, or 0.1 in the pipeline Remove this fraction of the displacement from the original trunk embeddings at each update.
--gradient_checkpointing off Reduce peak GPU memory in Boltz-2 and OpenFold3 by recomputing activations.
--seed / --seeds Backbone default Boltz-2 uses --seed. Protenix and OpenFold3 accept comma-separated values with --seeds.

MSA processing

Flag Default Meaning
--msa_mask_fraction 0 Fraction of protein MSA columns masked in non-query rows before profile calculation.
--subsample_msa off Enable the native Boltz-2 MSA row sampler.
--num_subsampled_msa 1024 Number of rows for the Boltz-2 sampler. Requires --subsample_msa.

Protenix and OpenFold3 retain their native MSA sampling settings.

Sampling and selection

Flag Default Meaning
--target_samples 0 Passing candidate target for the pipeline. Use 0 for a single run.
--sigmas 0.5,1.0,1.5,2.0,2.5 Bandwidths cycled through by pipeline rounds.
--plddt_filter off Apply the confidence filter using unguided reference predictions.
--num_reference 5 Number of unguided predictions used by the confidence filter.
--max_rounds 20 Maximum number of guided rounds.
--representatives 0 Number of Foldseek cluster representatives. Requires --target_samples. Use 0 for global pLDDT selection.
--foldseek_binary foldseek Path or command name of the Foldseek executable.
--foldseek_threads 4 CPU threads for Foldseek alignment and score export.

Container

Build a separate image for each backbone. Set BACKBONE to boltz (the default), protenix or openfold3.

apptainer build --build-arg BACKBONE=boltz conforflux.sif container/Singularity.def
docker build --build-arg BACKBONE=protenix -f container/Dockerfile -t conforflux .

See container instructions for running the images.


Citation

@inproceedings{suzuki2026conforflux,
  title     = {ConforFlux: Particle-Guided Trunk Repulsion for Diverse Protein Conformations},
  author    = {Suzuki, Shosuke and Amagasa, Toshiyuki},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}

About

[NeurlPS 2026] Particle-Guided Trunk Repulsion for Diverse Protein Conformations

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Contributors

Languages