TabMixNN: A Unified Deep Learning Framework for Structural Mixed Effects Modeling on Tabular Data
This work addresses the need for methods that handle hierarchical data structures while maintaining interpretability for researchers in fields like genomics and spatial analysis, though it appears incremental as it synthesizes existing techniques.
The authors tackled the problem of analyzing hierarchical tabular data with diverse outcome types by developing TabMixNN, a deep learning framework that integrates classical mixed-effects modeling with neural networks, resulting in a flexible tool that supports regression, classification, multitask learning, and applications like longitudinal data analysis and spatial-temporal modeling.
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.