NIMO: a Nonlinear Interpretable MOdel

Deep learning has achieved remarkable success across many domains, but it has also created a growing demand for interpretability in model predictions. Although many explainable machine learning methods have been proposed, post-hoc explanations lack guaranteed fidelity and are sensitive to hyperparameter choices, highlighting the appeal of inherently interpretable models. For example, linear regression provides clear feature effects through its coefficients. However, such models are often outperformed by more complex neural networks (NNs) that usually lack inherent interpretability. To address this dilemma, we introduce NIMO, a framework that combines inherent interpretability with the expressive power of neural networks. Building on the simple linear regression, NIMO is able to provide flexible and intelligible feature effects. Relevantly, we develop an optimization method based on parameter elimination, that allows for optimizing the NN parameters and linear coefficients effectively and efficiently. By relying on adaptive ridge regression we can easily incorporate sparsity as well. We show empirically that our model can provide faithful and intelligible feature effects while maintaining good predictive performance.

💡 Research Summary

The paper introduces NIMO (Nonlinear Interpretable Model), a hybrid architecture that preserves the global interpretability of linear regression while adding the expressive power of neural networks through instance‑specific nonlinear corrections. The core formulation augments each linear coefficient β_j with a multiplicative factor h_j(x) = 1 + g_{u_j}(x_{‑j}), where g_{u_j} is a scalar‑valued neural network evaluated on the input vector with the j‑th feature masked out. By enforcing g_{u_j}(0)=0, the model collapses to a pure linear regression when all features are at their standardized mean (zero), guaranteeing that the marginal effect at the mean (MEM) equals β_j. Consequently, β_j provides a transparent global feature effect, while h_j(x) supplies flexible local adjustments, unifying global and instance‑level explanations within a single model.

Training NIMO is non‑trivial because the linear coefficients β and the neural network parameters u are tightly coupled. The authors adopt a parameter‑elimination strategy inspired by profile likelihood: they solve for β in closed form as a function of u, β̂(u) = (B_uᵀB_u + λI)⁻¹ B_uᵀy, where B_u = X + X ∘ G_u and G_u = g_u(X) contains the per‑feature correction terms. Substituting β̂(u) back into the loss reduces the optimization to a problem over u alone, which can be tackled efficiently with gradient‑based methods. To encourage sparsity, the ℓ₂ penalty is replaced by ℓ₁. Since ℓ₁ does not admit a closed‑form solution, the authors employ adaptive ridge regression: at each iteration they compute a weighted ridge solution that is equivalent to the Lasso optimum, thus preserving sparsity while retaining the benefits of parameter elimination.

The architecture uses a single shared neural network g_u for all features. Positional encodings identify the target feature, and the masked input (with the j‑th component zeroed) allows the network to learn feature‑specific corrections without replicating d separate networks. Group ℓ₂ regularization on the first fully‑connected layer further promotes feature‑level sparsity, making it easier to interpret which inputs drive the nonlinear adjustments.

Empirical evaluation spans synthetic data (with known linear and interaction components) and several real‑world datasets from healthcare and finance. Results show that NIMO accurately recovers the true β coefficients on synthetic data, achieves lower mean‑squared error than pure neural networks, and outperforms existing hybrid methods such as Neural Additive Models (NAM), LassoNet, and Interpretable Mixture of Experts (IME) in both predictive performance and interpretability metrics. Sparsity experiments demonstrate that increasing the regularization strength prunes irrelevant features while preserving the most informative ones, confirming the effectiveness of the adaptive ridge/Lasso scheme.

The authors discuss limitations: (1) computational overhead grows with the number of features because each forward pass must mask and encode each dimension, which may be mitigated by dimensionality reduction or sparse masking; (2) MEM, while identical to β_j, may not fully capture feature importance in highly nonlinear regimes, but the complementary local corrections h_j(x) help alleviate this issue. Future work includes scaling the approach to high‑dimensional and multimodal data, integrating Bayesian uncertainty estimates, and exploring more efficient hardware implementations of the masking operation.

In summary, NIMO offers a principled solution to the accuracy‑interpretability trade‑off by combining closed‑form linear coefficient estimation, adaptive sparsity, and per‑instance neural corrections. Its ability to deliver faithful global feature effects alongside accurate local explanations makes it a compelling choice for high‑stakes applications where model transparency is as critical as predictive power.