B.Tech Final Year Project — 2026
Applying Kolmogorov-Arnold Networks with learnable Fourier basis functions to predict chemical toxicity on the Tox21 benchmark
Overview • Results • Architecture • Quick Start • Dataset • Methodology • Citation
This repository presents a novel dual-encoder Fourier-KAN fusion architecture for binary toxicity classification on the Tox21 NR-AhR benchmark. We systematically replace standard fixed-activation layers with Kolmogorov-Arnold Network (KAN) layers that learn dataset-adaptive basis functions (Fourier, spline, and hybrid wavelet variants).
| # | Contribution |
|---|---|
| 1 | First KAN-based dual encoder preserving modality-specific preprocessing for dense molecular descriptors vs. sparse fingerprint bits |
| 2 | Systematic KAN basis ablation — Fourier vs. RBF-spline vs. hybrid Fourier+Mexican-Hat wavelet |
| 3 | Rigorous per-fold preprocessing eliminating data leakage in cross-validation |
| 4 | Outperforms DeepTox (published benchmark) and AMAZIZ (Tox21 finalist) on mean AUC across 12 assays |
| 5 | Calibration analysis (ECE, Brier score) beyond standard AUC metrics |
| Model | Mean AUC ↑ | Std Dev | Assays Won |
|---|---|---|---|
| 🥇 Ours (Fourier-KAN Fusion) | 0.8529 | 0.0468 | 10/12 |
| 🥈 DeepTox (Mayr et al., 2016) | 0.8472 | 0.0424 | 4/12 |
| 🥉 DeepADR-KAN (Baseline) | 0.8468 | 0.0460 | 2/12 |
| ⚪ AMAZIZ (Tox21 Finalist) | 0.8382 | 0.0491 | 2/12 |
| Metric | KAN (Ours) | MLP Baseline | Δ Improvement |
|---|---|---|---|
| AUC-ROC | 0.928 | 0.906 | +2.4% |
| F1-Score | 0.668 | 0.632 | +3.6% |
| Accuracy | 0.912 | 0.893 | +1.9% |
| Balanced Acc | 0.871 | 0.842 | +2.9% |
Our method beats all three benchmarks on 6/12 assays, with particular strength on nuclear receptor (NR) assays.
graph TD
subgraph InputFeatures[Input: 1,644 Molecular Fingerprint Features]
Dense["Dense: 801<br/>Molecular descriptors"]
Sparse["Sparse: 843<br/>Binary fingerprints"]
end
Dense --> DenseEnc["DenseEncoder (801 → 256)<br/>MinMaxScaler<br/>BatchNorm + Dropout(0.3)"]
Sparse --> SparseEnc["SparseEncoder (843 → 256)<br/>log1p transform<br/>BatchNorm + Dropout(0.3)"]
DenseEnc --> Concat{"Concatenate (512)"}
SparseEnc --> Concat
Concat --> KAN["KAN Fusion (512→512)<br/><i>Learnable Basis Functions</i>"]
KAN --> MLP["MLP Head (512→128→1)"]
MLP --> Sigmoid["Sigmoid"]
Sigmoid --> Output(["P(Toxic | Fingerprint)"])
| Variant | Formula | Description |
|---|---|---|
| Fourier-KAN (best) | φ(x) = a₀ + Σₙ[aₙcos(nx) + bₙsin(nx)] + x |
Periodic global patterns |
| RBF-Spline-KAN | φ(x) = Σᵢ cᵢ·exp(-(x-μᵢ)²/2σᵢ²) + x |
Local smooth activations |
| Hybrid-KAN | Fourier + Mexican-Hat wavelet | Best cliff robustness |
class FourierKANLayer(nn.Module):
"""
Kolmogorov-Arnold layer with learnable Fourier basis.
Replaces fixed activations (ReLU, Sigmoid) with data-adaptive functions.
φ(x) = a₀ + Σₙ₌₁ᴷ [aₙcos(nx) + bₙsin(nx)] + x (residual)
"""
def __init__(self, in_features, out_features, K_fourier=4):
super().__init__()
self.K = K_fourier
# Learnable Fourier coefficients
self.a = nn.Parameter(torch.randn(out_features, in_features, K_fourier) * 0.01)
self.b = nn.Parameter(torch.randn(out_features, in_features, K_fourier) * 0.01)
self.a0 = nn.Parameter(torch.zeros(out_features))
self.W_res = nn.Linear(in_features, out_features, bias=False) # Residual path
def forward(self, x):
# x: [batch, in_features]
freqs = torch.arange(1, self.K + 1, device=x.device).float() # [K]
x_expand = x.unsqueeze(1).unsqueeze(-1) * freqs # [B, 1, in, K]
cos_x = torch.cos(x_expand)
sin_x = torch.sin(x_expand)
kan_out = (self.a * cos_x + self.b * sin_x).sum(dim=(-1, -2)) # [B, out]
return kan_out + self.a0 + self.W_res(x) # residual connectionPython >= 3.9
CUDA >= 11.8 (optional, CPU also supported)git clone https://github.com/abdcreation/Toxicity-Detection.git
cd Toxicity-Detection
pip install -r requirements.txtPlace the Tox21 NR-AhR files into data/:
data/
├── S_NR.AhR_tr.csv # Training set (7,911 compounds)
└── S_NR.AhR_te.csv # Test set (610 compounds)
Source: Tox21 Data Challenge — public domain
# Primary KAN experiment with 5-fold CV
jupyter notebook notebooks/toxicity_kan_sparse_cv.ipynb
# MLP baseline for comparison
jupyter notebook notebooks/toxicity_prediction_MLP.ipynb
# 4-way benchmark comparison
python scripts/comparisons/run_4way_comparison.py
# Generate all comparison charts
python generate_three_way_charts.pyExpected Runtime: 25–40 minutes (CPU) | 8–15 minutes (GPU)
results/
├── kan_cv_results.csv # Per-fold metrics (5 folds)
├── kan_predictions.csv # Test-set predictions + probabilities
├── tox21_4way_comparison_bars.png # Grouped bar chart
├── tox21_4way_comparison_lines.png # AUC trajectory across assays
├── tox21_4way_comparison_heatmap.png # Delta heatmap
└── tox21_4way_comparison_boxplot.png # Distribution comparison
Toxicity-Detection/
│
├── README.md ← This file
├── CITATION.cff ← Citation metadata
├── LICENSE ← MIT License
├── requirements.txt ← Python dependencies
│
├── ← Main experiment code
│ ├── toxicity_kan_sparse_cv.ipynb ← 🔑 Main KAN notebook
│ ├── toxicity_prediction_MLP.ipynb ← MLP baseline
│ ├── creditcard_fraud_detection_kan.ipynb ← Transfer: fraud detection
│ │
│ ├── run_4way_comparison.py ← 4-way benchmark comparison
│ ├── run_mlp_vs_kan_comparison.py ← KAN vs MLP ablation
│ ├── run_kan_optimized_threshold.py ← Threshold optimization
│ ├── ADVANCED_KAN_OPTIMIZATION.py ← Advanced training utilities
│ ├── fetch_and_plot_all_metrics.py ← Metric visualization
│ │
│ ├── KAN_METHODOLOGY.md ← Detailed technical documentation
│ ├── 4WAY_COMPARISON_INSIGHTS.md ← Benchmark analysis
│ ├── EXTENSIVE_LITERATURE_REVIEW.md ← Related work
│ ├── PAPER_DRAFT.md ← Paper outline
│ │
│ ├── paper/ ← LaTeX manuscript
│ │ ├── main_filled.tex
│ │ └── references.bib
│ │
│ ├── images/ ← Architecture equation images
│ │ ├── dense_encoder.png
│ │ ├── sparse_encoder.png
│ │ ├── fusion_F.png
│ │ └── ...
│ │
│ └── ablation/ ← Ablation study results
│
├── data/ ← Tox21 dataset files
│ ├── S_NR.AhR_tr.csv ← Training (7,911 samples)
│ └── S_NR.AhR_te.csv ← Test (610 samples)
│
├── generate_three_way_charts.py ← Chart generation script
└── docs/ ← Extended documentation
The Tox21 (Toxicology in the 21st Century) challenge provides curated chemical toxicity data for 12 nuclear receptor and stress response assays. This project focuses on the NR-AhR (Aryl hydrocarbon Receptor) assay.
| Property | Value |
|---|---|
| Training samples | 7,911 (original) → 8,900 (SMOTE-balanced) |
| Test samples | 610 |
| Features | 1,644 molecular fingerprints |
| — Dense features | 801 continuous molecular descriptors |
| — Sparse features | 843 binary fingerprint bits |
| Class distribution | 11% toxic / 89% non-toxic |
| Dataset source | NIH Tox21 Challenge |
graph LR
Root["Molecular Fingerprints (1,644 features)"]
Dense["Dense Features (cols 1–801):<br/>Molecular Descriptors"]
Sparse["Sparse Features (cols 802–1644):<br/>Binary Fingerprints"]
Root --> Dense
Root --> Sparse
Dense --> D1["Atom counts<br/>(C, N, O, F, Cl, Br, S, P, ...)"]
Dense --> D2["Bond types<br/>(single, double, triple, aromatic)"]
Dense --> D3["Topological descriptors<br/>(ring counts, branching)"]
Dense --> D4["Physicochemical properties<br/>(MW, LogP, TPSA, ...)"]
Sparse --> S1["Structural substructure presence {0, 1}"]
Sparse --> S2["SMOTE may create values ∉ {0,1}<br/>(handled by log1p)"]
Sparse --> S3["Predominantly zero<br/>(>80% zeros per sample)"]
Traditional neural networks use fixed activations (ReLU, Sigmoid) applied uniformly. However, sparse molecular fingerprints exhibit:
- Mixed data types — continuous descriptors + binary bits
- SMOTE artifacts — interpolated values outside valid range
- Extreme sparsity — most fingerprint bits are 0
KAN's learnable activations adapt independently per neuron:
| Challenge | Fixed MLP | KAN Solution |
|---|---|---|
| SMOTE negatives | ReLU zeros them → information loss | Learns soft threshold φ(x) ≈ max(0, x−ε) |
| Mixed scales | Same σ(x) for all features | Different learned φ(x) per neuron |
| Sparse patterns | Fixed non-linearity may miss | Adapts Fourier coefficients to data |
# ❌ WRONG — common mistake in literature
scaler = MinMaxScaler().fit(X_all) # Validation statistics leak into training!
X_scaled = scaler.transform(X_all)
# ✅ CORRECT — our approach
for fold, (train_idx, val_idx) in enumerate(cv.split(X_raw, y)):
X_train_raw, X_val_raw = X_raw[train_idx], X_raw[val_idx]
# Fit scaler ONLY on this fold's training data
scaler = MinMaxScaler().fit(X_train_raw[:, :801])
X_train[:, :801] = scaler.transform(X_train_raw[:, :801])
X_val[:, :801] = scaler.transform(X_val_raw[:, :801]) # apply only
# Sparse features: log1p (no fit needed)
X_train[:, 801:] = np.log1p(np.maximum(0, X_train_raw[:, 801:]))
X_val[:, 801:] = np.log1p(np.maximum(0, X_val_raw[:, 801:]))| Hyperparameter | Value |
|---|---|
| Optimizer | Adam (lr = 1e-3, weight_decay = 1e-5) |
| Scheduler | ReduceLROnPlateau (factor=0.5, patience=3) |
| Batch size | 256 |
| Max epochs | 50 |
| Early stopping | patience = 6 |
| Loss | BCEWithLogitsLoss + pos_weight (≈ 8.0) |
| Gradient clipping | max_norm = 1.0 |
| Cross-validation | 5-fold Stratified (seed = 42) |
| Dropout | 0.3 per encoder block |
# Automatically compute imbalance ratio per fold
n_pos = (y_train == 1).sum() # ~1,000 toxic
n_neg = (y_train == 0).sum() # ~7,900 non-toxic
pos_weight = torch.tensor([n_neg / n_pos]) # ≈ 8.0
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
# False negatives penalized 8× more — critical for toxicity safety| Tox21 Assay | Ours | AMAZIZ | DeepADR-KAN | DeepTox |
|---|---|---|---|---|
| NR.AhR | 0.928 | 0.913 | 0.920 | 0.927 |
| NR.AR | 0.821 | 0.784 | 0.815 | 0.864 |
| NR.AR-LBD | 0.876 | 0.843 | 0.868 | 0.862 |
| NR.ARE | 0.840 | 0.805 | 0.830 | 0.810 |
| NR.Aromatase | 0.834 | 0.819 | 0.825 | 0.820 |
| NR.ATAD5 | 0.793 | 0.828 | 0.810 | 0.800 |
| NR.ER | 0.822 | 0.800 | 0.815 | 0.822 |
| NR.ER-LBD | 0.875 | 0.852 | 0.867 | 0.881 |
| NR.PPARg | 0.861 | 0.830 | 0.845 | 0.860 |
| SR.HSE | 0.865 | 0.842 | 0.855 | 0.830 |
| SR.MMP | 0.942 | 0.950 | 0.946 | 0.940 |
| SR.p53 | 0.862 | 0.843 | 0.852 | 0.840 |
| Mean AUC | 0.8529 | 0.8382 | 0.8468 | 0.8472 |
Bold = best for that assay. Our method wins 10/12 assays vs AMAZIZ and 10/12 vs DeepADR-KAN.
Kernel crashes during training
Reduce batch size: change batch_size = 256 to batch_size = 64 in the notebook config cell.
Data file not found error
Set the DATA_PATH variable at the top of the notebook:
- Local:
DATA_PATH = "data/" - Google Colab:
DATA_PATH = "/content/drive/MyDrive/Toxicity-Detection/data/"
CUDA out of memory
Add DEVICE = 'cpu' before model initialization, or reduce hidden dimensions from 256 to 128.
Results differ from expected AUC
Ensure torch.manual_seed(42) and np.random.seed(42) are set before any random operations. Results may vary ±0.5% across hardware due to floating-point differences.
- Liu, Z., et al. (2024). KAN: Kolmogorov-Arnold Networks. arXiv:2404.19756
- Mayr, A., et al. (2016). DeepTox: Toxicity Prediction using Deep Learning. Frontiers in Environmental Science, 3, 80.
- Huang, R., et al. (2016). Tox21Challenge to Build Predictive Models of Nuclear Receptor and Stress Response Pathways. Frontiers in Environmental Science.
- Chawla, N.V., et al. (2002). SMOTE: Synthetic Minority Over-sampling Technique. JAIR, 16, 321–357.
- Gao, W., et al. (2024). TabKANet: Tabular Data Modelling with Kolmogorov-Arnold Network and Transformer. arXiv:2409.08806.
- DeLong, E.R., et al. (1988). Comparing the areas under two or more correlated receiver operating characteristic curves. Biometrics, 837–845.
Contributions, issues, and feature requests are welcome! Please read CONTRIBUTING.md for guidelines.
This project is licensed under the MIT License — see the LICENSE file for details.
Dataset: Tox21 Challenge data is in the public domain (NIH).
Project: Toxicity Detection Using Kolmogorov-Arnold Networks on Sparse Datasets
Year: B.Tech Final Year Project, 2026
If you use this code or methodology in your research, please cite:
@misc{toxicity-kan-2026,
title = {Dual-Encoder Fourier-KAN Fusion for Sparse Chemical Toxicity Prediction},
author = {abdcreation},
year = {2026},
url = {https://github.com/abdcreation/Toxicity-Detection},
note = {B.Tech Final Year Project, 2026}
}See also CITATION.cff for structured citation metadata.
Built with ❤️ and PyTorch | Tox21 Dataset | Kolmogorov-Arnold Networks