TabKD: What Happens When You Teach a Tiny Model to Think Like XGBoost — Without Seeing Any Real Data

Complete End-to-End TabKD Implementation (PyTorch)

The implementation below is a complete, runnable PyTorch implementation of TabKD, structured across 10 sections that map directly to the paper. It covers dynamic soft bin learning with boundary-focused sampling, the interaction diversity loss with entropy maximization over pairwise joint bin distributions, the hardness loss via KL divergence, the three-phase training loop (warmup → bin stabilization → adversarial distillation), replay buffer management, all four teacher wrappers (Neural Network, XGBoost, Random Forest, TabTransformer), dataset helpers for all four benchmark datasets, evaluation metrics including coverage quantification, and a smoke test that validates all components without requiring real data.

# ==============================================================================
# TabKD: Tabular Knowledge Distillation through Interaction Diversity
#        of Learned Feature Bins
# Paper: arXiv:2603.15481v1 [cs.LG] (2026)
# Authors: Shovon Niverd Pereira, Krishna Khadka, Yu Lei
# Affiliation: University of Texas at Arlington
# ==============================================================================
# Sections:
#   1.  Imports & Configuration
#   2.  Dynamic Soft Bin Learner (decision-boundary-aligned discretization)
#   3.  Interaction Diversity Loss (pairwise bin entropy maximization)
#   4.  Generator Network (noise → synthetic tabular samples)
#   5.  Student Network (lightweight MLP approximation)
#   6.  Teacher Wrappers (NN, XGBoost, Random Forest, TabTransformer)
#   7.  Replay Buffer (catastrophic forgetting prevention)
#   8.  TabKD Trainer (three-phase: warmup → bins → adversarial)
#   9.  Dataset Helpers & Evaluation Metrics
#  10.  Smoke Test
# ==============================================================================

from __future__ import annotations

import math
import random
import warnings
from collections import deque
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.utils.data import DataLoader, TensorDataset

warnings.filterwarnings("ignore")


# ─── SECTION 1: Configuration ─────────────────────────────────────────────────

@dataclass
class TabKDConfig:
    """
    Hyperparameter configuration for TabKD (Section 5.7).

    Attributes
    ----------
    n_features      : int   — number of input features (F)
    n_classes       : int   — number of output classes (C)
    n_bins          : int   — bins per feature K (paper: 8)
    noise_dim       : int   — generator input noise dimension
    gen_hidden      : int   — generator hidden layer size
    student_hidden  : int   — student MLP hidden units (paper: 32)
    batch_size      : int   — training batch size (paper: 128)
    lr              : float — Adam learning rate (paper: 0.001)
    lambda_cov      : float — diversity loss weight (paper: 10.0)
    lambda_hard     : float — hardness loss weight (paper: 2.0)
    lambda_intra    : float — intra-bin variance weight
    lambda_inter    : float — inter-bin inverse-variance weight
    lambda_div      : float — class diversity weight (bin phase generator)
    lambda_boundary : float — boundary entropy weight (bin phase generator)
    warmup_epochs   : int   — Stage 0 warmup epochs (paper: 30)
    bin_epochs      : int   — Stage 1 bin learning epochs (paper: 200)
    adv_epochs      : int   — Stage 2 adversarial epochs (paper: 400)
    tau_start       : float — temperature annealing start
    tau_end         : float — temperature annealing end
    replay_ratio    : float — fraction of replay samples in adversarial batch (paper: 0.1)
    replay_size     : int   — maximum replay buffer size
    query_budget    : int   — total synthetic query budget (paper: 9600)
    """
    n_features: int = 14
    n_classes: int = 2
    n_bins: int = 8
    noise_dim: int = 64
    gen_hidden: int = 128
    student_hidden: int = 32
    batch_size: int = 128
    lr: float = 1e-3
    lambda_cov: float = 10.0
    lambda_hard: float = 2.0
    lambda_intra: float = 1.0
    lambda_inter: float = 1.0
    lambda_div: float = 1.0
    lambda_boundary: float = 1.0
    warmup_epochs: int = 30
    bin_epochs: int = 200
    adv_epochs: int = 400
    tau_start: float = 1.0
    tau_end: float = 0.05
    replay_ratio: float = 0.1
    replay_size: int = 2000
    query_budget: int = 9600


# ─── SECTION 2: Dynamic Soft Bin Learner ──────────────────────────────────────

class SoftBinLearner(nn.Module):
    """
    Learns K adaptive bin boundaries per feature, aligned with teacher
    decision boundaries (Section 4.2).

    Instead of hard discrete assignments (which are non-differentiable),
    uses a temperature-annealed softmax over learned bin center offsets
    to produce soft membership vectors m ∈ [0,1]^K for each feature.

    Bin Learning Objective (Eq. 1):
        L_bin = λ_intra · Var_intra(M, P_T)
              + λ_inter · (1 / Var_inter(M, P_T))

    Intra-bin variance: for each bin, the variance of teacher predictions
    weighted by soft memberships. Low = homogeneous bins (desired).
    Inter-bin variance: variance across bin-mean predictions.
    High = bins capture different behaviors (desired).

    Parameters
    ----------
    n_features : int — number of input features
    n_bins     : int — number of bins per feature K
    tau        : float — initial temperature for soft assignment
    """

    def __init__(self, n_features: int, n_bins: int, tau: float = 1.0):
        super().__init__()
        self.n_features = n_features
        self.n_bins = n_bins
        self.tau = tau

        # Learnable bin centers: (F, K)
        # Initialized to spread uniformly in [-2, 2] (standardized feature space)
        centers_init = torch.linspace(-2.0, 2.0, n_bins).unsqueeze(0).expand(n_features, -1).clone()
        self.bin_centers = nn.Parameter(centers_init)

    def set_temperature(self, tau: float):
        """Update soft assignment temperature (called during annealing)."""
        self.tau = tau

    def forward(self, x: Tensor) -> Tensor:
        """
        Compute soft bin membership for each sample.

        Parameters
        ----------
        x : (N, F) — input samples (standardized)

        Returns
        -------
        m : (N, F, K) — soft membership vectors
            m[n, f, k] = probability that sample n, feature f belongs to bin k
        """
        N = x.shape[0]
        # x: (N, F) → (N, F, 1)
        x_expanded = x.unsqueeze(-1)
        # bin_centers: (F, K) → (1, F, K)
        centers = self.bin_centers.unsqueeze(0)
        # Negative squared distance to each bin center: (N, F, K)
        neg_dist = -((x_expanded - centers) ** 2) / (self.tau + 1e-8)
        # Soft membership via temperature-scaled softmax over bins
        m = F.softmax(neg_dist, dim=-1)  # (N, F, K)
        return m

    def hard_assignment(self, x: Tensor) -> Tensor:
        """
        Hard bin assignment for coverage measurement at evaluation time.
        Returns (N, F) integer bin indices.
        """
        with torch.no_grad():
            m = self.forward(x)
            return m.argmax(dim=-1)  # (N, F)


