SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shibal Ibrahim MIT
Cambridge, MA  
shibal@mit.edu
Gabriel Isaac Afriat  
MIT
Cambridge, MA afriatg@mit.edu
Kayhan Behdin  
MIT
Cambridge, MA behdink@mit.edu
Rahul Mazumder  
MIT
Cambridge, MA
rahulmaz@mit.edu
Generalized Additive Models (GAMs) are a family of flexible and interpretable models with old roots in statistics. GAMs are often used with pairwise interactions to improve model accuracy while still retaining flexibility and interpretability but lead to computational challenges as we are dealing with order of  $p^2$  terms. It is desirable to restrict the number of components (i.e., encourage sparsity) for easier interpretability, and better computational and statistical properties. Earlier approaches, considering sparse pairwise interactions, have limited scalability, especially when imposing additional structural interpretability constraints. We propose a flexible GRAND-SLAMIN framework that can learn GAMs with interactions under sparsity and additional structural constraints in a differentiable end-to-end fashion. We customize first-order gradient-based optimization to perform sparse backpropagation to exploit sparsity in additive effects for any differentiable loss function in a GPU-compatible manner. Additionally, we establish novel non-asymptotic prediction bounds for our estimators with tree-based shape functions. Numerical experiments on real-world datasets show that our toolkit performs favorably in terms of performance, variable selection and scalability when compared with popular toolkits to fit GAMs with interactions. Our work expands the landscape of interpretable modeling while maintaining prediction accuracy competitive with non-interpretable black-box models. Our code is available at https://github.com/mazumder-lab/grandslamin.
Many state-of-the-art learners e.g., tree ensembles, neural networks, kernel support vector machines, can be difficult to interpret. There have been various efforts to derive some post-training explainability from these models —see [5] for a survey. Post-hoc explainability attempts to explain black-box prediction with interpretable instance-specific approximations e.g., LIME [49] and SHAP [35]. However, such approximations are known to be unstable [12, 28], expensive [51] and inaccurate [32]. Hence, it is desirable to consider modeling approaches that are inherently interpretable.
Amongst classical approaches, there are some models that have inherent interpretability e.g., Linear models, CART [4] and Generalized Additive Models (GAMs)[14]. GAMs [15] which have old roots in statistics are considered a front-runner in the context of interpretable modeling. They consider an additive model of the main effects of the form:  $g(\mathbb{E}[y]) = \sum_{j\in [p]}f_j(x_j)$ , where  $x_{j}$  denotes the  $j$ th feature in input  $\boldsymbol {x}\in \mathbb{R}^p$ ,  $f_{j}$  is a univariate shape function and  $g$  denotes the link function that adapts the model to various settings such as regression or classification. GAMs are considered easy to interpret as the impact of each feature can be understood via visualizing the corresponding shape function e.g., plotting  $f_{j}(x_{j})$  vs  $x_{j}$ . However, such models often suffer in prediction performance when compared to black-box methods e.g., deep neural networks (DNNs). This can be attributed in part to the fact that GAMs do not consider interactions among covariates.
There has been some exciting work that aims to reduce this performance gap by considering GAMs with pairwise interactions [see, for example, 34, 60, 6, 46, 9, 10, and the references therein]. GAMs with pairwise interactions consider a model of the following form:
$$
g \left(\mathbb {E} [ y ]\right) = \sum_ {j \in [ p ]} f _ {j} \left(x _ {j}\right) + \sum_ {(j, k) \in \mathcal {I}} f _ {j, k} \left(x _ {j}, x _ {k}\right) \tag {1}
$$
where  $f_{j,k}$  is a bivariate shape function and  $\mathcal{I} \subseteq \{(1,2), (1,3), \dots, (p - 1,p)\}$  denotes the set of all pairwise interactions. Under this model,  $f_{j}(x_{j})$  is the  $j$ -main effect and  $f_{j,k}(x_{j}, x_{k})$  is the  $(j,k)$ -th interaction effect. Pairwise interactions are considered interpretable as each of the bivariate shape function  $f_{j,k}$  can be visualized as a heatmap on an  $x_{j}, x_{k}$ -plane. Despite their appeal, GAMs with pairwise interactions pose several challenges: (i) Learning all pairwise interaction effects of the order of  $p^2$  lead to computational and statistical challenges. (ii) Performing component selection such that only a few of the components  $\{f_{j}\}$  and  $\{f_{j,k}\}$  are nonzero in an end-to-end fashion (while training) is a hard combinatorial optimization problem. We remind the reader that component selection is needed to aid interpretability. (iii) Imposing structural constraints on the interaction effects, e.g., hierarchy [42, 13, 7, 39, 3, 60, 10] makes the associated optimization task more complex.
In this paper, we introduce a novel GRAND-SLAMIN framework that allows for a flexible way to do component selection in GAMs with interactions under additional structural constraints in an end-to-end fashion. In particular, we introduce an alternative formulation of GAMs with interactions with additional binary variables. Next, we smooth these binary variables so that we can effectively learn these components via first-order methods in smooth optimization (e.g., SGD). Our formulation appears to have an edge over existing methods in terms of (i) model flexibility for enhanced structural interpretability and (ii) computational efficiency. First, the binary variables allow us to impose in the model (a) component selection constraints and (b) additional structural constraints (e.g., hierarchy) via a unified optimization formulation. Both of constraints (a), (b) can aid interpretability, model compression, and result in faster inference and better statistical properties. Second, the our smoothing procedure for the binary variables allows us to have customized algorithms that exploit sparsity in the forward and backward pass of the backpropagation algorithm.
For structural interpretability, we study two notions: weak and strong hierarchy [42, 13, 7, 39, 3].
Weak Hierarchy:  $f_{j,k}\neq 0\Rightarrow f_j\neq 0$  or  $f_{k}\neq 0\quad \forall (j,k)\in \mathcal{I}, j\in [p], k\in [p].$  (2)
Strong Hierarchy:  $f_{j,k} \neq 0 \implies f_j \neq 0$  and  $f_k \neq 0 \quad \forall (j,k) \in \mathcal{I}, j \in [p], k \in [p]$ . (3)
Weak hierarchy allows an interaction effect  $f_{j,k}$  to be selected if either main effect  $f_{j}$  or  $f_{k}$  is selected. Strong hierarchy allows for an interaction effect  $f_{j,k}$  to be selected only if both main effects  $f_{j}$  and  $f_{k}$  are selected. Such hierarchy constraints are popular in high-dimensional statistics: (i) They lead to more interpretable models [36, 3, 10], (ii) They promote practical sparsity, i.e., reduce the number of features that need to be measured when making new predictions (see Sec. 6.2) — this can reduce future data collection costs [3, 61], (iii) Additional constraints can also help regularize a model, sometimes resulting in improved AUC (see Sec. 6.1). (iv) They can also reduce variance in estimation of main effects in the presence of interaction effects (see Sec. 6.4), allowing the user to have more "trust" on model interpretability explanations.
Contributions. To summarise, while it's well acknowledged that GAMs with sparse interactions are a useful flexible family of explainable models, learning them pose significant computational challenges due to the combinatorial nature of the associated optimization problem. Our technical contributions in this paper be summarized as:
1. We propose a novel optimization formulation that makes use of indicator (binary) variables. The indicator variables allow us to impose both (a) component selection and (b) structural constraints in an end-to-end fashion. We consider a smooth and continuous parameterization of the binary variables so that the optimization objective is differentiable (for a smooth training loss) and hence amenable to first order methods such as SGD.
2. We show the flexibility of our framework by considering two different notions of hierarchy. While these constraints improve interpretability, they make the combinatorial problem more challenging [3, 17]. We propose end-to-end algorithms to train these models, making our approach quite different from existing neural-based toolkits [60, 10].
3. We exploit sparsity in the indicator variables during the course of the training for sparse forward and backward passes in a customized backpropagation algorithm in a GPU-compatible manner. This provides speedups on training times up to a factor of  $10 \times$  over standard backpropagation.
4. We study novel statistical properties of our model, and present non-asymptotic prediction error bounds. Different from earlier work, our results apply to learning with tree-based shape functions (for both main and interaction effects).
5. We introduce a new open-source toolkit GRAND-SLAMIN and perform experiments on a collection of 16 real-world datasets to demonstrate the effectiveness of our toolkit in terms of prediction performance, variable selection and scalability.
GAMs. GAMs have a long history in statistics [14] and have been extensively studied. They're often studied with smooth spline shape functions [see, e.g., 37, 47, 20, 63, 62, and references therein]. Some works have studied tree-based shape functions [33] and neural basis functions [1, 58].
GAMs with all pairwise interactions. In this thread, [9] study low-rank decomposition with neural network shape functions; [46] fit all pairwise interactions using shared neural bases.
Sparse GAMs with interactions. [31] introduced COSSO, which penalizes the sum of the Sobolev norms of the functional components, producing sparse models. [23] propose ELAAN, which is an  $\ell_0$ -regularized formulation with smooth cubic splines. [23] demonstrate the usefulness of their approach in a regression setting in terms of compact component selection and efficiency on a large-scale Census survey response prediction. [34, 40, 6] explore tree-based shape functions for fitting additive models with sparse interactions. [34] consider all main effects and a subset of pairwise interactions; the subset of interactions are selected via greedy stage-wise interaction detection heuristics. [40] provide an efficient implementation of the above approach as Explainable Boosting Machines (EBMs). [6] propose NODE-GA $^2$ M: an end-to-end learning approach with differentiable neural oblivious decision (NODE) trees [44]. Component selection in NODE-GAM is achieved by constraining the number of trees, and each tree learns to use one or two features via entmax transformation [43].
Structural Constraints. Structural interpretability constraints such as hierarchical interactions, have been studied for both linear settings [3, 30, 59, 17] and nonparametric settings [45, 60, 10, 23]. We briefly review prior work on nonparametric hierarchical interactions as it relates to this paper. [60] proposed GAMI-Net, which is a multi-stage neural-based approach that fits all main effects and a subset of Top- $k$  interaction effects, selected via a fast interaction screening method [34]. Amongst this screened set, interactions that satisfy the weak hierarchy are later used to fit interaction effects. They also prune some main and interaction effects after training based on a variation-based ranking measure. [10] proposed SIAN, which uses Archipelago [54] to measure the strength of each main and interaction effect from a trained DNN, and then screens (i.e., selects a subset of candidate main and interaction effects) using Archipelago scores to identify main effects and interaction effects that obey strong hierarchy. Then, it fits a GAM model with the screened main and interaction effects. [60] and [10] only support screening of interactions obeying hierarchy before the training for interaction effects is done. None of these approaches impose hierarchy while training with interactions. [23] with their ELAAN framework also consider strong hierarchy in the presence of  $\ell_0$ -regularized formulation with splines for the least squares loss (regression). ELAAN has a two-stage approach: It selects a candidate set of interactions and then applies commercial mixed integer programming solvers to learn sparse interactions under a hierarchy constraint. This approach would require customized algorithms to adapt to different loss objectives e.g., multiclass classification. To our knowledge, current techniques for sparse hierarchical (nonparametric) interactions are not based on end-to-end differentiable training: they can be limited in flexibility and scalability—a gap we intend to fill in this work.
Table 1: Relevant work on sparse GAMs with Interactions. Models in rows 1-2 have some variable selection but no hierarchy; models in row 3-5 have screening-based approaches for hierarchy.
Paper	Selection Method	StructuralConstraints	Reg.	Classif.	ShapeFunctions	StatisticalProperties	Scalable
Main	Interactions	Weak Hier.	Strong Hier.	Bin.	Multi
EBM [40]	None	Greedy	X	X	✓	✓	✓	trees	X	✓
NODE-GA2M [6]	Entmax+Anneal	X	X	✓	✓	✓	trees	X	✓
GAMI-Net [60]	Prune	Screening	Screening	X	✓	✓	X	neural	X	X
SIAN [10]	None	Screening	X	Screening	✓	✓	X	neural	✓	X
ELAAN [23]	Group L0	X	Screening+Convex Relax.	✓	X	X	splines	✓	✓
GRAND-SLAMIN	Binary Variables	End-to-end	End-to-end	✓	✓	✓	trees	✓	✓
Hier.  $=$  Hierarchy, Reg.  $=$  Regression, Classif.  $=$  Classification, Bin.  $=$  Binary, Relax.  $=$  Relaxation
Note [9] and [10] also consider higher-order interactions (beyond two-way ones), which can be hard to interpret. For convenience, Table 1 summarizes some relevant work on Sparse GAMs with interactions and possible structural constraints.
We first present in Sec. 3.1 an alternative formulation of GAMs with interactions using binary variables for imposing sparsity and structural constraints. Next, in Sec. 3.2, we present a smooth reformulation of the objective that can be solved with first-order gradient-based methods.
3.1 An optimization formulation with binary variables
We first present an alternative formulation of GAMs with interactions under sparsity with/without additional structured hierarchy constraints. Let us consider the parameterization:
$$
f = \sum_ {j \in [ p ]} f _ {j} \left(x _ {j}\right) z _ {j} + \sum_ {(j, k) \in \mathcal {I}} f _ {j, k} \left(x _ {j}, x _ {k}\right) q \left(z _ {j}, z _ {k}, z _ {j, k}\right), \tag {4}
$$
with main effects  $f_{j}(\cdot)$ , interaction effects  $f_{j,k}(\cdot)$  and binary gates  $z_{j}$  and  $q(z_{j},z_{k},z_{j,k})$ . We consider three different parameterizations for  $q(\cdot)$ , satisfying the following different constraints:
No structural constraint:  $q(z_{j},z_{k},z_{j,k})\stackrel {\mathrm{def}}{=}z_{j,k},$  (5)
Weak hierarchy:  $q(z_{j},z_{k},z_{j,k})\stackrel {\mathrm{def}}{=}(z_{j} + z_{j} - z_{j}z_{k})z_{j,k},$  (6)
Strong hierarchy:  $q(z_{j},z_{k},z_{j,k})\stackrel {\mathrm{def}}{=}z_{j}z_{k}z_{j,k}.$  (7)
The binary gates  $z_{j}$  and  $q(z_{j},z_{k},z_{j,k})$  play the role of selection. In particular, when  $z_{j} = 0$ , the corresponding  $j$ -th main effect  $f_{j}(\cdot)$  is excluded from our additive model (4). Similarly, when  $q(z_{j},z_{k},z_{j,k}) = 0$ , the corresponding  $(j,k)$ -th interaction effect  $f_{j,k}(\cdot)$  is excluded. Then, we can formulate the regularized objective as:
$$
\min  _ {\left\{f _ {j} \right\}, \left\{f _ {j, k} \right\}, \left\{z _ {j} \right\} \in \{0, 1 \} ^ {p}, \left\{z _ {j, k} \right\} \in \{0, 1 \} ^ {| \mathcal {I} |}} \hat {\mathbf {E}} [ \ell (y, f) ] + \lambda \Big (\sum_ {j \in [ p ]} z _ {j} + \alpha \sum_ {(j, k) \in \mathcal {I}} z _ {j, k} \Big), \tag {8}
$$
where the first term denotes empirical loss over the training data, the penalty term controls model sparsity:  $\lambda \geq 0$  is the selection penalty,  $\alpha \in [1,\infty)$  controls the relative selection strength of main and interaction effects. We refer to the framework in (8) under the different constraints (5)-(7) as GRAND-SLAMIN<sup>1</sup>. We discuss extension of this framework to third-order interactions in Supplement Sec. D. However, we do not consider third-order interactions in our experiments as third-order interactions are hard to interpret.
The formulation in (8) with binary variables  $z_{j}, z_{j,k}$  and functions  $f_{j}, f_{j,k}$  with any of the constraint sets (5)-(7) is a challenging discrete optimization problem (with binary variables) and is not amenable to differentiable training via SGD (for example). Sec. 3.2 explores approximate solutions to (8) using a smooth reformulation of the binary variables. Intuitively, we rely on continuous relaxations of the binary variables  $z$ 's and parameterize  $f$ 's with smooth tree-based shape functions. The reformulation allows us to use first-order methods.
3.2 A Smooth Reformulation of Problem (8)
We discuss a smooth reformulation of the objective in (8). We describe an approach to parameterize the continuous relaxation of the binary variables  $z_{j}, z_{j,k}$  with a Smooth-Step function [18] and use smooth tree-based shape functions to model  $\{f_i\}$ ,  $\{f_{j,k}\}$ .
3.2.1 Relaxing Binary Variables with Smooth Gates
We present an approach to smooth the binary gates  $z$ 's in (8) using a smooth-step function [18], which we define next.
1GRAND-SLAMIN stands for GeNeRAlizeD Sparse Learning of Additive Models with INteractions.
Smooth-Step Function. Smooth-step function is a continuously differentiable function, similar in shape to the logistic function. However, unlike the logistic function, the smooth-step function can output 0 and 1 exactly for sufficiently large magnitudes of the input (see Appendix B for details). This function has been used for smoothing binary representations for conditional computation [18, 19, 21].
We parameterize each of the  $z$ 's in (8) as  $S(\mu)$ , where  $\mu \in \mathbb{R}$  is a learnable parameter and  $S(\cdot)$  denotes the Smooth-step function. We parameterize the additive function as:  $f = \sum_{j \in [p]} f_j(x_j) S(\mu_j) + \sum_{(j,k) \in \mathcal{I}} f_{j,k}(x_j, x_k) q(S(\mu_j), S(\mu_k), S(\mu_{j,k}))$  and optimize the following objective:
$$
\min  _ {\left\{f _ {j} \right\}, \left\{f _ {j, k} \right\}, \left. \{\mu_ {j} \} \in \mathbb {R} ^ {p}, \{\mu_ {j, k} \} \in \mathbb {R} ^ {| \mathcal {I} |} \right.} \hat {\mathbb {E}} [ \ell (y, f) ] + \lambda (\sum_ {j \in [ p ]} S (\mu_ {j}) + \alpha \sum_ {(j, k) \in \mathcal {I}} S (\mu_ {j, k})). \tag {9}
$$
Note that  $S(\mu_j)$  and  $S(\mu_{j,k})$  are continuously differentiable, so the formulation in (9) is amenable to first-order gradient-based methods (e.g., SGD).
Achieving binary gates. To encourage each of the  $S(\mu_j)$ 's and  $S(\mu_{j,k})$ 's to achieve binary state (and, not fractional) by the end of training, we add an entropy regularizer  $\tau (\sum_{j\in [p]}\Omega (S(\mu_j)) + \sum_{(j,k)\in \mathcal{I}}\Omega (S(\mu_{j,k})))$  where  $\Omega (S(\mu)) = -(S(\mu)\log S(\mu) + (1 - S(\mu))\log (1 - S(\mu)))$  and  $\tau \geq 0$  controls how quickly each of the gates  $S(\mu)$  converges to a binary  $z$ .
3.2.2 Soft trees
We use soft trees [25, 24, 53, 11]—based on hyperplane splits (univariate or bivariate) and constant leaf nodes — as shape functions to parameterize the main effects  $f_{j}(\cdot)$  and the pairwise interaction effects  $f_{j,k}(\cdot)$ . See Figure 1 for an illustration. Soft trees were introduced as hierarchical mixture of experts by [25] and further developed by [24, 53, 11]. They allow for end-to-end learning [27, 18, 22]. They also have efficient implementations when learning tree ensembles [22]. A detailed definition of soft tress is given in Appendix A.
4 Efficient Implementation
We discuss a fast implementation of our approach. The key elements are: (i) Tensor parameterization of trees, (ii) Sparse backpropagation, and (iii) Complementary screening heuristics.
Tensor Parameterization of Additive Effects.
Typically, in neural-based additive model implementations [1, 60], a separate network module is constructed for each shape function. The outputs from each model are sequentially computed and combined additively. This can create bottleneck in scaling these models. Some recent approaches e.g., [10] try to work around this ap
proach by constructing a large block-wise network to compute representations of all shape functions simultaneously. However, this comes at a cost of large memory footprint. For tree-based shape functions, drawing inspiration from [22], we can implement a tensor-based formulation of all shape functions, leveraging the fact that each tree has the same depth. This parameterization can exploit GPU-friendly parallelization in computing all shape functions simultaneously without increasing memory footprints.
Sparse Backpropagation. We use first-order optimization methods (e.g., SGD and its variants) to optimize GRAND-SLAMIN. Typically, a main computational bottleneck in optimizing GAMs with all pairwise interactions via standard gradient-based backpropagation methods is the computation
Figure 1: Modeling main and interaction effects with soft trees.  $\phi$  denotes the sigmoid activation function. We omit biases in split nodes for brevity. For interaction effect,  $W_{i,1}^{j,k}\in \mathbb{R}$  and  $W_{i,2}^{j,k}\in \mathbb{R}$  denote the weights in  $i$ -th node of the  $(j,k)$ -th tree.
of forward pass and gradient computations with respect to all additive components (both main and interaction effects). This can hinder training large GAMs with interactions. We exploit the sparsity in GRAND-SLAMIN via the sparsity in the smooth-step function and its gradient during training.
Recall that  $S(\mu_j)$ 's and  $S(\mu_{jk})$ 's play a role of selection in a smoothed fashion. In the early stages of training,  $S(\mu_j)$ 's and  $S(\mu_{jk})$ 's are all in the range (0, 1). As the optimization proceeds, due to the entropic regularization and selection regularization,  $S(\mu_j)$ 's and  $S(\mu_{jk})$ 's progressively achieve binary state  $\{0, 1\}$  — the gradient with respect to  $\mu_j$  and  $\mu_{jk}$  also reaches 0 because of the nature of smooth-step function  $S(\cdot)$ . All the additive components corresponding to the selection variables that reached 0 can be removed from both the forward and the backward computational graph. This sparse backpropagation approach can provide large speedups on training times up to a factor of  $10 \times$  over standard backpropagation. Additionally, there is progressively a reduced memory footprint during the course of training in comparison to standard backpropagation.
The approach outlined above (use of Smooth-Step function for selection) when specialized to additive models allows us to implement GPU-friendly sparse backpropagation — this makes our work different from [18], which does not support GPU training.