augtab 0.1.0__tar.gz

augtab-0.1.0/PKG-INFO

@@ -0,0 +1,2060 @@
+Metadata-Version: 2.4
+Name: augtab
+Version: 0.1.0
+Summary: AugTab: Learnable Feature Augmentation for Low-Dimensional Tabular Data
+Author: Al Zadid Sultan Bin Habib, Md Younus Ahamed, Md Asif Bin Syed, Md Samiul Islam, Muntasir Tabasum, Tanpia Tasnim, Md. Ekramul Islam
+Project-URL: Homepage, https://www.zadidhabib.com/augtab.html
+Project-URL: Repository, https://github.com/zadid6pretam/AugTab
+Project-URL: Paper, https://doi.org/10.1007/978-3-032-37670-1_24
+Project-URL: Issues, https://github.com/zadid6pretam/AugTab/issues
+Keywords: tabular-data,tabular-learning,feature-augmentation,deep-learning,machine-learning,low-dimensional-data,pytorch,feature-engineering,feature-interactions,input-adaptive-gating
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Information Analysis
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.23
+Requires-Dist: torch>=2.0
+Provides-Extra: experiments
+Requires-Dist: pandas>=1.5; extra == "experiments"
+Requires-Dist: scikit-learn>=1.2; extra == "experiments"
+Requires-Dist: optuna>=3.0; extra == "experiments"
+Requires-Dist: matplotlib>=3.6; extra == "experiments"
+Requires-Dist: tqdm>=4.64; extra == "experiments"
+Requires-Dist: scipy>=1.9; extra == "experiments"
+Provides-Extra: notebooks
+Requires-Dist: pandas>=1.5; extra == "notebooks"
+Requires-Dist: scikit-learn>=1.2; extra == "notebooks"
+Requires-Dist: optuna>=3.0; extra == "notebooks"
+Requires-Dist: matplotlib>=3.6; extra == "notebooks"
+Requires-Dist: jupyter>=1.0; extra == "notebooks"
+Requires-Dist: notebook>=6.5; extra == "notebooks"
+Requires-Dist: ipykernel>=6.0; extra == "notebooks"
+Requires-Dist: tqdm>=4.64; extra == "notebooks"
+Requires-Dist: scipy>=1.9; extra == "notebooks"
+Provides-Extra: all
+Requires-Dist: pandas>=1.5; extra == "all"
+Requires-Dist: scikit-learn>=1.2; extra == "all"
+Requires-Dist: optuna>=3.0; extra == "all"
+Requires-Dist: matplotlib>=3.6; extra == "all"
+Requires-Dist: jupyter>=1.0; extra == "all"
+Requires-Dist: notebook>=6.5; extra == "all"
+Requires-Dist: ipykernel>=6.0; extra == "all"
+Requires-Dist: tqdm>=4.64; extra == "all"
+Requires-Dist: scipy>=1.9; extra == "all"
+
+# AugTab: Learnable Feature Augmentation for Low-Dimensional Tabular Data
+
+![Status](https://img.shields.io/badge/Status-Accepted-success)
+![ECML PKDD 2026](https://img.shields.io/badge/ECML%20PKDD-2026-purple)
+![Track](https://img.shields.io/badge/Track-Research%20Track-blue)
+[![Paper](https://img.shields.io/badge/Paper-In%20Press-red)](https://doi.org/10.1007/978-3-032-37670-1_24)
+![Feature Augmentation](https://img.shields.io/badge/Focus-Feature%20Augmentation-blueviolet)
+![Low-Dimensional Tabular Data](https://img.shields.io/badge/Data-Low--Dimensional%20Tabular%20Data-teal)
+![Python](https://img.shields.io/badge/Python-3.10%2B-green)
+![PyTorch](https://img.shields.io/badge/PyTorch-2.x-red)
+![Optuna](https://img.shields.io/badge/Optuna-Hyperparameter%20Tuning-brightgreen)
+
+<p align="center">
+  <img src="AugTab_Architecture.png" alt="AugTab Architecture" width="1000">
+</p>
+
+AugTab is a learnable feature augmentation framework for **low-dimensional tabular learning**. It introduces a **Feature Augmentation Layer (FAL)** that enriches the original feature space through complementary **nonlinear projection, explicit cross-interaction, and gated recombination** branches, allowing feature augmentation to be learned jointly with the downstream predictor rather than treated as fixed preprocessing or manual feature engineering. FAL combines a nonlinear MLP branch with a cross-interaction branch to construct candidate augmented features, followed by a learned projection and **input-adaptive gating** mechanism that selectively activates useful augmented dimensions for each sample. The resulting augmented representation retains the original features while appending task-aligned learned features, enabling the model to increase expressive capacity without discarding the original input information.
+
+To keep this expansion controlled, AugTab incorporates regularization for **augmentation sparsity, feature diversity, orthogonality, computational budget, and robustness to distribution shift**. These objectives discourage redundant augmented features, encourage complementarity with the original feature space, and regulate the expected number of active augmented dimensions. The framework also supports perturbation-based drift consistency and missingness-aware learning. Although the Feature Augmentation Layer is designed to be **backbone-agnostic**, the ECML PKDD 2026 implementation instantiates AugTab with a lightweight **MLP backbone**. Theoretical analysis further connects learnable augmentation with increased expressivity, improved input conditioning, and generalization under a bounded gating budget. Across **15 low-dimensional classification and regression datasets**, AugTab achieves the best average rank across classification tasks and the top result on all evaluated regression tasks, with comparisons spanning **54 classification baselines and 16 regression baselines**.
+
+## Overview
+
+**AugTab** is a tabular learning framework for problems where the original feature space is relatively low-dimensional and may provide limited expressive structure because of noise, heterogeneity, weak feature interactions, missingness, or distribution shift.
+
+The key idea is to make **feature augmentation itself learnable**. Instead of applying a fixed polynomial expansion, manually constructing feature interactions, or relying solely on the downstream model to discover useful representations, AugTab inserts a trainable **Feature Augmentation Layer (FAL)** before the prediction backbone.
+
+FAL constructs augmented features using:
+
+- **Cross-interaction branch** - learns explicit second-order feature interactions through factorized transformations.
+- **Nonlinear MLP branch** - learns nonlinear feature lifts from the original input.
+- **Learned projection** - combines and re-mixes complementary augmented representations.
+- **Input-adaptive gating** - determines which augmented dimensions should be activated for each sample.
+- **Budget-aware regularization** - controls the effective augmentation width and computational cost.
+- **Diversity and orthogonality regularization** - encourages augmented features to remain complementary rather than redundant.
+- **Drift consistency** - promotes stable augmented representations under distribution shifts and feature perturbations.
+
+For an input feature vector $x$, AugTab preserves the original features and concatenates them with the gated learned augmentation:
+
+$$
+\tilde{x} = x \oplus \big(g(x) \odot z(x)\big)
+$$
+
+where $z(x)$ represents the candidate augmented features and $g(x)$ provides input-dependent gates controlling their activation.
+
+AugTab is **not tied to a specific low-dimensional dataset**. The Feature Augmentation Layer is designed as a modular front end that can, in principle, be combined with different downstream tabular architectures. In the experiments reported in the paper, AugTab is instantiated with an **MLP backbone**.
+
+## Citation
+
+Al Zadid Sultan Bin Habib, Md Younus Ahamed, Md Asif Bin Syed, Md Samiul Islam, Muntasir Tabasum, Tanpia Tasnim, and Md. Ekramul Islam. **“AugTab: Learnable Feature Augmentation for Low-Dimensional Tabular Data.”** In *Machine Learning and Knowledge Discovery in Databases. Research Track (ECML PKDD 2026)*, Springer, 2027. https://doi.org/10.1007/978-3-032-37670-1_24
+
+BibTeX:
+
+```bibtex
+@inproceedings{habib2026augtab,
+  title     = {AugTab: Learnable Feature Augmentation for Low-Dimensional Tabular Data},
+  author    = {Habib, Al Zadid Sultan Bin and Ahamed, Md Younus and Syed, Md Asif Bin and Islam, Md Samiul and Tabasum, Muntasir and Tasnim, Tanpia and Islam, Md. Ekramul},
+  booktitle = {Machine Learning and Knowledge Discovery in Databases. Research Track},
+  year      = {2027},
+  publisher = {Springer Nature Switzerland},
+  doi       = {10.1007/978-3-032-37670-1_24}
+}
+```
+
+- **Paper:** https://doi.org/10.1007/978-3-032-37670-1_24 *(In Press)*
+- **Project Page:** https://www.zadidhabib.com/augtab.html
+- **GitHub:** https://github.com/zadid6pretam/AugTab
+
+This repository contains the official implementation of **AugTab** together with example notebooks for reproducing the training workflow, hyperparameter tuning, evaluation, and package installation.
+
+The repository currently includes:
+
+- `augtab/augtab.py` - core implementation of the **AugTab** framework, including the Feature Augmentation Layer (FAL), MLP backbone, regularization components, and task-specific interfaces for binary classification, multiclass classification, and regression.
+- `augtab/__init__.py` - package initialization file exposing the main AugTab classes and configuration objects through the `augtab` package.
+- `AugTab Try.ipynb` - example notebook demonstrating how to use AugTab with **Optuna-based hyperparameter tuning** and **5-fold cross-validation**, reporting **mean accuracy ± standard deviation** on the **Water Potability** dataset.
+- `PIP_Install_Check.ipynb` - lightweight notebook for verifying the **PyPI installation**, package imports, and basic AugTab functionality after installation with `pip install augtab`.
+- `pyproject.toml` - modern Python project and build-system configuration used for packaging and PyPI distribution.
+- `setup.cfg` - package metadata and configuration for installation and distribution.
+- `requirements.txt` - Python dependencies required to run AugTab and the accompanying experiments.
+- `AugTab_Architecture.png` - High-level architecture diagram of the AugTab framework.
+
+The example notebook is written so that the dataset file and target column can be changed easily, allowing the same workflow to be reused for other low-dimensional tabular datasets.
+
+In addition, `AugTab Try.ipynb` contains supplementary diagnostic and analysis code used to further inspect AugTab's behavior during training and evaluation.
+
+---
+
+## Repository Contents
+
+```text
+.
+├── augtab/
+│   ├── __init__.py
+│   └── augtab.py
+│
+├── AugTab Try.ipynb
+├── AugTab_Architecture.png
+├── PIP_Install_Check.ipynb
+├── pyproject.toml
+├── setup.cfg
+├── requirements.txt
+├── .gitignore
+├── LICENSE
+└── README.md
+```
+
+### Package Structure
+
+The installable Python package is contained in the `augtab/` directory:
+
+```text
+augtab/
+├── __init__.py
+└── augtab.py
+```
+
+The public API can therefore be imported directly from the package:
+
+```python
+from augtab import (
+    AugTabClassifier,
+    AugTabMulti,
+    AugTabRegressor,
+)
+```
+
+The three primary interfaces correspond to:
+
+- `AugTabClassifier` - binary classification
+- `AugTabMulti` - multiclass classification
+- `AugTabRegressor` - regression
+
+Lower-level configuration and model components are also exposed for users who want more control over the AugTab architecture and training configuration.
+
+### Main dependencies
+
+The repository uses the following main dependencies:
+
+```
+numpy>=1.23
+pandas>=1.5
+scikit-learn>=1.2
+optuna>=3.0
+torch>=2.0
+matplotlib>=3.6
+jupyter>=1.0
+notebook>=6.5
+ipykernel>=6.0
+tqdm>=4.64
+scipy>=1.9
+```
+## Installation
+
+You can install **AugTab** in several ways depending on your workflow.
+
+---
+
+### Option 1: Clone the Repository (Recommended for Development)
+
+```bash
+git clone https://github.com/zadid6pretam/AugTab.git
+cd AugTab
+pip install -r requirements.txt
+pip install -e .
+```
+
+- This is the recommended option if you want to modify the source code, run the provided experiment notebooks, reproduce the Optuna-based tuning workflow, or develop additional AugTab extensions.
+- Editable installation (`-e`) allows changes made inside the local `augtab/` package to be reflected immediately without reinstalling the package.
+
+---
+
+### Option 2: Install Directly from GitHub (No Cloning Needed)
+
+```bash
+pip install "git+https://github.com/zadid6pretam/AugTab.git"
+```
+
+- This installs the latest version of **AugTab** directly from the GitHub repository.
+
+After installation, the task-specific AugTab interfaces can be imported as:
+
+```python
+from augtab import AugTabClassifier, AugTabMulti, AugTabRegressor
+```
+
+---
+
+### Option 3: Use a Virtual Environment
+
+```bash
+python -m venv augtab-env
+
+# macOS / Linux
+source augtab-env/bin/activate
+
+# Windows
+# augtab-env\Scripts\activate
+
+git clone https://github.com/zadid6pretam/AugTab.git
+cd AugTab
+pip install -r requirements.txt
+pip install -e .
+```
+
+- Using a virtual environment is recommended to keep **AugTab** and its dependencies isolated from other Python projects.
+- Once activated, all AugTab dependencies and experiments will run within the isolated environment.
+
+---
+
+### Option 4: Local Install Without Editable Mode
+
+```bash
+git clone https://github.com/zadid6pretam/AugTab.git
+cd AugTab
+pip install -r requirements.txt
+pip install .
+```
+
+- This performs a standard local installation of AugTab.
+- Unlike editable mode, subsequent changes to the source code require reinstalling the package.
+
+After installation:
+
+```python
+from augtab import AugTabClassifier, AugTabMulti, AugTabRegressor
+```
+
+---
+
+### Option 5: Install from PyPI
+
+```bash
+pip install augtab
+```
+
+After installation, the main AugTab interfaces can be imported as:
+
+```python
+from augtab import AugTabClassifier, AugTabMulti, AugTabRegressor
+```
+
+The three high-level interfaces correspond to:
+
+- `AugTabClassifier` - binary classification
+- `AugTabMulti` - multiclass classification
+- `AugTabRegressor` - regression
+
+For example:
+
+```python
+from augtab import AugTabClassifier
+
+model = AugTabClassifier(
+    d_features=10,
+    k_aug=32,
+    device="cuda"
+)
+```
+
+Lower-level configuration and model components are also directly available:
+
+```python
+from augtab import (
+    FALConfig,
+    BackboneConfig,
+    RegularizerConfig,
+    AugTabConfig,
+    AugTabCore,
+)
+```
+
+A lightweight installation check is provided in:
+
+```text
+PIP_Install_Check.ipynb
+```
+
+which can be used to verify the PyPI installation, package imports, and basic AugTab functionality.
+
+## Example Usage
+
+AugTab supports **binary classification**, **multiclass classification**, and **regression** through three task-specific interfaces:
+
+```python
+from augtab import AugTabClassifier, AugTabMulti, AugTabRegressor
+```
+
+For a **new dataset**, we recommend using the **Optuna-tuned workflow** because the optimal augmentation width, FAL capacity, MLP backbone size, regularization strengths, learning rate, batch size, and number of training epochs may vary across datasets.
+
+The examples below therefore provide six common configurations:
+
+1. Binary classification without hyperparameter tuning
+2. Binary classification with Optuna tuning
+3. Multiclass classification without hyperparameter tuning
+4. Multiclass classification with Optuna tuning
+5. Regression with Optuna tuning
+6. Regression without hyperparameter tuning
+
+> **Recommended for new datasets:** Use Examples **2, 4, and 5**, which perform Optuna-based hyperparameter tuning with **5-fold cross-validation**. Keep the test set completely separate from Optuna and use it only for final evaluation.
+
+> **Note:** The examples use synthetic low-dimensional datasets so that they can run without external data files. Replace the synthetic `X` and `y` arrays with your own tabular dataset.
+
+---
+
+### Example 1: Binary Classification Without Hyperparameter Tuning
+
+This example trains `AugTabClassifier` using a fixed configuration.
+
+```python
+import random
+import numpy as np
+import torch
+
+from sklearn.datasets import make_classification
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
+
+from augtab import AugTabClassifier, RegularizerConfig
+
+
+# ============================================================
+# Reproducibility
+# ============================================================
+
+SEED = 42
+
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(SEED)
+
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+print("Device:", DEVICE)
+
+
+# ============================================================
+# Example binary dataset
+# Replace X and y with your own data
+# ============================================================
+
+X, y = make_classification(
+    n_samples=800,
+    n_features=12,
+    n_informative=8,
+    n_redundant=2,
+    n_classes=2,
+    random_state=SEED,
+)
+
+X = X.astype(np.float32)
+y = y.astype(np.int64)
+
+
+# ============================================================
+# Train / test split
+# ============================================================
+
+X_train, X_test, y_train, y_test = train_test_split(
+    X,
+    y,
+    test_size=0.20,
+    random_state=SEED,
+    stratify=y,
+)
+
+
+# ============================================================
+# Standardization
+# Fit preprocessing on training data only
+# ============================================================
+
+scaler = StandardScaler()
+
+X_train = scaler.fit_transform(X_train).astype(np.float32)
+X_test = scaler.transform(X_test).astype(np.float32)
+
+
+# ============================================================
+# Regularization
+# ============================================================
+
+regs = RegularizerConfig(
+    lambda_sparse=1e-3,
+    lambda_div=1e-3,
+    lambda_orth=1e-3,
+    lambda_budget=1e-3,
+    lambda_drift=0.0,
+)
+
+
+# ============================================================
+# Initialize AugTab
+# ============================================================
+
+model = AugTabClassifier(
+    d_features=X_train.shape[1],
+    k_aug=32,
+    kprime=64,
+    h_hidden=64,
+    widths=(128, 128),
+    activation="gelu",
+    append_mask=False,
+    gating="basic",
+    regs=regs,
+    device=DEVICE,
+    lr=2e-3,
+    weight_decay=1e-4,
+)
+
+
+# ============================================================
+# Train
+# ============================================================
+
+model.fit(
+    X_train,
+    y_train,
+    epochs=80,
+    batch_size=64,
+    verbose=False,
+)
+
+
+# ============================================================
+# Evaluate
+# ============================================================
+
+y_pred = model.predict(X_test).numpy()
+y_prob = model.predict_proba(X_test).numpy().reshape(-1)
+
+print("\nBinary Classification Results")
+print("Accuracy :", accuracy_score(y_test, y_pred))
+print("F1       :", f1_score(y_test, y_pred))
+print("ROC-AUC  :", roc_auc_score(y_test, y_prob))
+```
+
+---
+
+### Example 2: Binary Classification with Optuna Hyperparameter Tuning
+
+For a **new binary classification dataset**, this is the recommended workflow.
+
+Optuna searches the AugTab architecture, regularization, and optimization hyperparameters using **5-fold stratified cross-validation** on the training data. The test set remains completely untouched during tuning.
+
+```python
+import random
+import numpy as np
+import optuna
+import torch
+
+from sklearn.datasets import make_classification
+from sklearn.model_selection import train_test_split, StratifiedKFold
+from sklearn.preprocessing import StandardScaler
+from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
+
+from augtab import AugTabClassifier, RegularizerConfig
+
+
+# ============================================================
+# Reproducibility
+# ============================================================
+
+SEED = 42
+
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(SEED)
+
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+print("Device:", DEVICE)
+
+
+# ============================================================
+# Example binary dataset
+# Replace X and y with your own data
+# ============================================================
+
+X, y = make_classification(
+    n_samples=800,
+    n_features=12,
+    n_informative=8,
+    n_redundant=2,
+    n_classes=2,
+    random_state=SEED,
+)
+
+X = X.astype(np.float32)
+y = y.astype(np.int64)
+
+
+# ============================================================
+# Hold out the final test set BEFORE Optuna
+# ============================================================
+
+X_train, X_test, y_train, y_test = train_test_split(
+    X,
+    y,
+    test_size=0.20,
+    random_state=SEED,
+    stratify=y,
+)
+
+
+# ============================================================
+# Optuna search space
+# ============================================================
+
+def sample_params(trial):
+
+    hidden_width = trial.suggest_categorical(
+        "hidden_width",
+        [64, 128, 256],
+    )
+
+    depth = trial.suggest_int(
+        "depth",
+        1,
+        3,
+    )
+
+    return {
+        "k_aug": trial.suggest_categorical(
+            "k_aug",
+            [8, 16, 24, 32, 48, 64],
+        ),
+
+        "kprime": trial.suggest_categorical(
+            "kprime",
+            [16, 32, 64, 128],
+        ),
+
+        "h_hidden": trial.suggest_categorical(
+            "h_hidden",
+            [32, 64, 128, 256],
+        ),
+
+        "activation": trial.suggest_categorical(
+            "activation",
+            ["gelu", "relu", "silu"],
+        ),
+
+        "widths": tuple(
+            [hidden_width] * depth
+        ),
+
+        "lambda_sparse": trial.suggest_float(
+            "lambda_sparse",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_div": trial.suggest_float(
+            "lambda_div",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_orth": trial.suggest_float(
+            "lambda_orth",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_budget": trial.suggest_float(
+            "lambda_budget",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_drift": trial.suggest_categorical(
+            "lambda_drift",
+            [0.0, 1e-5, 1e-4, 1e-3],
+        ),
+
+        "lr": trial.suggest_float(
+            "lr",
+            1e-4,
+            2e-2,
+            log=True,
+        ),
+
+        "weight_decay": trial.suggest_float(
+            "weight_decay",
+            1e-6,
+            3e-3,
+            log=True,
+        ),
+
+        "batch_size": trial.suggest_categorical(
+            "batch_size",
+            [32, 64, 128],
+        ),
+
+        "epochs": trial.suggest_categorical(
+            "epochs",
+            [50, 80, 120],
+        ),
+    }
+
+
+# ============================================================
+# Build AugTab
+# ============================================================
+
+def build_model(params, d_features):
+
+    regs = RegularizerConfig(
+        lambda_sparse=params["lambda_sparse"],
+        lambda_div=params["lambda_div"],
+        lambda_orth=params["lambda_orth"],
+        lambda_budget=params["lambda_budget"],
+        lambda_drift=params["lambda_drift"],
+    )
+
+    return AugTabClassifier(
+        d_features=d_features,
+        k_aug=params["k_aug"],
+        kprime=params["kprime"],
+        h_hidden=params["h_hidden"],
+        widths=params["widths"],
+        activation=params["activation"],
+        append_mask=False,
+        gating="basic",
+        regs=regs,
+        device=DEVICE,
+        lr=params["lr"],
+        weight_decay=params["weight_decay"],
+    )
+
+
+# ============================================================
+# 5-fold Optuna objective
+# ============================================================
+
+def objective(trial):
+
+    params = sample_params(trial)
+
+    cv = StratifiedKFold(
+        n_splits=5,
+        shuffle=True,
+        random_state=SEED,
+    )
+
+    scores = []
+
+    for fold, (train_idx, val_idx) in enumerate(
+        cv.split(X_train, y_train)
+    ):
+
+        # ----------------------------------------------------
+        # Fold-local preprocessing
+        # ----------------------------------------------------
+
+        scaler = StandardScaler()
+
+        X_tr = scaler.fit_transform(
+            X_train[train_idx]
+        ).astype(np.float32)
+
+        X_val = scaler.transform(
+            X_train[val_idx]
+        ).astype(np.float32)
+
+        y_tr = y_train[train_idx]
+        y_val = y_train[val_idx]
+
+
+        # ----------------------------------------------------
+        # Fresh model for every fold
+        # ----------------------------------------------------
+
+        torch.manual_seed(SEED + fold)
+
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(SEED + fold)
+
+        model = build_model(
+            params,
+            d_features=X_tr.shape[1],
+        )
+
+
+        # ----------------------------------------------------
+        # Train
+        # ----------------------------------------------------
+
+        model.fit(
+            X_tr,
+            y_tr,
+            epochs=params["epochs"],
+            batch_size=params["batch_size"],
+            verbose=False,
+        )
+
+
+        # ----------------------------------------------------
+        # Validation accuracy
+        # ----------------------------------------------------
+
+        score = model.score(
+            X_val,
+            y_val,
+        )
+
+        scores.append(score)
+
+
+    mean_score = float(np.mean(scores))
+    std_score = float(np.std(scores))
+
+    trial.set_user_attr(
+        "cv_mean",
+        mean_score,
+    )
+
+    trial.set_user_attr(
+        "cv_std",
+        std_score,
+    )
+
+    return mean_score
+
+
+# ============================================================
+# Run Optuna
+# ============================================================
+
+sampler = optuna.samplers.TPESampler(
+    seed=SEED
+)
+
+study = optuna.create_study(
+    direction="maximize",
+    sampler=sampler,
+)
+
+
+# Small value for demonstration.
+# Increase for full experiments.
+N_TRIALS = 20
+
+study.optimize(
+    objective,
+    n_trials=N_TRIALS,
+    show_progress_bar=True,
+)
+
+
+print("\nBest 5-fold CV Accuracy:")
+print(
+    f"{study.best_trial.user_attrs['cv_mean']:.4f} "
+    f"± {study.best_trial.user_attrs['cv_std']:.4f}"
+)
+
+print("\nBest Hyperparameters:")
+print(study.best_params)
+
+
+# ============================================================
+# Reconstruct widths from Optuna parameters
+# ============================================================
+
+best = study.best_params.copy()
+
+best["widths"] = tuple(
+    [best["hidden_width"]] * best["depth"]
+)
+
+
+# ============================================================
+# Train final model using all training data
+# ============================================================
+
+final_scaler = StandardScaler()
+
+X_train_final = final_scaler.fit_transform(
+    X_train
+).astype(np.float32)
+
+X_test_final = final_scaler.transform(
+    X_test
+).astype(np.float32)
+
+
+final_model = build_model(
+    best,
+    d_features=X_train_final.shape[1],
+)
+
+
+final_model.fit(
+    X_train_final,
+    y_train,
+    epochs=best["epochs"],
+    batch_size=best["batch_size"],
+    verbose=False,
+)
+
+
+# ============================================================
+# Final held-out test evaluation
+# ============================================================
+
+y_pred = final_model.predict(
+    X_test_final
+).numpy()
+
+y_prob = final_model.predict_proba(
+    X_test_final
+).numpy().reshape(-1)
+
+
+print("\nFinal Test Results")
+print(
+    "Accuracy :",
+    accuracy_score(y_test, y_pred),
+)
+
+print(
+    "F1       :",
+    f1_score(y_test, y_pred),
+)
+
+print(
+    "ROC-AUC  :",
+    roc_auc_score(y_test, y_prob),
+)
+```
+
+> **Recommended:** Increase `N_TRIALS` for full experiments. The value above is intentionally kept small so that the example remains practical as a quick-start demonstration.
+
+---
+
+### Example 3: Multiclass Classification Without Hyperparameter Tuning
+
+For multiclass problems, use `AugTabMulti` and provide the number of target classes through `n_classes`.
+
+```python
+import random
+import numpy as np
+import torch
+
+from sklearn.datasets import make_classification
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+from sklearn.metrics import accuracy_score, f1_score
+
+from augtab import AugTabMulti, RegularizerConfig
+
+
+# ============================================================
+# Reproducibility
+# ============================================================
+
+SEED = 42
+
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(SEED)
+
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# ============================================================
+# Example multiclass dataset
+# ============================================================
+
+X, y = make_classification(
+    n_samples=900,
+    n_features=15,
+    n_informative=10,
+    n_redundant=2,
+    n_classes=3,
+    n_clusters_per_class=1,
+    random_state=SEED,
+)
+
+X = X.astype(np.float32)
+y = y.astype(np.int64)
+
+N_CLASSES = len(np.unique(y))
+
+
+# ============================================================
+# Train / test split
+# ============================================================
+
+X_train, X_test, y_train, y_test = train_test_split(
+    X,
+    y,
+    test_size=0.20,
+    random_state=SEED,
+    stratify=y,
+)
+
+
+# ============================================================
+# Standardization
+# ============================================================
+
+scaler = StandardScaler()
+
+X_train = scaler.fit_transform(
+    X_train
+).astype(np.float32)
+
+X_test = scaler.transform(
+    X_test
+).astype(np.float32)
+
+
+# ============================================================
+# Regularization
+# ============================================================
+
+regs = RegularizerConfig(
+    lambda_sparse=1e-3,
+    lambda_div=1e-3,
+    lambda_orth=1e-3,
+    lambda_budget=1e-3,
+    lambda_drift=0.0,
+)
+
+
+# ============================================================
+# Initialize AugTab
+# ============================================================
+
+model = AugTabMulti(
+    d_features=X_train.shape[1],
+    n_classes=N_CLASSES,
+    k_aug=32,
+    kprime=64,
+    h_hidden=64,
+    widths=(128, 128),
+    activation="gelu",
+    append_mask=False,
+    gating="basic",
+    regs=regs,
+    device=DEVICE,
+    lr=2e-3,
+    weight_decay=1e-4,
+)
+
+
+# ============================================================
+# Train
+# ============================================================
+
+model.fit(
+    X_train,
+    y_train,
+    epochs=80,
+    batch_size=64,
+    verbose=False,
+)
+
+
+# ============================================================
+# Evaluate
+# ============================================================
+
+y_pred = model.predict(
+    X_test
+).numpy()
+
+y_prob = model.predict_proba(
+    X_test
+).numpy()
+
+
+print("\nMulticlass Classification Results")
+
+print(
+    "Accuracy :",
+    accuracy_score(y_test, y_pred),
+)
+
+print(
+    "Macro F1 :",
+    f1_score(
+        y_test,
+        y_pred,
+        average="macro",
+    ),
+)
+
+print(
+    "Probability matrix shape:",
+    y_prob.shape,
+)
+```
+
+---
+
+### Example 4: Multiclass Classification with Optuna Hyperparameter Tuning
+
+For a **new multiclass dataset**, this is the recommended workflow.
+
+```python
+import random
+import numpy as np
+import optuna
+import torch
+
+from sklearn.datasets import make_classification
+from sklearn.model_selection import train_test_split, StratifiedKFold
+from sklearn.preprocessing import StandardScaler
+from sklearn.metrics import accuracy_score, f1_score
+
+from augtab import AugTabMulti, RegularizerConfig
+
+
+# ============================================================
+# Reproducibility
+# ============================================================
+
+SEED = 42
+
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(SEED)
+
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# ============================================================
+# Example multiclass dataset
+# ============================================================
+
+X, y = make_classification(
+    n_samples=900,
+    n_features=15,
+    n_informative=10,
+    n_redundant=2,
+    n_classes=3,
+    n_clusters_per_class=1,
+    random_state=SEED,
+)
+
+X = X.astype(np.float32)
+y = y.astype(np.int64)
+
+N_CLASSES = len(np.unique(y))
+
+
+# ============================================================
+# Held-out test set
+# ============================================================
+
+X_train, X_test, y_train, y_test = train_test_split(
+    X,
+    y,
+    test_size=0.20,
+    random_state=SEED,
+    stratify=y,
+)
+
+
+# ============================================================
+# Optuna search space
+# ============================================================
+
+def sample_params(trial):
+
+    hidden_width = trial.suggest_categorical(
+        "hidden_width",
+        [64, 128, 256],
+    )
+
+    depth = trial.suggest_int(
+        "depth",
+        1,
+        3,
+    )
+
+    return {
+        "k_aug": trial.suggest_categorical(
+            "k_aug",
+            [8, 16, 24, 32, 48, 64],
+        ),
+
+        "kprime": trial.suggest_categorical(
+            "kprime",
+            [16, 32, 64, 128],
+        ),
+
+        "h_hidden": trial.suggest_categorical(
+            "h_hidden",
+            [32, 64, 128, 256],
+        ),
+
+        "activation": trial.suggest_categorical(
+            "activation",
+            ["gelu", "relu", "silu"],
+        ),
+
+        "widths": tuple(
+            [hidden_width] * depth
+        ),
+
+        "lambda_sparse": trial.suggest_float(
+            "lambda_sparse",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_div": trial.suggest_float(
+            "lambda_div",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_orth": trial.suggest_float(
+            "lambda_orth",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_budget": trial.suggest_float(
+            "lambda_budget",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_drift": trial.suggest_categorical(
+            "lambda_drift",
+            [0.0, 1e-5, 1e-4, 1e-3],
+        ),
+
+        "lr": trial.suggest_float(
+            "lr",
+            1e-4,
+            2e-2,
+            log=True,
+        ),
+
+        "weight_decay": trial.suggest_float(
+            "weight_decay",
+            1e-6,
+            3e-3,
+            log=True,
+        ),
+
+        "batch_size": trial.suggest_categorical(
+            "batch_size",
+            [32, 64, 128],
+        ),
+
+        "epochs": trial.suggest_categorical(
+            "epochs",
+            [50, 80, 120],
+        ),
+    }
+
+
+# ============================================================
+# Build model
+# ============================================================
+
+def build_model(params, d_features):
+
+    regs = RegularizerConfig(
+        lambda_sparse=params["lambda_sparse"],
+        lambda_div=params["lambda_div"],
+        lambda_orth=params["lambda_orth"],
+        lambda_budget=params["lambda_budget"],
+        lambda_drift=params["lambda_drift"],
+    )
+
+    return AugTabMulti(
+        d_features=d_features,
+        n_classes=N_CLASSES,
+        k_aug=params["k_aug"],
+        kprime=params["kprime"],
+        h_hidden=params["h_hidden"],
+        widths=params["widths"],
+        activation=params["activation"],
+        append_mask=False,
+        gating="basic",
+        regs=regs,
+        device=DEVICE,
+        lr=params["lr"],
+        weight_decay=params["weight_decay"],
+    )
+
+
+# ============================================================
+# 5-fold Optuna objective
+# ============================================================
+
+def objective(trial):
+
+    params = sample_params(trial)
+
+    cv = StratifiedKFold(
+        n_splits=5,
+        shuffle=True,
+        random_state=SEED,
+    )
+
+    scores = []
+
+    for fold, (train_idx, val_idx) in enumerate(
+        cv.split(X_train, y_train)
+    ):
+
+        scaler = StandardScaler()
+
+        X_tr = scaler.fit_transform(
+            X_train[train_idx]
+        ).astype(np.float32)
+
+        X_val = scaler.transform(
+            X_train[val_idx]
+        ).astype(np.float32)
+
+        y_tr = y_train[train_idx]
+        y_val = y_train[val_idx]
+
+
+        torch.manual_seed(SEED + fold)
+
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(SEED + fold)
+
+
+        model = build_model(
+            params,
+            d_features=X_tr.shape[1],
+        )
+
+
+        model.fit(
+            X_tr,
+            y_tr,
+            epochs=params["epochs"],
+            batch_size=params["batch_size"],
+            verbose=False,
+        )
+
+
+        score = model.score(
+            X_val,
+            y_val,
+        )
+
+        scores.append(score)
+
+
+    mean_score = float(np.mean(scores))
+    std_score = float(np.std(scores))
+
+    trial.set_user_attr(
+        "cv_mean",
+        mean_score,
+    )
+
+    trial.set_user_attr(
+        "cv_std",
+        std_score,
+    )
+
+    return mean_score
+
+
+# ============================================================
+# Run Optuna
+# ============================================================
+
+study = optuna.create_study(
+    direction="maximize",
+    sampler=optuna.samplers.TPESampler(
+        seed=SEED
+    ),
+)
+
+
+N_TRIALS = 20
+
+study.optimize(
+    objective,
+    n_trials=N_TRIALS,
+    show_progress_bar=True,
+)
+
+
+print("\nBest 5-fold CV Accuracy:")
+
+print(
+    f"{study.best_trial.user_attrs['cv_mean']:.4f} "
+    f"± {study.best_trial.user_attrs['cv_std']:.4f}"
+)
+
+print("\nBest Hyperparameters:")
+print(study.best_params)
+
+
+# ============================================================
+# Final training
+# ============================================================
+
+best = study.best_params.copy()
+
+best["widths"] = tuple(
+    [best["hidden_width"]] * best["depth"]
+)
+
+
+scaler = StandardScaler()
+
+X_train_final = scaler.fit_transform(
+    X_train
+).astype(np.float32)
+
+X_test_final = scaler.transform(
+    X_test
+).astype(np.float32)
+
+
+final_model = build_model(
+    best,
+    d_features=X_train_final.shape[1],
+)
+
+
+final_model.fit(
+    X_train_final,
+    y_train,
+    epochs=best["epochs"],
+    batch_size=best["batch_size"],
+    verbose=False,
+)
+
+
+# ============================================================
+# Final test evaluation
+# ============================================================
+
+y_pred = final_model.predict(
+    X_test_final
+).numpy()
+
+
+print("\nFinal Test Results")
+
+print(
+    "Accuracy :",
+    accuracy_score(y_test, y_pred),
+)
+
+print(
+    "Macro F1 :",
+    f1_score(
+        y_test,
+        y_pred,
+        average="macro",
+    ),
+)
+```
+
+---
+
+### Example 5: Regression with Optuna Hyperparameter Tuning
+
+For a **new regression dataset**, we recommend tuning `AugTabRegressor` with Optuna and **5-fold cross-validation**.
+
+The optimization objective below maximizes mean cross-validation $R^2$.
+
+```python
+import random
+import numpy as np
+import optuna
+import torch
+
+from sklearn.datasets import make_regression
+from sklearn.model_selection import train_test_split, KFold
+from sklearn.preprocessing import StandardScaler
+from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
+
+from augtab import AugTabRegressor, RegularizerConfig
+
+
+# ============================================================
+# Reproducibility
+# ============================================================
+
+SEED = 42
+
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(SEED)
+
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# ============================================================
+# Example regression dataset
+# ============================================================
+
+X, y = make_regression(
+    n_samples=800,
+    n_features=12,
+    n_informative=8,
+    noise=15.0,
+    random_state=SEED,
+)
+
+X = X.astype(np.float32)
+y = y.astype(np.float32)
+
+
+# ============================================================
+# Held-out test set
+# ============================================================
+
+X_train, X_test, y_train, y_test = train_test_split(
+    X,
+    y,
+    test_size=0.20,
+    random_state=SEED,
+)
+
+
+# ============================================================
+# Optuna search space
+# ============================================================
+
+def sample_params(trial):
+
+    hidden_width = trial.suggest_categorical(
+        "hidden_width",
+        [64, 128, 256],
+    )
+
+    depth = trial.suggest_int(
+        "depth",
+        1,
+        3,
+    )
+
+    return {
+        "k_aug": trial.suggest_categorical(
+            "k_aug",
+            [8, 16, 24, 32, 48, 64],
+        ),
+
+        "kprime": trial.suggest_categorical(
+            "kprime",
+            [16, 32, 64, 128],
+        ),
+
+        "h_hidden": trial.suggest_categorical(
+            "h_hidden",
+            [32, 64, 128, 256],
+        ),
+
+        "activation": trial.suggest_categorical(
+            "activation",
+            ["gelu", "relu", "silu"],
+        ),
+
+        "widths": tuple(
+            [hidden_width] * depth
+        ),
+
+        "lambda_sparse": trial.suggest_float(
+            "lambda_sparse",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_div": trial.suggest_float(
+            "lambda_div",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_orth": trial.suggest_float(
+            "lambda_orth",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_budget": trial.suggest_float(
+            "lambda_budget",
+            1e-6,
+            1e-2,
+            log=True,
+        ),
+
+        "lambda_drift": trial.suggest_categorical(
+            "lambda_drift",
+            [0.0, 1e-5, 1e-4, 1e-3],
+        ),
+
+        "lr": trial.suggest_float(
+            "lr",
+            1e-4,
+            2e-2,
+            log=True,
+        ),
+
+        "weight_decay": trial.suggest_float(
+            "weight_decay",
+            1e-6,
+            3e-3,
+            log=True,
+        ),
+
+        "batch_size": trial.suggest_categorical(
+            "batch_size",
+            [32, 64, 128],
+        ),
+
+        "epochs": trial.suggest_categorical(
+            "epochs",
+            [50, 80, 120],
+        ),
+    }
+
+
+# ============================================================
+# Build AugTab regressor
+# ============================================================
+
+def build_model(params, d_features):
+
+    regs = RegularizerConfig(
+        lambda_sparse=params["lambda_sparse"],
+        lambda_div=params["lambda_div"],
+        lambda_orth=params["lambda_orth"],
+        lambda_budget=params["lambda_budget"],
+        lambda_drift=params["lambda_drift"],
+    )
+
+    return AugTabRegressor(
+        d_features=d_features,
+        k_aug=params["k_aug"],
+        kprime=params["kprime"],
+        h_hidden=params["h_hidden"],
+        widths=params["widths"],
+        activation=params["activation"],
+        append_mask=False,
+        gating="basic",
+        regs=regs,
+        device=DEVICE,
+        lr=params["lr"],
+        weight_decay=params["weight_decay"],
+    )
+
+
+# ============================================================
+# 5-fold Optuna objective
+# ============================================================
+
+def objective(trial):
+
+    params = sample_params(trial)
+
+    cv = KFold(
+        n_splits=5,
+        shuffle=True,
+        random_state=SEED,
+    )
+
+    scores = []
+
+    for fold, (train_idx, val_idx) in enumerate(
+        cv.split(X_train)
+    ):
+
+        scaler = StandardScaler()
+
+        X_tr = scaler.fit_transform(
+            X_train[train_idx]
+        ).astype(np.float32)
+
+        X_val = scaler.transform(
+            X_train[val_idx]
+        ).astype(np.float32)
+
+        y_tr = y_train[train_idx]
+        y_val = y_train[val_idx]
+
+
+        torch.manual_seed(SEED + fold)
+
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(SEED + fold)
+
+
+        model = build_model(
+            params,
+            d_features=X_tr.shape[1],
+        )
+
+
+        model.fit(
+            X_tr,
+            y_tr,
+            epochs=params["epochs"],
+            batch_size=params["batch_size"],
+            verbose=False,
+        )
+
+
+        score = model.score(
+            X_val,
+            y_val,
+        )
+
+        scores.append(score)
+
+
+    mean_score = float(np.mean(scores))
+    std_score = float(np.std(scores))
+
+    trial.set_user_attr(
+        "cv_mean",
+        mean_score,
+    )
+
+    trial.set_user_attr(
+        "cv_std",
+        std_score,
+    )
+
+    return mean_score
+
+
+# ============================================================
+# Run Optuna
+# ============================================================
+
+study = optuna.create_study(
+    direction="maximize",
+    sampler=optuna.samplers.TPESampler(
+        seed=SEED
+    ),
+)
+
+
+N_TRIALS = 20
+
+study.optimize(
+    objective,
+    n_trials=N_TRIALS,
+    show_progress_bar=True,
+)
+
+
+print("\nBest 5-fold CV R2:")
+
+print(
+    f"{study.best_trial.user_attrs['cv_mean']:.4f} "
+    f"± {study.best_trial.user_attrs['cv_std']:.4f}"
+)
+
+print("\nBest Hyperparameters:")
+print(study.best_params)
+
+
+# ============================================================
+# Final training
+# ============================================================
+
+best = study.best_params.copy()
+
+best["widths"] = tuple(
+    [best["hidden_width"]] * best["depth"]
+)
+
+
+scaler = StandardScaler()
+
+X_train_final = scaler.fit_transform(
+    X_train
+).astype(np.float32)
+
+X_test_final = scaler.transform(
+    X_test
+).astype(np.float32)
+
+
+final_model = build_model(
+    best,
+    d_features=X_train_final.shape[1],
+)
+
+
+final_model.fit(
+    X_train_final,
+    y_train,
+    epochs=best["epochs"],
+    batch_size=best["batch_size"],
+    verbose=False,
+)
+
+
+# ============================================================
+# Final held-out test evaluation
+# ============================================================
+
+y_pred = final_model.predict(
+    X_test_final
+).numpy().reshape(-1)
+
+
+rmse = np.sqrt(
+    mean_squared_error(
+        y_test,
+        y_pred,
+    )
+)
+
+mae = mean_absolute_error(
+    y_test,
+    y_pred,
+)
+
+r2 = r2_score(
+    y_test,
+    y_pred,
+)
+
+
+print("\nFinal Test Results")
+print("RMSE :", rmse)
+print("MAE  :", mae)
+print("R2   :", r2)
+```
+
+---
+
+### Example 6: Regression Without Hyperparameter Tuning
+
+This example trains `AugTabRegressor` using a fixed configuration without Optuna.
+
+```python
+import random
+import numpy as np
+import torch
+
+from sklearn.datasets import make_regression
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
+
+from augtab import AugTabRegressor, RegularizerConfig
+
+
+# ============================================================
+# Reproducibility
+# ============================================================
+
+SEED = 42
+
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(SEED)
+
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# ============================================================
+# Example regression dataset
+# ============================================================
+
+X, y = make_regression(
+    n_samples=800,
+    n_features=12,
+    n_informative=8,
+    noise=15.0,
+    random_state=SEED,
+)
+
+X = X.astype(np.float32)
+y = y.astype(np.float32)
+
+
+# ============================================================
+# Train / test split
+# ============================================================
+
+X_train, X_test, y_train, y_test = train_test_split(
+    X,
+    y,
+    test_size=0.20,
+    random_state=SEED,
+)
+
+
+# ============================================================
+# Standardization
+# ============================================================
+
+scaler = StandardScaler()
+
+X_train = scaler.fit_transform(
+    X_train
+).astype(np.float32)
+
+X_test = scaler.transform(
+    X_test
+).astype(np.float32)
+
+
+# ============================================================
+# Regularization
+# ============================================================
+
+regs = RegularizerConfig(
+    lambda_sparse=1e-3,
+    lambda_div=1e-3,
+    lambda_orth=1e-3,
+    lambda_budget=1e-3,
+    lambda_drift=0.0,
+)
+
+
+# ============================================================
+# Initialize AugTab
+# ============================================================
+
+model = AugTabRegressor(
+    d_features=X_train.shape[1],
+    k_aug=32,
+    kprime=64,
+    h_hidden=64,
+    widths=(128, 128),
+    activation="gelu",
+    append_mask=False,
+    gating="basic",
+    regs=regs,
+    device=DEVICE,
+    lr=2e-3,
+    weight_decay=1e-4,
+)
+
+
+# ============================================================
+# Train
+# ============================================================
+
+model.fit(
+    X_train,
+    y_train,
+    epochs=80,
+    batch_size=64,
+    verbose=False,
+)
+
+
+# ============================================================
+# Evaluate
+# ============================================================
+
+y_pred = model.predict(
+    X_test
+).numpy().reshape(-1)
+
+
+rmse = np.sqrt(
+    mean_squared_error(
+        y_test,
+        y_pred,
+    )
+)
+
+mae = mean_absolute_error(
+    y_test,
+    y_pred,
+)
+
+r2 = r2_score(
+    y_test,
+    y_pred,
+)
+
+
+print("\nRegression Results")
+print("RMSE :", rmse)
+print("MAE  :", mae)
+print("R2   :", r2)
+```
+
+---
+
+### Adapting AugTab to Your Own Dataset
+
+To apply AugTab to a new low-dimensional tabular dataset:
+
+1. Replace the synthetic `X` and `y` arrays with your own feature matrix and target.
+2. Encode categorical variables before passing them to AugTab.
+3. Encode classification labels as integer class IDs.
+4. Standardize numerical features using statistics computed from the training split only.
+5. Set `d_features=X.shape[1]` after preprocessing.
+6. Use `AugTabClassifier` for binary classification.
+7. Use `AugTabMulti` for multiclass classification and specify `n_classes`.
+8. Use `AugTabRegressor` for regression.
+9. Use `append_mask=False` when no explicit missingness mask is required.
+10. For a new dataset, prefer the **Optuna workflow** to select AugTab and optimization hyperparameters.
+11. During Optuna tuning, perform preprocessing independently inside each cross-validation fold.
+12. Initialize a **fresh AugTab model for every cross-validation fold**.
+13. Keep the test set completely separate from hyperparameter tuning and use it only for final evaluation.
+
+> **Recommended for new datasets:** The fixed-configuration examples are useful for quick experiments and installation checks, while the **Optuna + 5-fold cross-validation workflows** are recommended when reporting dataset-specific AugTab results.
+
+## Related Work and Project Context
+
+AugTab is part of my broader work on tabular deep learning, but it was developed as a **separate collaborative project outside my PhD dissertation research**. Like **ZAYAN**, it explores a complementary direction beyond my primary PhD research on feature ordering, sequencing, and high-dimensional tabular learning.
+
+For broader context, some of our related tabular and multimodal learning projects are listed below.
+
+### GOTabPFN - ICML 2026
+
+**GOTabPFN: From Feature Ordering to Compact Tokenization for Tabular Foundation Models on High-Dimensional Data**
+
+- **Venue:** International Conference on Machine Learning (ICML 2026)
+- **GitHub:** https://github.com/zadid6pretam/GOTabPFN
+
+### iSyncTab - ECCV 2026
+
+**iSyncTab: Learning Cross-Modal Feature Sequencing for Image-Tabular Data via Neural Synchrony**
+
+- **Venue:** European Conference on Computer Vision (ECCV 2026)
+- **GitHub:** https://github.com/zadid6pretam/iSyncTab
+
+### BSTabDiff - ICLR 2026 DeLTa Workshop
+
+**BSTabDiff: Block-Subunit Diffusion Priors for High-Dimensional Tabular Data Generation**
+
+- **Venue:** ICLR 2026 Workshop on Deep Generative Models in Machine Learning: Theory, Principle and Efficacy (DeLTa)
+- **GitHub:** https://github.com/zadid6pretam/BSTabDiff
+
+### iStructTab - ICPR 2026
+
+**iStructTab: Structured Feature Sequencing for Multimodal Learning of Image and Tabular Data**
+
+- **Venue:** International Conference on Pattern Recognition (ICPR 2026)
+- **GitHub:** https://github.com/zadid6pretam/iStructTab
+
+### DynaTab - AAAI 2026 NeurAI Workshop
+
+**DynaTab: Dynamic Feature Ordering as Neural Rewiring for High-Dimensional Tabular Data**
+
+- **Venue:** AAAI 2026 Workshop on NeuroAI Multimodal Intelligence
+- **GitHub:** https://github.com/zadid6pretam/DynaTab
+
+### TabSeq - ICPR 2024
+
+**TabSeq: A Framework for Deep Learning on Tabular Data via Sequential Ordering**
+
+- **Venue:** International Conference on Pattern Recognition (ICPR 2024)
+- **GitHub:** https://github.com/zadid6pretam/TabSeq
+
+---
+
+### ZAYAN - ICPR 2026
+
+**ZAYAN: Disentangled Contrastive Transformer for Tabular Remote Sensing Data**
+
+- **Venue:** International Conference on Pattern Recognition (ICPR 2026)
+- **GitHub:** https://github.com/zadid6pretam/ZAYAN
+
+> **Note:** AugTab and ZAYAN were developed as separate collaborative projects outside my PhD dissertation research.
+
+## Contact
+
+For any questions, issues, or suggestions related to AugTab, please feel free to open an issue on GitHub.