NODE-GAM: Neural Generalized Additive Model for Interpretable Deep Learning

Deployment of machine learning models in real high-risk settings ( e.g . healthcare ) often depends not only on model ’ s accuracy but also on its fairness , robustness and interpretability . Generalized Additive Models ( GAMs ) are a class of interpretable models with a long history of use in these high-risk domains , but they lack desirable features of deep learning such as differentiability and scalability . In this work , we propose a neural GAM ( NODE-GAM ) and neural GA2M ( NODE-GA2M ) that scale well and perform better than other GAMs on large datasets , while remaining interpretable compared to other ensemble and deep learning models . We demonstrate that our models find interesting patterns in the data . Lastly , we show that we improve model accuracy via self-supervised pre-training , an improvement that is not possible for non-differentiable GAMs . 1 INTRODUCTION . As machine learning models become increasingly adopted in everyday life , we begin to require models to not just be accurate , but also satisfy other constraints such as fairness , bias discovery , and robustness under distribution shifts for high-stakes decisions ( e.g. , in healthcare , finance and criminal justice ) . These needs call for an easier ability to inspect and understand a model ’ s predictions . Generalized Additive Models ( GAMs ) ( Hastie & Tibshirani , 1990 ) have a long history of being used to detect and understand tabular data patterns in a variety of fields including medicine ( Hastie & Tibshirani , 1995 ; Izadi , 2020 ) , business ( Sapra , 2013 ) and ecology ( Pedersen et al. , 2019 ) . Recently proposed tree-based GAMs and GA2Ms models ( Lou et al. , 2013 ) further improve on original GAMs ( Spline ) having higher accuracy and better ability to discover data patterns ( Caruana et al. , 2015 ) . These models are increasingly used to detect dataset bias ( Chang et al. , 2021 ) or audit black-box models ( Tan et al. , 2018a ; b ) . As a powerful class of models , they still lack some desirable features of deep learning that made these models popular and effective , such as differentiability and scalability . In this work , we propose a deep learning version of GAM and GA2M that enjoy the benefits of both worlds . Our models are comparable to other deep learning approaches in performance on tabular data while remaining interpretable . Compared to other GAMs , our models can be optimized using GPUs and mini-batch training allowing for higher accuracy and more effective scaling on larger datasets . We also show that our models improve performance when labeled data is limited by self-supervised pretraining and finetuning , where other non-differentiable GAMs can not be applied . Several works have focused on building interpretable deep learning models that are effective for tabular data . TabNet ( Arik & Pfister , 2020 ) achieves state-of-the-art performance on tabular data while also providing feature importance per example by its attention mechanism . Although attention seems to be correlated with input importance ( Xu et al. , 2015 ) , in the worst case they might not correlate well ( Wiegreffe & Pinter , 2019 ) . Yoon et al . ( 2020 ) proposes to use self-supervised learning on tabular data and achieves state-of-the-art performance but does not address interpretability . NIT ( Tsang et al. , 2018 ) focuses on building a neural network that produces at most K-order interactions and thus include GAM and GA2M . However , NIT requires a two-stage iterative training process that requires longer computations . And their performance is slightly lower to DNNs while ours are overall on par with it . They also do not perform purification that makes GA2M graphs unique when showing them . The most relevant approaches to our work are NODE ( Popov et al. , 2019 ) and NAM ( Agarwal et al. , 2020 ) . Popov et al . ( 2019 ) developed NODE that mimics an ensemble of decision trees but permits differentiability and achieves state-of-the-art performance on tabular data . Unfortunately , NODE suffers from a lack of interpretability similarly to other ensemble and deep learning models . On the other hand , Neural Additive Model ( NAM ) whose deep learning architecture is a GAM , similar to our proposal , thus assuring interpretability . However , NAM can not model the pairwise interactions and thus do not allow GA2M . Also , because NAM builds a small feedforward net per feature , in high-dimensional datasets NAM may require large memory and computation . Finally , NAM requires training of 10s to 100s of models and ensemble them which incurs large computations and memory , while ours only trains once ; our model is also better than NAM without the ensemble ( Supp . A ) . To make our deep GAM scalable and effective , we modify NODE architecture ( Popov et al. , 2019 ) to be a GAM and GA2M , since NODE achieves state-of-the-art performance on tabular data , and its tree-like nature allows GAM to learn quick , non-linear jumps that better match patterns seen in real data ( Chang et al. , 2021 ) . We thus call our models NODE-GAM and NODE-GA2M respectively . One of our key contributions is that we design several novel gating mechanisms that gradually reduce higher-order feature interactions learned in the representation . This also enables our NODE-GAM and NODE-GA2M to automatically perform feature selection via back-propagation for both marginal and pairwise features . This is a substantial improvement on tree-based GA2M that requires an additional algorithm to select which set of pairwise feature interactions to learn ( Lou et al. , 2013 ) . Overall , our contributions can be summarized as follows : • Novel architectures for neural GAM and GA2M thus creating interpretable deep learning models . • Compared to state-of-the-art GAM methods , our NODE-GAM and NODE-GA2M achieve similar performance on medium-sized datasets while outperforming other GAMs on larger datasets . • We demonstrate that NODE-GAM and NODE-GA2M discover interesting data patterns . • Lastly , we show that NODE-GAM benefits from self-supervised learning that improves performance when labeled data is limited , and performs better than other GAMs . We foresee our novel deep learning formulation of the GAMs to be very useful in high-risk domains , such as healthcare , where GAMs have already proved to be useful but stopped short from being applied to new large data collections due to scalability or accuracy issues , as well as settings where access to labeled data is limited . Our novel approach also benefits the deep learning community by adding high accuracy interpretable models to the deep learning repertoire . 2 BACKGROUND . GAM and GA2M : GAMs and GA2Ms are interpretable by design because of their functional forms . Given an input x ∈ RD , a label y , a link function g ( e.g . g is log p1−p in binary classification ) , main effects fj for each feature j , and feature interactions fjj′ , GAM and GA2M are expressed as : GAM : g ( y ) = f0 + D∑ j=1 fj ( xj ) , GA2M : g ( y ) = f0 + D∑ j=1 fj ( xj ) + D∑ j=1 ∑ j′ > j fjj′ ( xj , xj′ ) . Unlike full complexity models ( e.g . DNNs ) that have y = f ( x1 , ... , xj ) , GAMs and GA2M are interpretable because the impact of each feature fj and each feature interaction fjj′ can be visualized as a graph ( i.e . for fj , x-axis shows xj and y-axis shows fj ( xj ) ) . Humans can easily simulate how they work by reading fjs and fjj′ off different features from the graph and adding them together . GAM baselines : We compare with Explainable Boosting Machine ( EBM ) ( Nori et al. , 2019 ) that implements tree-based GAM and GA2M . We also compare with splines proposed in the 80s ( Hastie & Tibshirani , 1990 ) using Cubic splines in pygam package ( Servén & Brummitt , 2018 ) . Neural Oblivious Decision Trees ( NODEs ) : We describe NODEs for completeness and refer the readers to Popov et al . ( 2019 ) for more details . NODE consists of L layers where each layer has m differentiable oblivious decision trees ( ODT ) of equal depth C. Below we describe a single ODT . Differentiable Oblivious Decision Trees : An ODT works like a traditional decision tree except for all nodes in the same depth share the same input features and thresholds , which allows parallel computation and makes it suitable for deep learning . Specifically , an ODT of depth C compares C chosen input feature to C thresholds , and returns one of the 2C possible responses . Mathmatically , for feature functions F c which choose what features to split , splitting thresholds bc , and a response vector R ∈ R2C , the tree output h ( x ) is defined as : h ( x ) = R · ( [ I ( F 1 ( x ) ≤ b1 ) I ( F 1 ( x ) > b1 ) ] ⊗ [ I ( F 2 ( x ) ≤ b2 ) I ( F 2 ( x ) > b2 ) ] ⊗ · · · ⊗ [ I ( FC ( x ) ≤ bC ) I ( FC ( x ) > bC ) ] ) ( 1 ) Here I is the indicator function , ⊗ is the outer product , and · is the inner product . Both feature functions F c and I prevent differentiability . To make them differentiable , Popov et al . ( 2019 ) replace F c ( x ) as a weighted sum of features : F c ( x ) = D∑ j=1 xjentmaxα ( F c ) j = x · entmaxα ( F c ) . ( 2 ) Here F c ∈ RD are the logits for which features to choose , and entmaxα ( Peters et al. , 2019 ) is the entmax transformation which works like a sparse version of softmax such that the sum of the output equals to 1 . They also replace the I with entmoid which works like a sparse sigmoid that has output values between 0 and 1 . Since all operations are differentiable ( entmax , entmoid , outer and inner products ) , the ODT is differentiable . Stacking trees into deep layers : Popov et al . ( 2019 ) follow the design similar to DenseNet where all tree outputs h ( x ) from previous layers ( each layer consists of total I trees ) become the inputs to the next layer . For input features x , the inputs xl to each layer l becomes : x1 = x , xl = [ x , h1 ( x1 ) , ... , h ( l−1 ) ( x ( l−1 ) ) ] for l > 1 . ( 3 ) And the final output of the model ŷ ( x ) is the average of all tree outputs h1 , ... , hL of all L layers : ŷ ( x ) = 1 LI L∑ l=1 I∑ i=1 hli ( x l ) ( 4 ) 3 OUR MODEL DESIGN . GAM design : See Supp . B for a complete pseudo code . To make NODE a GAM , we make three key changes to avoid any feature interactions in the architecture ( Fig . 1 ) . First , instead of letting F c ( x ) be a weighted sum of features ( Eq . 2 ) , we make it only pick 1 feature . We introduce a temperature annealing parameter T that linearly decreases from 1 to 0 for the first S learning steps to make entmaxα ( F c/T ) gradually become one-hot : F c ( x ) = x · entmaxα ( F c/T ) , T S steps−−−−→ 0 . ( 5 ) Second , within each tree , we make the logits F c the same across depth C i.e . F 1 = · · · = FC = F to avoid any feature interaction within a tree . Third , we avoid the DenseNet connection between two trees that focus on different features j , j′ , since they create feature interactions between features j and j′ if two trees connect . Thus we introduce a gate that only allows connections between trees that take the same features . Let Gi = entmaxα ( Fi/T ) of the tree i . For tree i in layer l and another tree î in layer l̂ for l̂ < l , the gating weight gl̂ii and the feature function Fli for tree i become : gl̂ii = Gî ·Gi , Fli ( x ) = x ·Gi + l−1∑ l̂=1 I∑ î=1 hl̂̂i ( x ) gl̂ii . ( 6 ) Since G becomes gradually one-hot by Eq . 5 , after S steps gîi would only become 1 when Gî = Gi and 0 otherwise . This enforces no feature interaction between tree connections . Attention-based GAMs ( AB-GAMs ) : To make the above GAM more expressive , we add an attention weight al̂ii in the feature function Fli ( x ) to decide which previous tree to focus on : Fli ( x ) = D∑ j=1 xjGij + l−1∑ l̂=1 I∑ î=1 hl̂̂i ( x ) gl̂iial̂ii where l−1∑ l̂=1 I∑ î=1 gl̂iial̂ii = 1 . ( 7 ) To achieve this , we introduce attention logits Ali for each tree i that after entmax it produces al̂ii : al̂ii = gl̂iientmaxα ( log ( gi ) +Ali ) î . ( 8 ) This forces the attention of a tree i that ∑ î al̂ii = 1 for all î that gl̂ii = 1 and al̂ii = 0 when gl̂ii = 0 . The attention logits A requires a large matrix size [ I , ( l−1 ) I ] for each layer l > 1 which explodes the memory . We instead make A as the inner product of two smaller matrices such that A = BC where B is of size [ I , E ] and C is of size [ E , ( l − 1 ) I ] , where E is a hyperparameter for the embedding dimension of the attention . Last Linear layer : Lastly , instead of averaging the outputs of all trees as the output of the model ( Eq . 4 ) , we add the last linear layer to be a weighted sum of all outputs : ŷ ( x ) = L∑ l=1 I∑ i=1 hli ( xl ) wli . ( 9 ) Note that in self-supervised learning , wli has multiple output heads to predict multiple tasks . Regularization : We also include other changes that improves performance . First , we add Dropout ( rate p1 ) on the outputs of trees hli ( xl ) , and Dropout ( rate p2 ) on the final weights wli . Also , to increase diversity of trees , each tree can only model on a random subset of features ( η ) , an idea similar to Random Forest . We also add an ` 2 penalization ( λ ) on hli ( xl ) . In binary classification task where labels y are imbalanced between class 0 and 1 , we set a constant as log p ( y ) 1−p ( y ) that is added to the final output of the model such that after sigmoid it becomes the p ( y ) if the output of the model is 0 . We find it ’ s crucial for ` 2 penalization to work since ` 2 induces the model to output 0 . NODE-GA2Ms — extending NODE-GAMs to two-way interactions : To allow two-way interactions , for each tree we introduce two logits F 1 and F 2 instead of just one , and let F c = F ( c−1 ) mod 2+1 for c > 2 ; this allows at most 2 features to interact within each tree ( Fig . 7 ) . Besides temperature annealing ( Eq . 5 ) , we make the gating weights gîi = 1 only if the combination of F 1 , F 2 is the same between tree î and i ( i.e . both trees î and i focus on the same 2 features ) . We set gîi as : gîi = min ( ( G 1 i ·G1î ) × ( G 2 i ·G2î ) + ( G 1 i ·G2î ) × ( G 2 i ·G1î ) , 1 ) . ( 10 ) We cap the value at 1 to avoid uneven amplifications as gîi = 2 when G 1 i = G 2 i = G 1 î = G2 î . Data Preprocessing and Hyperparameters : We follow Popov et al . ( 2019 ) to do target encoding for categorical features , and do quantile transform for all features to Gaussian distribution ( we find Gaussian works better than Uniform ) . We use random search to search the architecture space for NODE , NODE-GAM and NODE-GA2M . We use QHAdam ( Ma & Yarats , 2018 ) and average the most recent 5 checkpoints ( Izmailov et al. , 2018 ) . In addition , we adopt learning rate warmup ( Goyal et al. , 2017 ) , and do early stopping and learning rate decay on the plateau . More details in Supp . F. Extracting shape graphs from GAMs : We follow Chang et al . ( 2021 ) to implement a function that extracts main effects fj from any GAM model including NODE-GAM , Spline and EBM . The main idea is to take the difference between the model ’ s outputs of two examples ( x1 , x2 ) that have the same values except for feature j . Since the intercept and other main effects are canceled out when taking the difference , the difference f ( x2 ) − f ( x1 ) is equal to fj ( x2j ) − fj ( x1j ) . If we query all the unique values of xj , we get all values of fj relative to fj ( x1j ) . Then we center the graph of fj by setting the average of fj ( xj ) across the dataset as 0 and add the average to the intercept term f0 . Extracting shape graphs from GA2Ms : Designing a black box function to extract from any GA2M is non-trivial , as each changed feature xj would change not just main effect term fj but also every interactions ∀j′fjj′ that involve feature j . Instead , since we know which features each tree takes , we can aggregate the output of trees into corresponding main fj and interaction terms fjj′ . Note that GA2M can have many representations that result in the same function . For example , for a prediction value v associated with x2 , we can move v to the main effect f2 ( x2 ) = v , or the interaction effect f23 ( x2 , · ) = v that involves x2 . To solve this ambiguity , we adopt `` purification '' ( Lengerich et al. , 2020 ) that pushes interaction effects into main effects if possible . See Supp . C for details .

Summary:  - The authors develop the NODE-GAM and NODE-GA$^2$M architectures by modifying the previously-developed NODE architecture with constraints and a gating mechanism to ensure that, both within each tree and across layers, the model is only allowed to learn feature interactions of order 1 (for NODE-GAM) or 2 for (NODE-GA$^2$M). - They show that their proposed NODE-GAM and NODE-GA$^2$M architectures perform comparably to the baselines of Explainable Boosting Machines, tree-based GAM/GA$^2$M and the traditional spline-based GAMs on medium-sized datasets while outperforming them on larger datasets. They also show that NODE-GAM can benefit from self-supervised learning by pretraining a model to reconstruct the original input from a masked input. - They apply NODE-GAM and NODE-GA$^2$M to real-world datasets and showcase the patterns uncovered.