Research · 2025

BioGraphVAE
Bio-Inspired Graph Autoencoder
for Molecular Reconstruction

A lightweight extension of GraphVAE with a biologically inspired funnel attention mechanism that achieves better reconstruction fidelity, smoothness, and diversity on molecular graph datasets.

🚀 Get Started ⭐ View on GitHub
✨ Features

What makes BioGraphVAE special?

BioGraphVAE introduces bio-inspired mechanisms on top of classical graph autoencoders, enabling richer and more faithful molecular reconstructions.

🌿

Funnel Attention

A learnable funnel vector re-weights latent representations, inspired by photosynthetic energy funnels found in nature.

📉

Noise Annealing

Deterministic linear noise schedule prevents posterior collapse and stabilises training across diverse molecular graph distributions.

🧪

Molecular Datasets

Plug-and-play support for QM9 (quantum molecular) and ZINC (drug-like compounds) via PyTorch Geometric.

📊

Rich Evaluation

MSE, Diversity, Smoothness, and Kolmogorov-Smirnov metrics give a multi-dimensional view of reconstruction quality.

🗺️

Latent Visualisation

Built-in PCA, t-SNE, and UMAP projections reveal the geometric structure of learned molecular latent spaces.

⚖️

Baseline Benchmarks

Fair comparison against vanilla VAE, GraphVAE, and GAN baselines on identical train/test splits.

🏗️ Architecture

Model Architecture

BioGraphVAE stacks a two-layer GCN encoder with global mean pooling, applies the bio-inspired funnel attention, then decodes with an MLP.

Molecular Graph
(nodes, edges)
→
GCNConv ×2
+ ReLU
→
Global Mean Pool
z ∈ ℝ64
→
🌿 Funnel Attention
z · funnel
→
MLP Decoder
64 → node_dim
→
Reconstructed
Node Features

🔵 Encoder

  • GCNConv(node_dim → 64) + ReLU
  • GCNConv(64 → latent_dim)
  • global_mean_pool → z ∈ ℝ64

🌿 Funnel Attention (Novel)

  • funnel = nn.Parameter(ones(latent_dim))
  • z_funnel = z ⊙ funnel
  • Learnable, element-wise modulation
  • Bio-inspired by photosynthesis

🔴 Decoder

  • Linear(latent_dim → 64) + ReLU
  • Linear(64 → node_dim)
  • MSE reconstruction loss

⚙️ Hyperparameters

  • Latent dim: 64
  • Batch size: 64
  • Optimizer: Adam (lr = 1e-3)
  • GPU: CUDA required
📊 Evaluation

Performance Metrics

BioGraphVAE is evaluated on four complementary metrics that together capture reconstruction quality, sample diversity, and distributional fidelity.

📐
MSE
↓ Lower is Better

Mean-Squared Error measuring reconstruction fidelity of node features.

🌈
Diversity
↑ Higher is Better

Standard deviation of generated node features — more diversity means richer outputs.

〰️
Smoothness
↓ Lower is Better

Mean squared finite difference along canonical node indices.

📈
KS Statistic
↓ Lower is Better

Kolmogorov-Smirnov test measuring distributional alignment with real data.

Model Comparison

Model Architecture Graph Structure Funnel Attention
VAE MLP Encoder/Decoder ✗ ✗
GraphVAE GCN Encoder + MLP Decoder ✓ ✗
GAN Generator + Discriminator ✗ ✗
BioGraphVAE Ours GCN + Funnel Attn + MLP ✓ ✓
BioGraphVAE results — performance comparison across models and metrics
Figure 1. Performance comparison of BioGraphVAE against baseline models (VAE, GraphVAE, GAN) across MSE, Diversity, Smoothness, and KS Statistic on molecular graph datasets. BioGraphVAE achieves the best balance of reconstruction fidelity and sample diversity thanks to its bio-inspired funnel attention mechanism.
🗄️ Datasets

Supported Molecular Datasets

Both datasets are loaded natively through PyTorch Geometric and preprocessed automatically by the extraction scripts.

⚛️ QM9

A quantum chemistry dataset of ~134,000 small organic molecules with up to 9 heavy atoms (C, H, O, N, F). Each molecule comes with 19 quantum mechanical properties computed via DFT.

Quantum Chemistry ~134k molecules 9 atom types

💊 ZINC

A curated database of commercially available drug-like compounds used for virtual screening. Contains millions of molecules with rich structural diversity relevant to drug discovery.

Drug Discovery Drug-like High Diversity
🚀 Quickstart

Get Up and Running

Follow these steps to reproduce the BioGraphVAE experiments on your machine.

1

Create a conda environment

Python 3.10 is required for compatibility with PyTorch Geometric.

conda create -n biographvae python=3.10 -y
conda activate biographvae
2

Install dependencies

All required packages are listed in requirements.txt.

pip install torch torch-geometric numpy pandas tqdm pyyaml \
            matplotlib plotly seaborn umap-learn scikit-learn rdkit
3

Preprocess datasets

Downloads and processes QM9 and ZINC datasets.

python main/extract_qm9.py    # Download & preprocess QM9
python main/extract_zinc.py   # Download & preprocess ZINC
4

Train BioGraphVAE

Train on either dataset. GPU (CUDA) is recommended.

python main/experiment_qm9.py    # Train on QM9
python main/experiment_zinc.py   # Train on ZINC
5

Visualise results

Generate PCA, t-SNE, and UMAP projections of the latent space.

python main/visualise.py    # Generates PCA / t-SNE / UMAP plots
📄 Citation

How to Cite

If you use BioGraphVAE in your research, please cite the following paper. 

@inproceedings{chakraborty2025biographvae,
  title     = {BioGraphVAE: A Bio-inspired Autoencoder for Structured Data Reconstruction},
  author    = {Chakraborty, Somyajit and Chakraborty, Debashis and Mondal, Khokan
               and Jana, Angshuman and Gayen, Avijit},
  booktitle = {Proceedings of ...},
  year      = {2025}
}