Source-linked AI summary
Implementing neural network mixed-effects models in Template Model Builder (TMB)
Nan Zheng, Hoi Yiu Cheung, Vibhu Sharma, James T. Thorson, Noel G. Cadigan
TL;DR
The paper addresses limitations in implementing neural network mixed-effects models by proposing their implementation in Template Model Builder. The approach uses automatic differentiation and full-likelihood methods, with reported gains in estimation accuracy and implementation efficiency, while the current implementation remains limited to manually coded feedforward networks.
Problem
Existing neural network implementations face limitations related to manually derived objective functions and gradients, while standard feedforward networks do not explicitly encode structural priors.
Method
The paper proposes implementing neural network mixed-effects models in Template Model Builder, requiring users to specify the relevant likelihood components while using automatic differentiation and Laplace approximation.
Results
The proposed methodology offers enhanced estimation accuracy and improved implementation efficiency, with advantages highlighted from adopting a full likelihood.
Takeaways & Limitations
TMB provides a framework for implementing neural network mixed-effects models with reported efficiency, flexibility, and statistical-performance advantages.
Takeaways & Limitations
The current implementation manually codes feedforward neural networks, and extending this approach to convolutional, recurrent, or transformer architectures remains an implementation boundary.
Abstract
from arXiv · showhide
Neural network mixed-effects models (NMMs) have gained traction by combining the strong representation and predictive power of artificial neural networks with the capacity of mixed-effects modeling to capture complex correlation structures. However, existing estimation approaches rely heavily on manual derivations of objective functions and gradients, which inherently forces simplifying approximations and severely constrains the complexity and accuracy of NMMs. In this work, we introduce a general framework for implementing NMMs using Template Model Builder (TMB). By leveraging automatic differentiation and Laplace approximation, TMB requires users to specify only the negative joint log-likelihood and any regularization terms. The framework automatically integrates out random effects and evaluates the marginal objective function alongside its exact gradients, eliminating the need for manual derivations or ad hoc approximations. We demonstrate the efficiency, flexibility, and statistical performance of TMB-based NMMs across two numerical examples, including an application to monotonic NMMs. Reproducible code is provided to facilitate broader adoption.
1 Introduction
Neural network mixed-effects models combine neural networks’ representational power with mixed-effects models’ flexibility for irregular and complex data. The paper proposes TMB-based estimation to avoid restrictive manual derivations and improve efficiency, accuracy, implementation, and scalability.
- Standard and specialized neural architectures can remain data-intensive or rely on structural characteristics that many statistical datasets lack.
- Mixed-effects models accommodate unbalanced longitudinal data, irregular observation schedules, complex sampling designs, and relatively small sample sizes.
- NMMs combine neural networks’ representational power with mixed-effects models’ flexibility and interpretability.
- NMMs can represent observational-model parameters through a neural-network output combined with random effects and an elementwise transformation.
- Manual NMM estimation requires simplifying approximations, omits a covariance-and-Hessian determinant term, and becomes increasingly difficult and error-prone as network depth grows.
- TMB accepts the joint objective function and automatically computes the Laplace-approximated marginal objective and gradients without omitting terms.
- The proposed approach improves computational efficiency, inferential accuracy, implementation ease, and scalability relative to manual derivation, with demonstrations on simulations and monotonic neural networks.
2 Methods
The framework replaces the fixed-effects component of mixed-effects models with a feed-forward neural network and allows broad observational distributions. TMB integrates out random effects using Laplace approximation while automatically deriving the marginal objective and gradients.
- Random effects follow multivariate normal distributions, inducing within-observation dependence and enabling dependence across observations through shared components or correlated distributions.
- Unlike the compared framework, the observational distribution is not restricted to exponential-family distributions.
- The network maps inputs and random-effects design information to an observational parameter through elementwise transformations.The parameter need not be a conditional mean; it may govern any observational distribution parameter.
- NMMs replace the fixed-effects portion of a GLMM linear predictor with a feed-forward neural network.
- TMB approximates the marginal likelihood with Laplace approximation and exploits sparsity to scale to tens of thousands of random effects.
- An L2 penalty on neural-network weights and biases mitigates overfitting, enhances generalization, and prioritizes predictive accuracy over feature selection.
3 Simulation study
The simulation evaluates TMB-based GNMM-2 models on nonlinear continuous and binary responses using raw predictors and leave-one-out prediction. TMB substantially outperforms the reported GNMM-2 results, with continuous-response MSPE near the noise floor and binary-response AUC increasing with random-intercept variance.
- The study reproduces a nonlinear simulation for both continuous and binary responses, assessing prediction by leave-one-out cross-validation.The final observation for each subject is withheld, and each configuration is replicated 500 times.
- The network receives only raw predictors and must recover XOR interactions that a linear predictor cannot represent.
- TMB integrates out subject random intercepts with Laplace approximation and optimizes the resulting marginal objective.
- TMB achieves approximately 0.003 MSPE across all τ values, versus approximately 0.45–0.55 for reported GNMM-2 results.The TMB MSPE tracks the irreducible noise floor σ^2_ϵ = 0.0025 while recovering the XOR signal.
- Binary-response AUC rises from approximately 0.89 at τ = 0 to 0.91 at τ = 20.The reported increase is attributed to more informative subject-specific random-intercept estimates for withheld observations.
- Across both response settings, the TMB implementation shows markedly better predictive performance than the reported GNMM-2 results.The comparison uses the nonlinear data-generating mechanism and a first hidden layer of 10 units.
4 Real data example: maturity data for 3LNO Amer-
The real-data application fits NMMs to female American plaice maturation data across 58 cohorts, using cohort segments and AR(1) random effects. Compared with a GLMM, the NMM allows nonlinear age effects and achieves slightly lower test BCE while capturing mild prediction asymmetry.
- Data and model setup: The dataset contains female American plaice maturation records from Northwest Atlantic surveys spanning 58 cohorts from 1958 to 2015.Maturity is binary, with 1 indicating mature fish and 0 indicating immature fish.
- Data and model setup: The data are partitioned into seven cohort segments, with six segments containing eight consecutive cohorts and the seventh containing 11.The final small segment is merged with the penultimate segment to avoid a three-cohort segment.
- Data and model setup: The NMM uses two hidden layers with 20 nodes each and an eight-dimensional input combining age with seven one-hot cohort-segment indicators.The logistic sigmoid is used in the hidden layers and the output yields conditional maturation probabilities given random effects.
5 Conclusion
The conclusion presents TMB as a framework for implementing NMMs with automatically differentiated marginal objectives and gradients. The reported advantages include estimation accuracy, implementation efficiency, scalability, and support for constrained nonlinear models, while broader architectures remain future work.
- Framework: TMB-based NMM implementation requires specifying the negative joint log-likelihood and regularization terms, then automatically integrates out random effects and computes exact marginal gradients.This removes the need for manual objective-function and gradient derivations or ad hoc approximations.
- Reported advantages: The authors report that full-likelihood formulation improves statistical efficiency relative to a quasi-likelihood approach.The conclusion identifies this as one advantage highlighted by the simulation studies.
- Reported advantages: The approach is reported to provide enhanced estimation accuracy, improved implementation efficiency, and scalability toward complex NMMs.These advantages are presented as consequences of automating derivations and avoiding heuristic omissions or approximations.
- Empirical demonstration: Empirical results demonstrate that monotonicity-constrained NMM objectives and gradients can be computed efficiently for downstream optimization.The application also supports NMM representation of intricate nonlinear input–output dynamics without extensive architecture tuning.
- Limitations and future work: The current implementation manually codes feedforward networks, making extension to convolutional, recurrent, or transformer architectures impractical.Future work proposes integrating Keras with RTMB and extending random effects to complex spatiotemporal structures and hybrid CPU/GPU execution.