class BinLearningLoss(nn.Module):
    """
    Computes the bin learning objective (Eq. 1 in paper):
        L_bin = λ_intra · Var_intra(M, P_T) + λ_inter · (1 / Var_inter(M, P_T))

    Intra-bin variance:
        For each feature f and bin k, compute the variance of teacher
        predictions P_T weighted by soft memberships m_k:
            var_intra_fk = Σ_n m_nfk · (P_T_n − μ_fk)²
            where μ_fk = Σ_n m_nfk · P_T_n / Σ_n m_nfk
        Average over all features and bins.

    Inter-bin variance:
        Variance of bin-mean predictions across bins (per feature).
        We want bins to have different mean predictions → maximize this.
    """

    def __init__(self, config: TabKDConfig):
        super().__init__()
        self.config = config

    def forward(self, m: Tensor, teacher_probs: Tensor) -> Tensor:
        """
        Parameters
        ----------
        m            : (N, F, K) — soft bin memberships
        teacher_probs: (N, C) — teacher output probabilities

        Returns
        -------
        loss : scalar — bin learning loss
        """
        N, F, K = m.shape
        # Use the max teacher probability as the scalar prediction signal
        p = teacher_probs.max(dim=-1).values  # (N,)

        # Expand p for broadcasting: (N, 1, 1)
        p_exp = p.unsqueeze(-1).unsqueeze(-1)

        # Weighted sum of memberships: (F, K)
        m_sum = m.sum(dim=0) + 1e-8  # (F, K)

        # Bin means μ_fk = Σ_n m_nfk · p_n / Σ_n m_nfk : (F, K)
        # m: (N, F, K), p: (N,)
        bin_means = (m * p.view(N, 1, 1)).sum(dim=0) / m_sum  # (F, K)

        # Intra-bin variance: (N, F, K) → scalar
        diff_sq = (p.view(N, 1, 1) - bin_means.unsqueeze(0)) ** 2
        intra_var = (m * diff_sq).sum(dim=0) / m_sum  # (F, K)
        loss_intra = intra_var.mean()

        # Inter-bin variance: variance of bin_means across bins, per feature
        inter_var = bin_means.var(dim=-1) + 1e-8  # (F,)
        loss_inter = (1.0 / inter_var).mean()

        loss = self.config.lambda_intra * loss_intra + self.config.lambda_inter * loss_inter
        return loss


# ─── SECTION 3: Interaction Diversity Loss ────────────────────────────────────

