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

TL;DR

TabMixNN is a versatile deep learning framework that combines classical mixed-effects models with neural networks for diverse tabular data analysis, supporting hierarchical structures, causal inference, and interpretability.

Contribution

It introduces a modular, PyTorch-based framework integrating mixed-effects modeling with modern neural architectures, including causal constraints and spatial kernels, for the first time.

Findings

Demonstrates flexibility across longitudinal, genomic, and spatial-temporal data.

Provides interpretability tools like SHAP and variance decomposition.

Supports diverse outcome types including regression and classification.

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.