Statistical-Neural Interaction Networks for Interpretable Mixed-Type Data Imputation

TL;DR

Statistical-Neural Interaction (SNI), an interpretable mixed-type imputation framework that couples correlation-derived statistical priors with neural feature attention through a Controllable-Prior Feature Attention (CPFA) module, is proposed.

Abstract

Real-world tabular databases routinely combine continuous measurements and categorical records, yet missing entries are pervasive and can distort downstream analysis. We propose Statistical-Neural Interaction (SNI), an interpretable mixed-type imputation framework that couples correlation-derived statistical priors with neural feature attention through a Controllable-Prior Feature Attention (CPFA) module. CPFA learns head-wise prior-strength coefficients $\{λ_h\}$ that softly regularize attention toward the prior while allowing data-driven deviations when nonlinear patterns appear to be present in the data. Beyond imputation, SNI aggregates attention maps into a directed feature-dependency matrix that summarizes which variables the imputer relied on, without requiring post-hoc explainers. We evaluate SNI against six baselines (Mean/Mode, MICE, KNN, MissForest, GAIN, MIWAE) on six datasets spanning ICU monitoring, population surveys, socio-economic statistics, and engineering applications. Under MCAR/strict-MAR at 30\% missingness, SNI is generally competitive on continuous metrics but is often outperformed by accuracy-first baselines (MissForest, MIWAE) on categorical variables; in return, it provides intrinsic dependency diagnostics and explicit statistical-neural trade-off parameters. We additionally report MNAR stress tests (with a mask-aware variant) and discuss computational cost, limitations -- particularly for severely imbalanced categorical targets -- and deployment scenarios where interpretability may justify the trade-off.

Figures (10)

• Figure 1: Schematic of the Statistical-Neural Interaction (SNI) workflow. Grey icons show the observed/missing split $(\mathbf{X}_{\mathrm{obs}},\mathbf{X}_{\mathrm{mis}})$ and initial estimate $\mathbf{X}^{(0)}$. Each EM-inspired iteration comprises: (i) a Statistical step (blue) computing prior $\mathbf{P}_f$ from $\boldsymbol{\Sigma}^{(g-1)}$; (ii) Pseudo-masking via $\mathbf{M}_f\!\sim\!\mathrm{Bernoulli}(\rho)$; and (iii) a Neural step (orange) where CPFA regresses continuous targets or classifies categorical ones. Iterations terminate when relative change $\Delta<\varepsilon$. Outputs include the completed matrix $\mathbf{X}^{(g)}$, attention maps $\{\mathbf{A}^{(h)}\}$, and prior-confidence weights $\{\lambda_h\}$.
• Figure 2: Architecture of the Controllable-Prior Feature Attention (CPFA) module. Input features are value-projected and position-encoded (pink) before multi-head feature attention (blue), where each head $h$ produces weights $\mathbf{A}^{(h)}$ regularized toward prior $\mathbf{P}_f$ via $\lambda_h=\mathrm{softplus}(\theta_{\lambda,h})$. The forward path (orange) passes through residual Add & Norm and Feed-Forward layers (green), then weighted aggregation feeds the final Predictor. Grey arrows denote auxiliary components: importance-score extraction and a feature bypass preserving linear information. See Algorithm \ref{['alg:CPFA']} for formal definition.
• Figure 3: Compact cross-domain summary of SNI versus six baselines under MCAR and strict MAR at 30% missingness (mean over five seeds; panels A--B are further averaged over MCAR/MAR). (A) Continuous NRMSE by dataset (lower is better). (B) Continuous $R^2$ by dataset (higher is better; values below $-0.1$ are clipped for readability). (C) Average per-setting rank across metrics (lower is better), with SNI highlighted. (D) Pairwise win rate of SNI against each baseline, aggregated over all reported metrics and all dataset--mechanism settings (a "win" is defined as SNI outperforming the baseline on a given metric in a given setting). Categorical metrics are computed only on datasets with categorical variables.
• Figure 4: Model-reliance dependency network derived from SNI on MIMIC-IV (strict MAR, 30% missingness). Each directed edge $j\!\rightarrow\! i$ corresponds to an entry $D_{ij}$ and its thickness encodes the average attention mass, i.e., how strongly SNI relied on source feature $j$ when imputing target feature $i$. Node size reflects incoming mass $\Sigma_j=\sum_i D_{ij}$, highlighting globally informative source variables (RESP and ALARM in this subset). This visualization is intended as an imputation diagnostic and does not imply causal relationships.
• Figure 5: Feature dependency matrix $D$ for MIMIC-IV derived from CPFA attention (strict MAR, 30% missingness). $D_{ij}$ denotes the normalized contribution from source feature $j$ (column) when imputing target feature $i$ (row), with $D_{ii}=0$. The first row illustrates that imputing SpO$_2$ relies most on RESP (0.60), SBP (0.19), and ALARM (0.19). As a model-reliance view, $D$ provides a compact audit trail of which columns were used by the imputer and may help identify unexpected reliance patterns or inform feature screening, even when an accuracy-first baseline is preferred.