|
|
--- |
|
|
license: mit |
|
|
tags: |
|
|
- protein-generation |
|
|
- antimicrobial-peptides |
|
|
- flow-matching |
|
|
- protein-design |
|
|
- esm |
|
|
- amp |
|
|
library_name: pytorch |
|
|
--- |
|
|
|
|
|
# FlowFinal: AMP Flow Matching Model |
|
|
|
|
|
FlowFinal is a state-of-the-art flow matching model for generating antimicrobial peptides (AMPs). The model uses continuous normalizing flows to generate protein sequences in the ESM-2 embedding space. |
|
|
|
|
|
## Model Description |
|
|
|
|
|
- **Model Type**: Flow Matching for Protein Generation |
|
|
- **Domain**: Antimicrobial Peptide (AMP) Generation |
|
|
- **Base Model**: ESM-2 (650M parameters) |
|
|
- **Architecture**: Transformer-based flow matching with classifier-free guidance (CFG) |
|
|
- **Training Data**: Curated AMP dataset with ~7K sequences |
|
|
|
|
|
## Key Features |
|
|
|
|
|
- **Classifier-Free Guidance (CFG)**: Enables controlled generation with different conditioning strengths |
|
|
- **ESM-2 Integration**: Leverages pre-trained protein language model embeddings |
|
|
- **Compression Architecture**: Efficient 16x compression of ESM-2 embeddings (1280 → 80 dimensions) |
|
|
- **Multiple CFG Scales**: Support for no conditioning (0.0), weak (3.0), strong (7.5), and very strong (15.0) guidance |
|
|
|
|
|
## Model Components |
|
|
|
|
|
### Core Architecture |
|
|
- `final_flow_model.py`: Main flow matching model implementation |
|
|
- `compressor_with_embeddings.py`: Embedding compression/decompression modules |
|
|
- `final_sequence_decoder.py`: ESM-2 embedding to sequence decoder |
|
|
|
|
|
### Trained Weights |
|
|
- `final_compressor_model.pth`: Trained compressor (315MB) |
|
|
- `final_decompressor_model.pth`: Trained decompressor (158MB) |
|
|
- `amp_flow_model_final_optimized.pth`: Main flow model checkpoint |
|
|
|
|
|
### Generated Samples (Today's Results) |
|
|
- Generated AMP sequences with different CFG scales |
|
|
- HMD-AMP validation results showing 8.8% AMP prediction rate |
|
|
|
|
|
## Performance Results |
|
|
|
|
|
### HMD-AMP Validation (80 sequences tested) |
|
|
- **Total AMPs Predicted**: 7/80 (8.8%) |
|
|
- **By CFG Configuration**: |
|
|
- No CFG: 1/20 (5.0%) |
|
|
- Weak CFG: 2/20 (10.0%) |
|
|
- Strong CFG: 4/20 (20.0%) ← Best performance |
|
|
- Very Strong CFG: 0/20 (0.0%) |
|
|
|
|
|
### Best Performing Sequences |
|
|
1. `ILVLVLARRIVGVIVAKVVLYAIVRSVVAAAKSISAVTVAKVTVFFQTTA` (No CFG) |
|
|
2. `EDLSKAKAELQRYLLLSEIVSAFTALTRFYVVLTKIFQIRVKLIAVGQIL` (Weak CFG) |
|
|
3. `IKLSRIAGIIVKRIRVASGDAQRLITASIGFTLSVVLAARFITIILGIVI` (Strong CFG) |
|
|
|
|
|
## Usage |
|
|
|
|
|
```python |
|
|
from generate_amps import AMPGenerator |
|
|
|
|
|
# Initialize generator |
|
|
generator = AMPGenerator( |
|
|
model_path="amp_flow_model_final_optimized.pth", |
|
|
device='cuda' |
|
|
) |
|
|
|
|
|
# Generate AMP samples |
|
|
samples = generator.generate_amps( |
|
|
num_samples=20, |
|
|
num_steps=25, |
|
|
cfg_scale=7.5 # Strong CFG recommended |
|
|
) |
|
|
``` |
|
|
|
|
|
## Training Details |
|
|
|
|
|
- **Optimizer**: AdamW with cosine annealing |
|
|
- **Learning Rate**: 4e-4 (final) |
|
|
- **Epochs**: 2000 |
|
|
- **Final Loss**: 1.318 |
|
|
- **Training Time**: 2.3 hours on H100 |
|
|
- **Dataset Size**: 6,983 samples |
|
|
|
|
|
## Files Structure |
|
|
|
|
|
``` |
|
|
FlowFinal/ |
|
|
├── models/ |
|
|
│ ├── final_compressor_model.pth |
|
|
│ ├── final_decompressor_model.pth |
|
|
│ └── amp_flow_model_final_optimized.pth |
|
|
├── generated_samples/ |
|
|
│ ├── generated_sequences_20250829.fasta |
|
|
│ └── hmd_amp_detailed_results.csv |
|
|
├── src/ |
|
|
│ ├── final_flow_model.py |
|
|
│ ├── compressor_with_embeddings.py |
|
|
│ ├── final_sequence_decoder.py |
|
|
│ └── generate_amps.py |
|
|
└── README.md |
|
|
``` |
|
|
|
|
|
## Citation |
|
|
|
|
|
If you use FlowFinal in your research, please cite: |
|
|
|
|
|
```bibtex |
|
|
@misc{flowfinal2025, |
|
|
title={FlowFinal: Flow Matching for Antimicrobial Peptide Generation}, |
|
|
author={Edward Sun}, |
|
|
year={2025}, |
|
|
url={https://huggingface.co/esunAI/FlowFinal} |
|
|
} |
|
|
``` |
|
|
|
|
|
## License |
|
|
|
|
|
This model is released under the MIT License. |
|
|
|