class InteractionDiversityLoss(nn.Module):
    """
    Core TabKD contribution: maximizes entropy over pairwise bin-combination
    distributions to achieve systematic feature interaction coverage (Eq. 3–4).

    For each pair of features (i, j), computes the empirical joint distribution
    P(k1, k2 | i, j) over K² bin combinations (Eq. 3):
        P(k1, k2 | i, j) = (1/N) Σ_x m_k1^(i)(x) · m_k2^(j)(x)

    Diversity loss = negative average entropy over all feature pairs (Eq. 4):
        L_diversity = −(1/C(F,2)) Σ_{i

    def __init__(self, config: TabKDConfig):
        super().__init__()
        self.config = config
        F = config.n_features
        self.n_pairs = F * (F - 1) // 2
        # Precompute all unique feature pair indices
        self.pairs = [(i, j) for i in range(F) for j in range(i + 1, F)]

    def forward(self, m: Tensor) -> Tensor:
        """
        Parameters
        ----------
        m : (N, F, K) — soft bin memberships from the bin learner

        Returns
        -------
        loss : scalar — negative average entropy (minimize = maximize diversity)
        """
        total_entropy = 0.0
        count = 0

        for i, j in self.pairs:
            m_i = m[:, i, :]  # (N, K)
            m_j = m[:, j, :]  # (N, K)

            # Joint distribution over K² bin combinations (Eq. 3)
            # outer product per sample: (N, K, K)
            joint = m_i.unsqueeze(-1) * m_j.unsqueeze(-2)  # (N, K, K)

            # Average over batch: empirical joint distribution (K, K)
            p_joint = joint.mean(dim=0)  # (K, K)

            # Flatten to (K²,) and compute entropy
            p_flat = p_joint.reshape(-1) + 1e-10  # numerical stability
            p_flat = p_flat / p_flat.sum()         # normalize to valid distribution
            entropy = -(p_flat * p_flat.log()).sum()

            total_entropy = total_entropy + entropy
            count += 1

        # Negative entropy: minimizing this maximizes entropy (uniform coverage)
        loss = -(total_entropy / max(1, count))
        return loss

    @torch.no_grad()
    def compute_coverage(self, bin_assignments: Tensor) -> float:
        """
        Compute cumulative interaction coverage (Eq. 13) given hard bin assignments.

        Parameters
        ----------
        bin_assignments : (N, F) — integer bin indices (from SoftBinLearner.hard_assignment)

        Returns
        -------
        coverage : float in [0, 1] — fraction of pairwise bin cells visited
        """
        F = self.config.n_features
        K = self.config.n_bins
        visited = set()

        N = bin_assignments.shape[0]
        for n in range(N):
            row = bin_assignments[n]
            for i, j in self.pairs:
                k1 = row[i].item()
                k2 = row[j].item()
                visited.add((i, j, k1, k2))

        total_cells = len(self.pairs) * (K * K)
        return len(visited) / max(1, total_cells)


# ─── SECTION 4: Generator Network ─────────────────────────────────────────────

class TabularGenerator(nn.Module):
    """
    Generator G(z): maps noise z ~ N(0, I) to synthetic tabular samples.

    Architecture: 3-layer MLP with Batch Normalization and LeakyReLU activations.
    Output: F-dimensional synthetic samples in the standardized feature space.

    The generator is trained to simultaneously:
      1. Maximize interaction diversity (L_diversity) — cover all bin pairs
      2. Maximize student-teacher disagreement (L_hardness) — target weaknesses

    Parameters
    ----------
    config : TabKDConfig
    """

    def __init__(self, config: TabKDConfig):
        super().__init__()
        self.config = config
        H = config.gen_hidden
        F = config.n_features

        self.net = nn.Sequential(
            nn.Linear(config.noise_dim, H),
            nn.BatchNorm1d(H),
            nn.LeakyReLU(0.2),
            nn.Linear(H, H * 2),
            nn.BatchNorm1d(H * 2),
            nn.LeakyReLU(0.2),
            nn.Linear(H * 2, H),
            nn.BatchNorm1d(H),
            nn.LeakyReLU(0.2),
            nn.Linear(H, F),
            nn.Tanh(),  # Output in [-1, 1] ≈ standardized feature range
        )
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, z: Tensor) -> Tensor:
        """
        Parameters
        ----------
        z : (N, noise_dim) — Gaussian noise

        Returns
        -------
        x_gen : (N, F) — synthetic tabular samples
        """
        return self.net(z)

    def sample(self, n: int, device: torch.device) -> Tensor:
        """Generate n synthetic samples."""
        z = torch.randn(n, self.config.noise_dim, device=device)
        return self.forward(z)


# ─── SECTION 5: Student Network ───────────────────────────────────────────────

class TabularStudent(nn.Module):
    """
    Student S(x): single-hidden-layer MLP (32 units, ReLU) as per Section 5.3.

    Deliberately lightweight — the same architecture is used regardless of
    teacher type, making results fair and comparable across configurations.

    Training objective (Eq. 7):
        L_student = (1/N) Σ_x D_KL(S(x) ‖ T(x))
    where T(x) provides soft label targets from the teacher.

    Parameters
    ----------
    config : TabKDConfig
    """

    def __init__(self, config: TabKDConfig):
        super().__init__()
        self.config = config
        self.net = nn.Sequential(
            nn.Linear(config.n_features, config.student_hidden),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(config.student_hidden, config.n_classes),
        )

    def forward(self, x: Tensor) -> Tensor:
        """Returns log-probabilities via log_softmax."""
        return F.log_softmax(self.net(x), dim=-1)

    def predict(self, x: Tensor) -> Tensor:
        """Returns class predictions."""
        with torch.no_grad():
            return self.forward(x).argmax(dim=-1)


# ─── SECTION 6: Teacher Wrappers ──────────────────────────────────────────────

class NeuralNetworkTeacher(nn.Module):
    """
    Pre-trained neural network teacher (128 → 64 units, ReLU, dropout 0.2).
    As per Section 5.3.
    """

    def __init__(self, n_features: int, n_classes: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_features, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, n_classes),
        )

    def forward(self, x: Tensor) -> Tensor:
        """Returns class probabilities (softmax)."""
        logits = self.net(x)
        return F.softmax(logits, dim=-1)

    def predict(self, x: Tensor) -> Tensor:
        with torch.no_grad():
            return self.forward(x).argmax(dim=-1)

    def train_on_data(
        self,
        X_train: Tensor,
        y_train: Tensor,
        epochs: int = 30,
        lr: float = 0.001,
    ) -> float:
        """Quick teacher pre-training. Returns final loss."""
        optimizer = torch.optim.Adam(self.parameters(), lr=lr)
        criterion = nn.CrossEntropyLoss()
        dataset = TensorDataset(X_train, y_train)
        loader = DataLoader(dataset, batch_size=64, shuffle=True)
        final_loss = 0.0
        for epoch in range(epochs):
            for xb, yb in loader:
                optimizer.zero_grad()
                loss = criterion(self.net(xb), yb)
                loss.backward()
                optimizer.step()
                final_loss = loss.item()
        return final_loss


class XGBoostTeacher:
    """
    XGBoost teacher wrapper. Non-differentiable gradient-free ensemble.
    TabKD handles this via black-box querying (no backprop through teacher).

    Requires: pip install xgboost
    Falls back to mock implementation for smoke testing.
    """

    def __init__(self, n_estimators: int = 100, max_depth: int = 6, lr: float = 0.1):
        self.params = {"n_estimators": n_estimators, "max_depth": max_depth, "learning_rate": lr}
        self._model = None
        self._mock = False

    def fit(self, X: np.ndarray, y: np.ndarray):
        try:
            from xgboost import XGBClassifier
            self._model = XGBClassifier(**self.params, use_label_encoder=False, eval_metric="logloss")
            self._model.fit(X, y)
        except ImportError:
            print("  [XGBoostTeacher] xgboost not installed. Using mock teacher.")
            self._mock = True

    def predict_proba_tensor(self, x: Tensor) -> Tensor:
        """Returns teacher probabilities as a Tensor (N, C)."""
        if self._mock or self._model is None:
            N = x.shape[0]
            probs = torch.rand(N, 2)
            return probs / probs.sum(dim=-1, keepdim=True)
        X_np = x.detach().cpu().numpy()
        probs = self._model.predict_proba(X_np)
        return torch.tensor(probs, dtype=torch.float32, device=x.device)

    def predict(self, x: Tensor) -> Tensor:
        return self.predict_proba_tensor(x).argmax(dim=-1)


class RandomForestTeacher:
    """
    Random Forest teacher wrapper. Non-differentiable ensemble.

    Requires: pip install scikit-learn
    Falls back to mock for smoke testing.
    """

    def __init__(self, n_estimators: int = 100, max_depth: int = 10, min_samples_split: int = 5):
        self.params = {"n_estimators": n_estimators, "max_depth": max_depth,
                       "min_samples_split": min_samples_split}
        self._model = None
        self._mock = False

    def fit(self, X: np.ndarray, y: np.ndarray):
        try:
            from sklearn.ensemble import RandomForestClassifier
            self._model = RandomForestClassifier(**self.params, random_state=42)
            self._model.fit(X, y)
        except ImportError:
            print("  [RandomForestTeacher] sklearn not installed. Using mock teacher.")
            self._mock = True

    def predict_proba_tensor(self, x: Tensor) -> Tensor:
        if self._mock or self._model is None:
            N = x.shape[0]
            probs = torch.rand(N, 2)
            return probs / probs.sum(dim=-1, keepdim=True)
        X_np = x.detach().cpu().numpy()
        probs = self._model.predict_proba(X_np)
        return torch.tensor(probs, dtype=torch.float32, device=x.device)

    def predict(self, x: Tensor) -> Tensor:
        return self.predict_proba_tensor(x).argmax(dim=-1)


class TabTransformerTeacher(nn.Module):
    """
    Lightweight TabTransformer teacher (Section 5.3).
    Multi-head self-attention over feature embeddings.

    Production: use the full TabTransformer implementation:
        pip install tab-transformer-pytorch
    """

    def __init__(self, n_features: int, n_classes: int, d_model: int = 32, n_heads: int = 4, n_layers: int = 2):
        super().__init__()
        self.feature_embed = nn.Linear(1, d_model)
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=n_heads, batch_first=True, dropout=0.1)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        self.head = nn.Linear(d_model * n_features, n_classes)
        self.n_features = n_features

    def forward(self, x: Tensor) -> Tensor:
        # x: (N, F) → (N, F, 1) → embed → (N, F, d_model)
        x_emb = self.feature_embed(x.unsqueeze(-1))
        x_out = self.transformer(x_emb)
        x_flat = x_out.reshape(x.shape[0], -1)
        return F.softmax(self.head(x_flat), dim=-1)

    def predict(self, x: Tensor) -> Tensor:
        with torch.no_grad():
            return self.forward(x).argmax(dim=-1)

    def train_on_data(self, X_train: Tensor, y_train: Tensor, epochs: int = 20) -> float:
        optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
        dataset = TensorDataset(X_train, y_train)
        loader = DataLoader(dataset, batch_size=64, shuffle=True)
        final_loss = 0.0
        for _ in range(epochs):
            for xb, yb in loader:
                optimizer.zero_grad()
                logits = self.head(
                    self.transformer(self.feature_embed(xb.unsqueeze(-1))).reshape(xb.shape[0], -1)
                )
                loss = F.cross_entropy(logits, yb)
                loss.backward()
                optimizer.step()
                final_loss = loss.item()
        return final_loss


def get_teacher_probs(teacher, x: Tensor) -> Tensor:
    """Unified interface to get teacher soft label probabilities."""
    if isinstance(teacher, (NeuralNetworkTeacher, TabTransformerTeacher)):
        with torch.no_grad():
            return teacher(x)
    elif isinstance(teacher, (XGBoostTeacher, RandomForestTeacher)):
        return teacher.predict_proba_tensor(x)
    raise ValueError(f"Unknown teacher type: {type(teacher)}")


# ─── SECTION 7: Replay Buffer ─────────────────────────────────────────────────

class ReplayBuffer:
    """
    Experience replay buffer populated during warmup (Section 4.5).

    Prevents catastrophic forgetting: the student is trained on a mix of
      - 90% adversarial samples (hard, diverse, from adversarial training)
      - 10% replay samples  (typical, from warmup — stability anchors)

    This ensures the student does not overfit to edge cases during adversarial
    training while forgetting typical inputs.
    """

    def __init__(self, max_size: int = 2000):
        self.buffer = deque(maxlen=max_size)
        self.max_size = max_size

    def add(self, x: Tensor):
        """Add samples to the buffer."""
        for i in range(x.shape[0]):
            self.buffer.append(x[i].detach().cpu())

    def sample(self, n: int, device: torch.device) -> Optional[Tensor]:
        """Sample n items from the buffer. Returns None if buffer is empty."""
        if len(self.buffer) == 0:
            return None
        n = min(n, len(self.buffer))
        samples = random.sample(list(self.buffer), n)
        return torch.stack(samples, dim=0).to(device)

    def __len__(self):
        return len(self.buffer)


# ─── SECTION 8: TabKD Trainer ─────────────────────────────────────────────────

class TabKDTrainer:
    """
    Full TabKD training pipeline implementing all three phases (Section 5.5):

    Phase 0 — Warmup (30 epochs):
        Train student on random samples, populate replay buffer.
        Objective: Eq. 8 — standard KL divergence distillation.

    Phase 1 — Bin Stabilization (200 epochs, temperature annealing):
        Train bin learner to align with teacher decision boundaries.
        Run boundary-focused generator to supply near-boundary samples.
        Objective: Eq. 1 (bin loss) + Eq. 2 (boundary generator).
        Freeze bin learner upon completion.

    Phase 2 — Adversarial Distillation (400 epochs):
        Generator maximizes interaction coverage + student hardness.
        Student minimizes KL divergence from teacher.
        Mix: 90% adversarial + 10% replay buffer samples.
        Objective: Eq. 6 (generator) + Eq. 7 (student).
    """

    def __init__(self, config: TabKDConfig, teacher, device: torch.device):
        self.config = config
        self.teacher = teacher
        self.device = device

        self.bin_learner = SoftBinLearner(config.n_features, config.n_bins).to(device)
        self.generator = TabularGenerator(config).to(device)
        self.student = TabularStudent(config).to(device)
        self.replay_buffer = ReplayBuffer(config.replay_size)

        self.bin_loss_fn = BinLearningLoss(config)
        self.diversity_loss_fn = InteractionDiversityLoss(config)

        self.gen_optimizer = torch.optim.Adam(self.generator.parameters(), lr=config.lr)
        self.student_optimizer = torch.optim.Adam(self.student.parameters(), lr=config.lr)
        self.bin_optimizer = torch.optim.Adam(self.bin_learner.parameters(), lr=config.lr * 0.5)

        # Cosine annealing schedulers
        total_epochs = config.warmup_epochs + config.bin_epochs + config.adv_epochs
        self.gen_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            self.gen_optimizer, T_max=total_epochs
        )
        self.student_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            self.student_optimizer, T_max=total_epochs
        )

        self.history = []

    def _generate_batch(self, n: int) -> Tensor:
        """Generate a batch of synthetic samples."""
        return self.generator.sample(n, self.device)

    def _teacher_probs(self, x: Tensor) -> Tensor:
        """Get teacher soft labels for a batch of samples."""
        return get_teacher_probs(self.teacher, x)

    def _annealed_temperature(self, epoch: int) -> float:
        """Linear temperature annealing during bin learning (Eq. 9)."""
        tau = self.config.tau_start - (
            (self.config.tau_start - self.config.tau_end) * epoch / max(1, self.config.bin_epochs)
        )
        return max(self.config.tau_end, tau)

    def _warmup(self):
        """
        Phase 0: Pre-train student on random samples.
        Populates replay buffer for later catastrophic forgetting prevention.
        (Eq. 8: L_warmup = KL(S(x) ‖ T(x)))
        """
        print(f"  [Phase 0] Warmup: {self.config.warmup_epochs} epochs...")
        self.student.train()
        total_loss = 0.0

        for epoch in range(self.config.warmup_epochs):
            x_rand = self._generate_batch(self.config.batch_size)
            teacher_probs = self._teacher_probs(x_rand)

            # KL divergence: KL(S(x) ‖ T(x))
            log_student = self.student(x_rand)
            loss = F.kl_div(log_student, teacher_probs, reduction="batchmean")

            self.student_optimizer.zero_grad()
            loss.backward()
            self.student_optimizer.step()
            total_loss += loss.item()

            # Populate replay buffer with diverse warmup samples
            self.replay_buffer.add(x_rand.detach())

        avg_loss = total_loss / self.config.warmup_epochs
        print(f"    Warmup complete | avg KL loss: {avg_loss:.4f} | replay buffer: {len(self.replay_buffer)} samples")

    def _bin_stabilization(self):
        """
        Phase 1: Learn adaptive bin boundaries aligned with teacher decisions.
        Temperature anneals from tau_start to tau_end across bin_epochs.
        Bins are frozen after this phase completes.
        (Eq. 1: L_bin; Eq. 2: L_gen^(1))
        """
        print(f"  [Phase 1] Bin Stabilization: {self.config.bin_epochs} epochs...")
        self.bin_learner.train()
        self.generator.train()
        total_bin_loss = 0.0

        for epoch in range(self.config.bin_epochs):
            # Update temperature for annealing
            tau = self._annealed_temperature(epoch)
            self.bin_learner.set_temperature(tau)

            # ── Generator: boundary-focused sampling (Eq. 2) ────────────────
            x_gen = self._generate_batch(self.config.batch_size)
            teacher_probs = self._teacher_probs(x_gen)
            student_log_probs = self.student(x_gen)

            # L_class-div: maximize class probability spread → class balance
            student_probs = student_log_probs.exp()
            class_mean = student_probs.mean(dim=0)
            # Maximize entropy of class marginals
            l_class_div = (class_mean * class_mean.log().clamp(-10)).sum()

            # L_entropy: maximize teacher output entropy (target decision boundaries)
            t_entropy = -(teacher_probs * (teacher_probs + 1e-10).log()).sum(dim=-1).mean()
            l_entropy = -t_entropy  # negative = minimize entropy = maximize boundary proximity

            gen_loss_phase1 = (self.config.lambda_div * l_class_div
                               + self.config.lambda_boundary * l_entropy)

            self.gen_optimizer.zero_grad()
            gen_loss_phase1.backward(retain_graph=True)
            self.gen_optimizer.step()

            # ── Bin Learner: align bins with teacher decision boundaries (Eq. 1)
            x_gen_detach = x_gen.detach()
            teacher_probs_detach = self._teacher_probs(x_gen_detach)
            m = self.bin_learner(x_gen_detach)
            bin_loss = self.bin_loss_fn(m, teacher_probs_detach)

            self.bin_optimizer.zero_grad()
            bin_loss.backward()
            self.bin_optimizer.step()
            total_bin_loss += bin_loss.item()

            if (epoch + 1) % 50 == 0:
                print(f"    Bin epoch {epoch+1}/{self.config.bin_epochs} | bin_loss={bin_loss.item():.4f} | τ={tau:.3f}")

        # Freeze bin learner
        for p in self.bin_learner.parameters():
            p.requires_grad = False
        self.bin_learner.eval()
        print(f"    Bins frozen. avg_bin_loss={total_bin_loss/self.config.bin_epochs:.4f}")

    def _adversarial_distillation(self):
        """
        Phase 2: Adversarial generator + student training with frozen bins.
        Generator maximizes interaction coverage + hardness (Eq. 6).
        Student minimizes KL from teacher (Eq. 7).
        Training mix: 90% adversarial + 10% replay.
        """
        print(f"  [Phase 2] Adversarial Distillation: {self.config.adv_epochs} epochs...")
        n_replay = max(1, int(self.config.batch_size * self.config.replay_ratio))
        n_adv = self.config.batch_size - n_replay

        all_bin_assignments = []  # collect for coverage tracking

        for epoch in range(self.config.adv_epochs):
            self.generator.train()
            self.student.train()

            # ── Generate adversarial samples ──────────────────────────────────
            x_adv = self._generate_batch(n_adv)
            teacher_probs = self._teacher_probs(x_adv)

            # Compute soft bin memberships with frozen bin learner
            with torch.no_grad():
                m = self.bin_learner(x_adv)

            # ── Generator Loss (Eq. 6): L_gen = L_diversity + λ_hard · L_hardness
            m_gen = self.bin_learner(x_adv)  # gradients flow to generator via x_adv
            l_diversity = self.diversity_loss_fn(m_gen)

            student_log_probs = self.student(x_adv).detach()
            # L_hardness (Eq. 5): reward generator for finding student failures
            l_hardness = -F.kl_div(
                student_log_probs, teacher_probs, reduction="batchmean"
            )

            gen_loss = (self.config.lambda_cov * l_diversity
                        + self.config.lambda_hard * l_hardness)

            self.gen_optimizer.zero_grad()
            gen_loss.backward()
            nn.utils.clip_grad_norm_(self.generator.parameters(), 5.0)
            self.gen_optimizer.step()

            # ── Student Loss (Eq. 7): minimize KL(S(x) ‖ T(x)) ───────────────
            x_adv_fresh = self._generate_batch(n_adv).detach()
            teacher_probs_fresh = self._teacher_probs(x_adv_fresh)

            # Mix with replay buffer samples
            replay = self.replay_buffer.sample(n_replay, self.device)
            if replay is not None:
                x_mix = torch.cat([x_adv_fresh, replay], dim=0)
                t_replay = self._teacher_probs(replay)
                t_mix = torch.cat([teacher_probs_fresh, t_replay], dim=0)
            else:
                x_mix, t_mix = x_adv_fresh, teacher_probs_fresh

            student_loss = F.kl_div(self.student(x_mix), t_mix, reduction="batchmean")
            self.student_optimizer.zero_grad()
            student_loss.backward()
            nn.utils.clip_grad_norm_(self.student.parameters(), 5.0)
            self.student_optimizer.step()

            # Track bin assignments for coverage measurement
            with torch.no_grad():
                hard_assignments = self.bin_learner.hard_assignment(x_adv_fresh)
                all_bin_assignments.append(hard_assignments.cpu())

            self.gen_scheduler.step()
            self.student_scheduler.step()

            if (epoch + 1) % 100 == 0:
                all_ba = torch.cat(all_bin_assignments, dim=0)
                coverage = self.diversity_loss_fn.compute_coverage(all_ba)
                print(
                    f"    Adv epoch {epoch+1}/{self.config.adv_epochs} | "
                    f"gen_loss={gen_loss.item():.4f} | "
                    f"student_kl={student_loss.item():.4f} | "
                    f"coverage={coverage:.1%}"
                )
                self.history.append({
                    "epoch": epoch + 1,
                    "gen_loss": gen_loss.item(),
                    "student_kl": student_loss.item(),
                    "coverage": coverage,
                })

    def train(self):
        """Execute full three-phase TabKD training."""
        print(f"\n{'='*60}")
        print(f"  TabKD Training | {self.config.n_features} features | {self.config.n_bins} bins")
        print(f"  Phases: Warmup({self.config.warmup_epochs}) → "
              f"Bins({self.config.bin_epochs}) → Adv({self.config.adv_epochs})")
        print(f"{'='*60}\n")
        self._warmup()
        self._bin_stabilization()
        self._adversarial_distillation()
        print("\n  TabKD training complete.")

    @torch.no_grad()
    def evaluate(self, X_test: Tensor, y_test: Tensor) -> Dict[str, float]:
        """
        Evaluate student model on held-out test data.

        Returns accuracy, F1, teacher-student agreement, and coverage.
        (Eqs. 10–13 in paper)
        """
        self.student.eval()
        student_preds = self.student.predict(X_test)
        teacher_preds = get_teacher_probs(self.teacher, X_test).argmax(dim=-1)
        y_np = y_test.cpu().numpy()
        s_np = student_preds.cpu().numpy()
        t_np = teacher_preds.cpu().numpy()

        acc = (s_np == y_np).mean()

        # F1 score (binary)
        tp = ((s_np == 1) & (y_np == 1)).sum()
        fp = ((s_np == 1) & (y_np == 0)).sum()
        fn = ((s_np == 0) & (y_np == 1)).sum()
        precision = tp / max(1, tp + fp)
        recall = tp / max(1, tp + fn)
        f1 = 2 * precision * recall / max(1e-8, precision + recall)

        # Teacher-student agreement (Eq. 12)
        agreement = (s_np == t_np).mean()

        # Coverage (Eq. 13)
        hard_assignments = self.bin_learner.hard_assignment(X_test)
        coverage = self.diversity_loss_fn.compute_coverage(hard_assignments.cpu())

        return {
            "accuracy": float(acc) * 100,
            "f1": float(f1) * 100,
            "agreement": float(agreement) * 100,
            "coverage": float(coverage) * 100,
        }


# ─── SECTION 9: Dataset Helpers & Evaluation ──────────────────────────────────

def load_benchmark_dataset(
    name: str = "adult",
    n_samples: int = 1000,
    test_ratio: float = 0.2,
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
    """
    Load one of the four TabKD benchmark datasets (Section 5.2).

    Production: download real datasets from UCI Machine Learning Repository:
      - Adult:         https://doi.org/10.24432/C5XW20
      - Breast Cancer: https://doi.org/10.24432/C51P4M
      - Credit:        https://doi.org/10.24432/C55S3H
      - Mushroom:      https://doi.org/10.24432/C5959T

    Production loading example (requires sklearn):
        from sklearn.datasets import load_breast_cancer
        from sklearn.preprocessing import StandardScaler
        data = load_breast_cancer()
        X = StandardScaler().fit_transform(data.data)
        y = data.target

    This function generates synthetic data with the same shape
    as each real dataset for demonstration purposes.

    Returns
    -------
    X_train, X_test, y_train, y_test — all as float32 Tensors
    """
    dataset_specs = {
        "adult":         {"n_features": 14, "name": "Adult Income (48K, 14 features)"},
        "credit":        {"n_features": 23, "name": "Credit Default (30K, 23 features)"},
        "breast_cancer": {"n_features": 30, "name": "Breast Cancer (569, 30 features)"},
        "mushroom":      {"n_features": 22, "name": "Mushroom (8K, 22 features)"},
    }

    if name not in dataset_specs:
        raise ValueError(f"Unknown dataset '{name}'. Choose from: {list(dataset_specs.keys())}")

    spec = dataset_specs[name]
    n_features = spec["n_features"]

    # Try to load real datasets via sklearn if available
    if name == "breast_cancer":
        try:
            from sklearn.datasets import load_breast_cancer
            from sklearn.preprocessing import StandardScaler
            from sklearn.model_selection import train_test_split
            data = load_breast_cancer()
            X = StandardScaler().fit_transform(data.data).astype(np.float32)
            y = data.target.astype(np.int64)
            X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=test_ratio, stratify=y, random_state=42)
            print(f"  Loaded real Breast Cancer dataset: {X.shape[0]} samples, {X.shape[1]} features")
            return (torch.tensor(X_tr), torch.tensor(X_te),
                    torch.tensor(y_tr), torch.tensor(y_te))
        except ImportError:
            pass

    # Synthetic fallback
    torch.manual_seed(42)
    X = torch.randn(n_samples, n_features)
    # Synthetic labels: XOR of first two features (non-linear interaction)
    y = ((X[:, 0] > 0).long() ^ (X[:, 1] > 0).long())

    n_test = int(n_samples * test_ratio)
    n_train = n_samples - n_test
    X_train, X_test = X[:n_train], X[n_train:]
    y_train, y_test = y[:n_train], y[n_train:]

    print(f"  Loaded synthetic '{name}' data: {n_train} train / {n_test} test | {n_features} features")
    return X_train, X_test, y_train, y_test


def run_tabkd_experiment(
    dataset_name: str = "breast_cancer",
    teacher_type: str = "nn",
    n_samples: int = 800,
    warmup_epochs: int = 5,
    bin_epochs: int = 20,
    adv_epochs: int = 40,
    n_bins: int = 8,
    verbose: bool = True,
) -> Dict[str, float]:
    """
    Run a complete TabKD experiment for a dataset + teacher combination.

    Parameters
    ----------
    dataset_name  : "adult" | "credit" | "breast_cancer" | "mushroom"
    teacher_type  : "nn" | "xgboost" | "rf" | "tabtransformer"
    n_samples     : total samples (use full dataset in production)
    warmup_epochs : Stage 0 epochs (paper: 30)
    bin_epochs    : Stage 1 epochs (paper: 200)
    adv_epochs    : Stage 2 epochs (paper: 400)
    n_bins        : bins per feature (paper: 8)

    Returns
    -------
    dict with accuracy, f1, agreement, coverage metrics
    """
    dataset_spec = {
        "adult": 14, "credit": 23, "breast_cancer": 30, "mushroom": 22
    }
    n_features = dataset_spec.get(dataset_name, 14)
    device = torch.device("cpu")

    print(f"\n{'─'*60}")
    print(f"  Dataset: {dataset_name} | Teacher: {teacher_type} | Bins: {n_bins}")
    print(f"{'─'*60}")

    # Load dataset
    X_train, X_test, y_train, y_test = load_benchmark_dataset(dataset_name, n_samples)
    n_features = X_train.shape[1]

    # Build and train teacher
    if teacher_type == "nn":
        teacher = NeuralNetworkTeacher(n_features, n_classes=2).to(device)
        teacher.train_on_data(X_train, y_train, epochs=15)
        teacher.eval()
        teacher_acc = (teacher.predict(X_test) == y_test).float().mean().item() * 100
    elif teacher_type == "xgboost":
        teacher = XGBoostTeacher()
        teacher.fit(X_train.numpy(), y_train.numpy())
        teacher_acc = (teacher.predict(X_test) == y_test).float().mean().item() * 100
    elif teacher_type == "rf":
        teacher = RandomForestTeacher()
        teacher.fit(X_train.numpy(), y_train.numpy())
        teacher_acc = (teacher.predict(X_test) == y_test).float().mean().item() * 100
    elif teacher_type == "tabtransformer":
        teacher = TabTransformerTeacher(n_features, n_classes=2).to(device)
        teacher.train_on_data(X_train, y_train, epochs=15)
        teacher.eval()
        teacher_acc = (teacher.predict(X_test) == y_test).float().mean().item() * 100
    else:
        raise ValueError(f"Unknown teacher type: {teacher_type}")

    print(f"  Teacher accuracy: {teacher_acc:.1f}%")

    # Build TabKD config
    config = TabKDConfig(
        n_features=n_features,
        n_classes=2,
        n_bins=n_bins,
        warmup_epochs=warmup_epochs,
        bin_epochs=bin_epochs,
        adv_epochs=adv_epochs,
    )

    # Run TabKD training
    trainer = TabKDTrainer(config, teacher, device)
    trainer.train()

    # Evaluate
    results = trainer.evaluate(X_test.to(device), y_test.to(device))
    results["teacher_accuracy"] = teacher_acc

    print(f"\n  Results | Acc={results['accuracy']:.1f}% | F1={results['f1']:.1f}% | "
          f"Agreement={results['agreement']:.1f}% | Coverage={results['coverage']:.1f}%")
    print(f"  (Paper target: {teacher_type}+{dataset_name} → see Table 2 for expected values)")
    return results


# ─── SECTION 10: Smoke Test ────────────────────────────────────────────────────

if __name__ == "__main__":
    print("=" * 60)
    print("TabKD — Full Framework Smoke Test")
    print("=" * 60)
    torch.manual_seed(42)
    np.random.seed(42)
    device = torch.device("cpu")

    # ── 1. Soft Bin Learner ──────────────────────────────────────────────────
    print("\n[1/7] Soft Bin Learner — shape and range checks...")
    config = TabKDConfig(n_features=14, n_classes=2, n_bins=8)
    bin_learner = SoftBinLearner(config.n_features, config.n_bins).to(device)
    x_test = torch.randn(32, config.n_features)
    m = bin_learner(x_test)
    assert m.shape == (32, 14, 8), f"Expected (32,14,8), got {m.shape}"
    assert (m >= 0).all() and (m <= 1).all(), "Memberships out of [0,1]"
    assert torch.allclose(m.sum(dim=-1), torch.ones(32, 14), atol=1e-5), "Memberships must sum to 1 per feature"
    print(f"  Membership tensor shape: {tuple(m.shape)} ✓")
    print(f"  Values in [0,1], sum to 1 per feature ✓")

    # ── 2. Bin Learning Loss ─────────────────────────────────────────────────
    print("\n[2/7] Bin Learning Loss...")
    bin_loss_fn = BinLearningLoss(config)
    teacher_probs = torch.softmax(torch.randn(32, 2), dim=-1)
    loss_bin = bin_loss_fn(m, teacher_probs)
    assert torch.isfinite(loss_bin), "Bin loss is not finite"
    print(f"  L_bin = {loss_bin.item():.4f} (finite, differentiable) ✓")

    # ── 3. Interaction Diversity Loss ────────────────────────────────────────
    print("\n[3/7] Interaction Diversity Loss & Coverage...")
    div_loss_fn = InteractionDiversityLoss(config)
    loss_div = div_loss_fn(m)
    assert torch.isfinite(loss_div), "Diversity loss is not finite"
    n_pairs = config.n_features * (config.n_features - 1) // 2
    print(f"  Feature pairs: {n_pairs} | L_diversity = {loss_div.item():.4f} ✓")

    hard_assignments = bin_learner.hard_assignment(x_test)
    coverage = div_loss_fn.compute_coverage(hard_assignments.cpu())
    print(f"  Pairwise coverage on 32 samples: {coverage:.1%}")
    total_cells = n_pairs * (config.n_bins ** 2)
    print(f"  Total pairwise bin cells: {total_cells:,}")

    # ── 4. Generator & Student Forward Pass ──────────────────────────────────
    print("\n[4/7] Generator & Student Networks...")
    gen = TabularGenerator(config).to(device)
    student = TabularStudent(config).to(device)
    x_gen = gen.sample(64, device)
    assert x_gen.shape == (64, 14), f"Generator output shape mismatch: {x_gen.shape}"
    log_probs = student(x_gen)
    assert log_probs.shape == (64, 2), f"Student output shape mismatch: {log_probs.shape}"
    assert torch.allclose(log_probs.exp().sum(dim=-1), torch.ones(64), atol=1e-5)
    print(f"  Generator: noise({config.noise_dim}) → samples({config.n_features}) ✓")
    print(f"  Student: samples({config.n_features}) → log_probs({config.n_classes}) ✓")
    n_gen = sum(p.numel() for p in gen.parameters())
    n_stu = sum(p.numel() for p in student.parameters())
    print(f"  Generator params: {n_gen:,} | Student params: {n_stu:,}")

    # ── 5. Neural Network Teacher ─────────────────────────────────────────────
    print("\n[5/7] Teacher Training (Neural Network)...")
    teacher_nn = NeuralNetworkTeacher(config.n_features, config.n_classes).to(device)
    X_mock = torch.randn(200, config.n_features)
    y_mock = ((X_mock[:, 0] > 0).long() ^ (X_mock[:, 1] > 0).long())
    loss_final = teacher_nn.train_on_data(X_mock, y_mock, epochs=10)
    teacher_nn.eval()
    teacher_probs_mock = get_teacher_probs(teacher_nn, X_mock[:8])
    assert teacher_probs_mock.shape == (8, 2)
    assert torch.allclose(teacher_probs_mock.sum(dim=-1), torch.ones(8), atol=1e-5)
    print(f"  Teacher trained | final loss: {loss_final:.4f}")
    print(f"  Teacher probs shape: {tuple(teacher_probs_mock.shape)} | sum to 1 ✓")

    # ── 6. Replay Buffer ─────────────────────────────────────────────────────
    print("\n[6/7] Replay Buffer...")
    buf = ReplayBuffer(max_size=500)
    for _ in range(10):
        buf.add(torch.randn(20, config.n_features))
    sample = buf.sample(32, device)
    assert sample is not None and sample.shape[1] == config.n_features
    print(f"  Buffer size: {len(buf)} | Sample shape: {tuple(sample.shape)} ✓")

    # ── 7. Full Training Run (short) ──────────────────────────────────────────
    print("\n[7/7] Full TabKD Training Run (Breast Cancer, NN teacher, mini-epochs)...")
    results = run_tabkd_experiment(
        dataset_name="breast_cancer",
        teacher_type="nn",
        n_samples=400,
        warmup_epochs=3,
        bin_epochs=10,
        adv_epochs=20,
        n_bins=8,
    )

    print("\n" + "=" * 60)
    print("✓  All TabKD checks passed. Framework is ready for use.")
    print("=" * 60)
    print("""
