mirus08/dann-cost-model

DANN Cost Model — E2ETune

A Domain-Adversarial Neural Network (DANN) surrogate for database query cost estimation, trained to generalise across MySQL and PostgreSQL engines using only 30% target-domain labels (semi-supervised).

This model is the production surrogate used by the E2ETune / GenKnob pipeline to generate synthetic preference rankings for LLM fine-tuning of database knob-tuning advisors.

---

Performance (5-Fold Cross-Validation)

Metric	Pooled	Domain 0 (MySQL-S)	Domain 1 (MySQL-L)	Domain 2 (PG-S)	Domain 3 (PG-S src)
Spearman ρ	0.9846 ± 0.0045	0.9726	0.9282	0.9811	0.9945
MAPE	9.47% ± 1.88%	10.53%	21.84%	12.26%	7.67%
RMSE	1.82 ± 0.34	3.46	4.66	0.66	0.31

Training used only 30% of target-domain labels — a 70% reduction in labelling cost vs. fully supervised.

---

Architecture

Input: tabular knob features + query-plan embedding (MiniLM-L6-v2, dim=384)
  |
  +-- SAINT Encoder (Gf): 2-layer transformer, d_model=128, 4 heads
  |     Feature tokens with value projection -> mean-pooled latent z (dim=128)
  |
  +-- Regression Head (Gy): z || plan_emb || masks || db_oh || hw_oh || ram
  |     256 -> 128 -> 64 -> 1  (LayerNorm + GELU + Dropout)
  |
  +-- Domain Classifier (Gd) via GRL: z -> 128 -> 64 -> 1 (source/target BCE)
        Gradient Reversal Layer forces domain-invariant z

Training objective: Huber (task) + CORAL (covariance alignment) + DANN GRL (adversarial)

Hyperparameter	Value
d_model	128
n_layers	2
n_heads	4
num_features	65
n_masks	43
n_db_engines	4
n_hardware	2
labeled_target_frac	0.3
n_folds	5
best_epoch	69
train_minutes	~183 min

---

Usage

Load (full model — no class import needed)

import torch

ckpt = torch.load("dann_frac03_best_FULL.pt", map_location="cpu")
model = ckpt["model"]
model.eval()

# Inspect inferred arch
print(ckpt["arch"])

Load (state_dict — requires model class)

import torch
from models import DANNCostModel   # from src/models.py in E2ETune-AI4DB

ckpt = torch.load("dann_frac03_best.pt", map_location="cpu")
arch = {  # see dann_frac03_arch.json
    "num_features": 65, "n_masks": 43, "n_db_engines": 4, "n_hardware": 2,
    "d_model": 128, "n_heads": 4, "n_layers": 2, "n_domains": 4, "dropout": 0.15,
}
model = DANNCostModel(**arch)
model.load_state_dict(ckpt["state_dict"])
model.eval()

Inference

import torch

# All inputs are float32 tensors, batch dimension first
cost_pred, z, _ = model(
    x_feat      = torch.zeros(1, 65),   # scaled tabular knob features
    x_mask      = torch.zeros(1, 43),   # binary mask features
    qp_emb      = torch.zeros(1, 384),  # MiniLM-L6-v2 query plan embedding
    db_engine_oh= torch.zeros(1, 4),    # one-hot: [mysql_s, mysql_l, pg_s, pg_l]
    hardware_oh = torch.zeros(1, 2),    # one-hot: [cloud, local]
    ram_gb      = torch.zeros(1, 1),    # normalised RAM in GB
    alpha_grl   = 0.0,                  # 0 = inference mode (GRL inactive)
)
# cost_pred: normalised log-cost, shape (batch,)

---

Files

File	Description
dann_frac03_best_FULL.pt	Full model object (torch.save(model, ...)) — load without class import
dann_frac03_best.pt	State-dict checkpoint (state_dict + config + metrics)
dann_frac03_arch.json	Inferred architecture config
metrics_dann_frac03.json	Full 5-fold CV metrics (per-domain, per-fold)
history_dann_frac03.csv	Per-epoch training history
models.py	Model class definitions (SAINT encoder, GRL, DANNCostModel)

---

Citation

If you use this model, please cite the E2ETune / GenKnob project:

Mendis, M. (2026). GenKnob: LLM-guided database knob tuning via
domain-adapted surrogate cost models. Final Year Project, University of Moratuwa.