Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TabMixNN: A Unified Deep Learning Framework for Structural Mixed Effects Modeling on Tabular Data

Deniz Akdemir

TL;DR

Key innovations include an R-style formula interface for accessibility, support for directed acyclic graph constraints for causal structure learning, Stochastic Partial Differential Equation kernels for spatial modeling, and comprehensive interpretability tools including SHAP values and variance decomposition.

Abstract

We present TabMixNN, a flexible PyTorch-based deep learning framework that synthesizes classical mixed-effects modeling with modern neural network architectures for tabular data analysis. TabMixNN addresses the growing need for methods that can handle hierarchical data structures while supporting diverse outcome types including regression, classification, and multitask learning. The framework implements a modular three-stage architecture: (1) a mixed-effects encoder with variational random effects and flexible covariance structures, (2) backbone architectures including Generalized Structural Equation Models (GSEM) and spatial-temporal manifold networks, and (3) outcome-specific prediction heads supporting multiple outcome families. Key innovations include an R-style formula interface for accessibility, support for directed acyclic graph (DAG) constraints for causal structure learning, Stochastic Partial Differential Equation (SPDE) kernels for spatial modeling, and comprehensive interpretability tools including SHAP values and variance decomposition. We demonstrate the framework's flexibility through applications to longitudinal data analysis, genomic prediction, and spatial-temporal modeling. TabMixNN provides a unified interface for researchers to leverage deep learning while maintaining the interpretability and theoretical grounding of classical mixed-effects models.

TabMixNN: A Unified Deep Learning Framework for Structural Mixed Effects Modeling on Tabular Data

TL;DR

Key innovations include an R-style formula interface for accessibility, support for directed acyclic graph constraints for causal structure learning, Stochastic Partial Differential Equation kernels for spatial modeling, and comprehensive interpretability tools including SHAP values and variance decomposition.

Abstract

We present TabMixNN, a flexible PyTorch-based deep learning framework that synthesizes classical mixed-effects modeling with modern neural network architectures for tabular data analysis. TabMixNN addresses the growing need for methods that can handle hierarchical data structures while supporting diverse outcome types including regression, classification, and multitask learning. The framework implements a modular three-stage architecture: (1) a mixed-effects encoder with variational random effects and flexible covariance structures, (2) backbone architectures including Generalized Structural Equation Models (GSEM) and spatial-temporal manifold networks, and (3) outcome-specific prediction heads supporting multiple outcome families. Key innovations include an R-style formula interface for accessibility, support for directed acyclic graph (DAG) constraints for causal structure learning, Stochastic Partial Differential Equation (SPDE) kernels for spatial modeling, and comprehensive interpretability tools including SHAP values and variance decomposition. We demonstrate the framework's flexibility through applications to longitudinal data analysis, genomic prediction, and spatial-temporal modeling. TabMixNN provides a unified interface for researchers to leverage deep learning while maintaining the interpretability and theoretical grounding of classical mixed-effects models.
Paper Structure (106 sections, 48 equations, 3 figures, 7 tables, 1 algorithm)

This paper contains 106 sections, 48 equations, 3 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions
  3. Paper Organization
  4. Background and Related Work
  5. Classical Mixed-Effects Models
  6. Deep Learning for Tabular Data
  7. Structural Equation Models and Causal Learning
  8. Spatial-Temporal Deep Learning
  9. Methods
  10. Mixed-Effects Encoder
  11. Fixed Effects
  12. Random Effects
  13. Constraint Enforcement and Identifiability
  14. Covariance Structures
  15. IID Covariance
  16. ...and 91 more sections

Figures (3)

  • Figure 1: Per-subject effect of Days (sleep deprivation) on Reaction Time using SHAP values. Each panel shows one subject's SHAP values plotted against Days, with fitted lines indicating the strength of the Days effect for that individual. Positive slopes indicate subjects whose reaction times worsen more with sleep deprivation, while negative or flat slopes indicate resilient subjects. This demonstrates heterogeneity in random slopes: different subjects show different sensitivities to sleep deprivation.
  • Figure 2: Comparison of subject-specific Days effects across all subjects. Left: Bar plot showing the SHAP slope (Days effect) for each subject, sorted by magnitude. Positive values indicate subjects who get worse with more sleep deprivation. The color gradient visualizes the range of heterogeneity. Right: Distribution of subject-specific effects showing mean and median. This quantifies individual differences in sensitivity to sleep deprivation.
  • Figure 3: SHAP waterfall plots showing feature contributions for four subjects at different time points (Days 0, 5, and 9). Each bar represents the SHAP contribution of a feature to the prediction at that time point. This visualization demonstrates how fixed effects (Days) and other features combine to produce predictions, and how these contributions change over time within subjects.