Next steps to reproduce paper results:
  1. Download real datasets from UCI Repository:
       Adult:         https://doi.org/10.24432/C5XW20
       Breast Cancer: https://doi.org/10.24432/C51P4M
       Credit:        https://doi.org/10.24432/C55S3H
       Mushroom:      https://doi.org/10.24432/C5959T

  2. Install full dependencies:
       pip install scikit-learn xgboost tab-transformer-pytorch

  3. Run full training with paper hyperparameters:
       config = TabKDConfig(
           n_bins=8, batch_size=128, lr=0.001,
           warmup_epochs=30, bin_epochs=200, adv_epochs=400,
           lambda_cov=10.0, lambda_hard=2.0,
       )

  4. Apply teacher-specific temperature schedules (Table 1):
       NN:              tau_start=1.0, tau_end=0.05
       Random Forest:   tau_start=1.2, tau_end=0.08
       XGBoost:         tau_start=1.5, tau_end=0.10
       TabTransformer:  tau_start=1.2, tau_end=0.08

  5. Access open-source implementation:
       https://anonymous.4open.science/r/int_div-0413/README.md

  6. Expected results (agreement %):
       Breast Cancer + NN:       ~95.6%
       Breast Cancer + XGBoost:  ~96.5%
       Credit + NN:              ~97.0%
       Mushroom + NN:            ~96.5%
""")