1]Microsoft Research

2]University of California, Berkeley

# Augmenting Interpretable Models with LLMs during Training

Chandan Singh<EMAIL_ADDRESS>Armin Askari<EMAIL_ADDRESS>Rich
Caruana<EMAIL_ADDRESS>Jianfeng Gao<EMAIL_ADDRESS>[  [

###### Abstract

Recent large language models (LLMs) have demonstrated remarkable prediction
performance for a growing array of tasks. However, their proliferation into
high-stakes domains (e.g. medicine) and compute-limited settings has created a
burgeoning need for interpretability and efficiency. We address this need by
proposing Augmented Interpretable Models (Aug-imodels), a framework for
leveraging the knowledge learned by LLMs to build extremely efficient and
interpretable models. Aug-imodels use LLMs during fitting but not during
inference, allowing complete transparency and often a speed/memory improvement
of greater than 1,000x for inference compared to LLMs. We explore two
instantiations of Aug-imodels in natural-language processing: (i) Aug-GAM,
which augments a generalized additive model with decoupled embeddings from an
LLM and (ii) Aug-Tree, which augments a decision tree with LLM feature
expansions. Across a variety of text-classification datasets, both outperform
their non-augmented counterparts. Aug-GAM can even outperform much larger
models (e.g. a 6-billion parameter GPT-J model), despite having 10,000x fewer
parameters and being fully transparent. We further explore Aug-imodels in a
natural-language fMRI study, where they generate interesting interpretations
from scientific data. All code for using Aug-imodels and reproducing results
is made available on Github.111Scikit-learn-compatible API available at
github.com/csinva/imodelsX and experiments code available at
github.com/microsoft/augmented-interpretable-models.

###### keywords:

Explainability, Interpretability, Transparent models, XAI, Large language
models

##  1 Introduction

Large language models (LLMs) have demonstrated remarkable predictive
performance across a growing range of diverse tasks [1, 2, 3]. However, their
proliferation has led to two burgeoning problems. First, like most deep neural
nets, LLMs have become increasingly difficult to interpret, often leading to
them being characterized as black boxes and debilitating their use in high-
stakes applications such as science [4], medicine [5], and policy-making [6].
Moreover, the use of black-box models such as LLMs has come under increasing
scrutiny in settings where users require explanations or where models struggle
with issues such as fairness [7] and regulatory pressure [8]. Second, black-
box LLMs have grown to massive sizes, incurring enormous energy costs [9] and
making them costly and difficult to deploy, particularly in low-compute
settings (e.g. edge devices).

As an alternative to large black-box models, transparent models, such as
generalized additive models [10] and decision trees [11] can maintain complete
interpretability. Additionally, transparent models tend to be dramatically
more computationally efficient than LLMs. While transparent models can
sometimes perform as well as black-box LLMs [12, 13, 14, 15], in many settings
(such as natural-language processing (NLP)), there is often a considerable gap
between the performance of transparent models and black-box LLMs.

Figure 1:  Aug-imodels use an LLM to augment an interpretable model during
fitting but not inference (toy example for movie-review classification). (A)
Aug-GAM fits an augmented additive model by extracting fixed-size embeddings
for decoupled ngrams in a given sequence, summing them, and using them to
train a supervised linear model. (B) At test time, Aug-GAM can be interpreted
exactly as a generalized additive model. A linear coefficient for each ngram
in the input is obtained by taking the dot product between the ngram’s
embedding and the shared vector $w$. (C) Aug-Tree improves each split of a
decision tree during fitting by (D) augmenting each keyphrase found by CART
with similar keyphrases generated by an LLM.

We address this gap by proposing Augmented Interpretable Models (Aug-imodels),
a framework to leverage the knowledge learned by LLMs to build extremely
efficient and interpretable models. Specifically, we define an Aug-imodel as a
method that leverages an LLM to fit an interpretable model, but _does not use
the LLM during inference_. This allows complete transparency and often a
substantial efficiency improvement (both in terms of speed and memory). Aug-
imodels can address shortcomings in existing transparent models by using the
world knowledge present in modern LLMs, such as information about feature
correlations.

We explore two instantiations of Aug-imodels: (i) Aug-GAM, which augments a
generalized additive model via decoupled embeddings from an LLM and (ii) Aug-
Tree, which augments a decision tree with improved features generated by
calling an LLM (see Fig 1). At inference time, both are completely transparent
and efficient: Aug-GAM requires only summing coefficients from a fixed
dictionary while Aug-Tree requires checking for the presence of keyphrases in
an input.

Across a variety of natural-language-processing datasets, both proposed Aug-
imodels outperform their non-augmented counterparts. Aug-GAM can even
outperform much larger models (e.g. a 6-billion parameter GPT-J model [16]),
despite having 10,000x fewer parameters and no nonlinearities. We further
explore Aug-imodels in a natural-language fMRI context, where we find that
they can predict well and generate interesting interpretations. In what
follows, Sec 2 formally introduces Aug-imodels, Sec 3 and Sec 4 shows results
for predictive performance and interpretation, Sec 5 explores Aug-imodels in
an fMRI prediction setting, Sec 6 reviews related work, and Sec 7 concludes
with a discussion.

##  2 Aug-imodels methodology: Aug-GAM and Aug-Tree

In this section, Sec 2.1 overviews limitations of existing transparent
methods, Sec 2.2 introduces Aug-GAM, and Sec 2.3 introduces Aug-Tree.

###  2.1 Limitations of existing transparent methods

We are given a dataset of $n$ natural language strings $X_{\text{text}}$ and
corresponding labels $\mathbf{y}\in\mathbb{R}^{n}$. In transparent modeling,
often each string $x$ is represented by a bag-of-words, in which each feature
$x_{i}$ is a binary indicator (or count) of the presence of a single token
(e.g. the word good). To model interactions between tokens, one can instead
use a bag-of-ngrams representation, whereby each feature is formed by
concatenating multiple tokens (e.g. the phrase not good). Using a bag-of-
ngrams representation maps $X_{\text{text}}$ to a feature matrix
$X\in\mathbb{R}^{n\times p}$, where $p$ is the number of unique ngrams in
$X_{\text{text}}$. While this representation enables interpretability, the
number of ngrams in a dataset grows exponentially with the size of the ngram
(how many tokens it contains) and the vocab-size; even for a modest vocab-size
of 10,000 tokens, the number of possible trigrams is already $10^{12}$. This
makes it difficult for existing transparent methods to model all trigrams
without overfitting. Moreover, existing transparent methods completely fail to
learn about ngrams not seen in the training set.

#### Preliminaries: GAMs

Generalized additive models, or GAMs [10] take the form:

$\displaystyle
g(\mathbb{E}[y])=\beta+f_{1}\left(x_{1}\right)+f_{2}\left(x_{2}\right)+\cdots+f_{K}\left(x_{p}\right),$
(1)

where $\left(x_{1},x_{2},\ldots,x_{p}\right)$ are the input features (i.e.
ngrams), $y$ is the target variable, $g(\cdot)$ is the link function (e.g.,
logistic function) and each $f_{i}$ is a univariate shape function with
$\mathbb{E}\left[f_{i}\right]=0$. Due to the function’s additivity, each
component function $f_{i}$ can be interpreted independently. Generalized
linear models, such as logistic regression, are a special form of GAMs where
each $f_{i}$ is restricted to be linear.

#### Preliminaries: decision trees

CART [11] fits a binary decision tree via recursive partitioning. When growing
a tree, CART chooses for each node $\mathfrak{t}$ the split $s$ that maximizes
the impurity decrease in the responses $\mathbf{y}$. For a given node $t$, the
impurity decrease has the expression

$\hat{\Delta}(s,\mathfrak{t},\mathbf{y}):=\sum_{\mathbf{x}_{i}\in\mathfrak{t}}h\left(y_{i},\bar{y}_{\mathfrak{t}}\right)-\sum_{\mathbf{x}_{i}\in\mathfrak{t}_{L}}h\left(y_{i},\bar{y}_{\mathfrak{t}_{L}}\right)-\sum_{\mathbf{x}_{i}\in\mathfrak{t}_{R}}h\left(y_{i},\bar{y}_{\mathfrak{t}_{R}}\right),$
(2)

where $\mathfrak{t}_{L}$ and $\mathfrak{t}_{R}$ denote the left and right
child nodes of $\mathfrak{t}$ respectively, and
$\bar{y}_{\mathfrak{t}},\bar{y}_{\mathfrak{t}_{L}},\bar{y}_{\mathfrak{t}_{R}}$
denote the mean responses in each of the nodes. For classification,
$h(\cdot,\cdot)$ corresponds to the Gini impurity, and for regression,
$h(\cdot,\cdot)$ is the mean-squared error. Each split $s$ is a partition of
the data based on a feature in $X$. To grow the tree, the splitting process is
repeated recursively for each child node until a stopping criteria (e.g. a max
depth) is satisfied. At inference time, we predict the response of an example
by following its path from the root to a leaf and then predicting with the
mean value found in that leaf.

###  2.2 Aug-GAM method description

To remedy the issues with the GAM model in Eq. 1, we propose Aug-GAM, an
intuitive model which leverages a pre-trained LLM to extract a feature
representation $\phi(x_{i})$ for each input ngram $x_{i}$. This allows
learning only a single linear weight vector $w$ with a fixed dimension (which
depends on the embedding dimension produced by the LLM), regardless of the
number of ngrams. As a result, Aug-GAM can learn efficiently as the number of
input features grows, and can also infer coefficients for unseen features. The
fitted model is still a GAM, ensuring that the model can be precisely
interpreted as a linear function of its inputs:

$g(\mathbb{E}[y])=\beta+w^{T}\sum_{i}\phi(x_{i})$ (3)

Fitting Aug-GAM is similar to the popular approach of finetuning a single
linear layer on top of LLM embeddings. However, it requires extra steps that
separately extract/embed each ngram to keep the contributions to the
prediction strictly additive across ngrams (see Fig 1A): (i) Extracting
ngrams: To ensure input ngrams are interpretable, ngrams are constructed using
a word-level tokenizer (here, spaCy [17]). We select the size of ngrams to be
used via cross-validation. (ii) Extracting embeddings: Each ngram is fed
through the LLM to retrieve a fixed-size embedding.222If a transformer returns
a variable-length embedding (e.g. the embedding is the size of the sequence
length), we average over its variable-length dimension. A common alternative
for bi-directional (masked) language models is to use the embedding for a
special token (i.e. [CLS]), but we aim to keep the approach here more general.
(iii) Summing embeddings: The embeddings of each ngram in the input are summed
to yield a single fixed-size vector, ensuring additivity of the final model.
(iv) Fitting the final linear model to make predictions: A linear model is fit
on the summed embedding vector. We choose the link function $g$ to be the
logit function (or the softmax for multi-class) for classification and the
identity function for regression. In both cases, we add $\ell_{2}$
regularization over the parameters $w$ in Eq. 3.

#### Computational considerations

During fitting, Aug-GAM is inexpensive to fit as (1) the pre-trained LLM is
not modified in any way, and can be any existing off-the-shelf model and (2)
Aug-GAM only requires fitting a fixed-size linear model. After training, the
model can be converted to a dictionary of scalar coefficients for each ngram,
where the coefficient is the dot product between the ngram’s embedding and the
fitted weight vector $w$ (Fig 1B). This makes inference extremely fast and
converts the model to have size equal to the number of fitted ngrams. When new
ngrams are encountered at test-time, the coefficients for these ngrams can
optionally be inferred by again taking the dot product between the ngram’s
embedding and the fitted weight vector $w$;

###  2.3 Aug-Tree method description

Aug-Tree improves upon a CART decision tree by augmenting features with
generations from an LLM. This helps capture correlations between ngrams,
including correlations with ngrams that are not present in the training data.
Aug-Tree does not modify the objective in Eq. 2 but rather modifies the
procedure for fitting each individual split $s$ (Fig 1D). While CART restricts
each split to a single ngram, Aug-Tree lets each split fit a disjunction of
ngrams (e.g. ngram1 $\lor$ ngram2 $\lor$ ngram3). The disjunction allows a
split to capture sparse interactions, such as synonyms in natural language.
This can help mitigate overfitting by allowing individual splits to capture
concrete concepts, rather than requiring many interacting splits.

When fitting each split, Aug-Tree starts with the ngram which maximizes the
objective in Eq. 2, just as CART would do, e.g. not good. Then, we query an
LLM to generate similar ngrams to include in the split, e.g. bad, poor, awful,
…, horrendous. Specifically, we query GPT-3 (text-davinci-003) [1] with the
prompt Generate 100 concise phrases that are very similar to the
keyphrase:\nKeyphrase: “{keyphrase}”\n1. and parse the outputs into a list of
ngrams. We greedily screen each ngram by evaluating the impurity of the split
when including the ngram in the disjunction; we then exclude any ngram which
increases the split’s impurity, resulting in a shortened list of ngrams, e.g.
bad, poor, dull. See extended algorithm details in Algorithm B2.

#### Computational considerations

As opposed to Aug-GAM, Aug-Tree uses an LLM API rather than LLM embeddings,
which may be more desirable depending on access to compute. The number of LLM
calls required is proportional to the number of nodes in the tree. During
inference, the LLM is no longer needed and making a prediction simply requires
checking an input for the presence of specific ngrams along one path in the
tree. Storing an Aug-GAM model requires memory proportional to the number of
raw strings associated with tree splits, usually substantially reducing memory
over the already small Aug-GAM model. We explore variations of Aug-Tree (such
as using LLM embeddings rather than an API) in Appendix B.

##  3 Results: Prediction performance

###  3.1 Experimental setup

#### Datasets

Table 1 shows the datasets we study: 4 widely used text classification
datasets spanning different domains (e.g. classifying the emotion of tweets
[18], the sentiment of financial news sentences [19], or the sentiment of
movie reviews [20, 21]), and 1 scientific text regression dataset (described
in Sec 5) [22]. Across datasets, the number of unique ngrams grows quickly
from unigrams to bigrams to trigrams. Moreover, many ngrams appear very
rarely, e.g., in the Financial Phrasebank (FPB) dataset, 91% of trigrams
appear only once in the training dataset.

Table 1: Overview of datasets studied here. The number of ngrams grows quickly with the size of the ngram.  |  FPB |  Rotten tomatoes | SST2 | Emotion | fMRI
---|---|---|---|---|---
Samples (train) | 2,313 | 8,530 | 67,349 | 16,000 | 9,461
Samples (val) | 1,140 | 1,066 | 872 | 2,000 | 291
Classes | 3 | 2 | 2 | 6 | Regression
Unigrams | 7,169 | 16,631 | 13,887 | 15,165 | 4,980
Bigrams | 28,481 | 93,921 | 72,501 | 106,201 | 27,247
Trigrams | 39,597 | 147,426 | 108,800 | 201,404 | 46,834
Trigrams that appear only once |  91% |  93% |  13% |  89% |  71%

#### Aug-GAM settings

We compare Aug-GAM to four interpretable baseline models: Bag of ngrams, TF-
IDF (Term frequency–inverse document frequency) [23], GloVE [24]333We use the
pre-trained Glove embeddings trained on Common Crawl (840 billion tokens, 2.2
million vocab-size, cased, 300-dimensional vectors)., and a model trained on
BERT embeddings for each unigram in the input (which can be viewed as running
Aug-GAM with only unigrams). We use BERT (bert-base-uncased) [3] as the LLM
for extracting embeddings, after finetuning on each dataset.444Pre-trained
language models are retrieved from HuggingFace [25]. See Table 4 for details
on all models and downloadable checkpoints. In each case, a model is fit via
cross-validation on the training set (to tune the amount of $\ell_{2}$
regularization added) and its accuracy is evaluated on the test set.

#### Aug-Tree settings

We compare Aug-Tree to two decision tree baselines: CART [11] and ID3 [26],
and we use bigram features. In addition to individual trees, we fit bagging
ensembles, where each tree is created using a bootstrap sample the same size
as the original dataset (as done in Random Forest [27]) and has depth 8. This
hurts interpretability, but can improve predictive performance and
calibration. For simplicity, we run Aug-GAM only in a binary classification
setting; to do so, we take two opposite classes from each multiclass dataset
(Negative/Positive for FPB and Sadness/Joy for Emotion).

###  3.2 Aug-GAM text-classification performance

#### Generalization as a function of ngram size

Fig 2A shows the test accuracy of Aug-GAM as a function of the ngram size used
for computing features. Aug-GAM outperforms the interpretable baselines,
achieving a considerable increase in accuracy across three of the four
datasets. Notably, Aug-GAM accuracy increases with ngram size, whereas the
accuracy of baseline methods decreases or plateaus. This is likely due to Aug-
GAM fitting only a fixed-size parameter vector, helping to prevent
overfitting.

\begin{overpic}[width=173.44534pt,tics=10]{figs/acc_by_ngrams_full.pdf} \put(2.0,95.0){\large{A}} \end{overpic} | \begin{overpic}[width=195.12767pt,tics=10]{figs/acc_interpolate_gam.pdf} \put(0.0,83.0){\large{B}} \end{overpic}
---|---
Figure 2:  (A) Test accuracy as a function of ngram size. As the ngram size
(i.e. the number of tokens in the ngram) increases, the gap between Aug-GAM
and the baselines grows. Averaged over three random cross-validation splits;
error bars are standard errors of the mean (many are within the points). (B)
Accuracy when using Aug-GAM in combination with BERT. A large percentage of
samples can be accurately predicted with Aug-GAM.

#### Comparing Aug-GAM performance with black-box baselines

Table 2 shows the test accuracy results for various models when choosing the
size of ngrams via cross-validation. Compared with interpretable baselines,
Aug-GAM shows considerable gains on three of the datasets and only a minor
gain on the tweet dataset (Emotion), likely because this dataset requires
fitting less high-order interactions.

Compared with the zero-shot performance of the much larger GPT models
(6-billion parameter GPT-J [16] and 175-billion parameter GPT-3, text-
davinci-002 [1])555Accuracy for GPT models is computed by averaging over
human-written prompts take from PromptSource ([28]); see details in Appendix
A)., Aug-GAM outperforms GPT-J. Aug-GAM lags slightly behind GPT-3 for binary
classification problems (Rotten Tomatoes and SST2) but outperforms GPT-3 for
multi-class classification problems (FPB and Emotion). The best black-box
baseline (a BERT finetuned model) outperforms Aug-GAM by 4%-6% accuracy. This
is potentially a reasonable tradeoff in settings where interpretability,
speed, or memory bottlenecks are critical.

Table 2:  Test accuracy for different models. Aug-GAM yields improvements over interpretable baselines and is competitive with some black-box baselines. Errors show standard error of the mean over 3 random data splits (or 3 different prompts for GPT models).  |  | FPB | Rotten tomatoes | SST2 | Emotion
---|---|---|---|---|---
Ours | Aug-GAM |  92.8 $\pm$ 0.37 |  81.6 $\pm$0.05 |  86.9 $\pm$ 0.10 |  89.5 $\pm$ 0.03
Interpretable baselines | Bag of ngrams |  85.0 $\pm$ 0.11 |  75.0 $\pm$ 0.09 |  82.8 $\pm$ 0.00 |  89.0 $\pm$ 0.09
TF-IDF |  84.9 $\pm$ 0.16 |  75.9 $\pm$ 0.06 |  83.4 $\pm$ 0.11 |  89.2 $\pm$ 0.04
GloVe |  80.5 $\pm$ 0.06 |  78.7 $\pm$ 0.03 |  80.1 $\pm$ 0.10 |  73.1 $\pm$ 0.09
BERT unigram embeddings |  86.4 $\pm$ 0.13 |  76.8 $\pm$ 0.19 |  81.7 $\pm$ 0.07 |  87.2 $\pm$ 0.06
Black-box baselines | BERT finetuned | 98.0 | 87.5 | 92.4 | 93.6
GPT-3 |  39.6 $\pm$1.6 | 82.7 $\pm$3.3 | 90.5 $\pm$3.9 |  45.1 $\pm$4.1
GPT-J |  27.0 $\pm$1.9 |  58.9 $\pm$3.1 |  58.4 $\pm$2.8 |  19.3 $\pm$1.9

#### Complementing Aug-GAM with a black-box model

In some settings, it may be useful to use Aug-GAM on relatively simple samples
(for interpretability/memory/speed) but relegate relatively difficult samples
to a black-box model. To study this setting, we first predict each sample with
Aug-GAM, then assess its confidence (how close its predicted probability for
the top class is to 1). If the confidence is above a pre-specified threshold,
we use the Aug-GAM prediction. Otherwise, we compute the sample’s prediction
using a finetuned BERT model. Fig 2B shows the accuracy for the entire test
set as we vary the percentage of samples predicted with Aug-GAM. Since Aug-GAM
yields probabilities that are reasonably calibrated (see Fig 6), rather than
the accuracy linearly interpolating between Aug-GAM and BERT, a large
percentage of samples can be predicted with Aug-GAM while incurring little to
no drop in accuracy. For example, when using Aug-GAM on 50% of samples, the
average drop in test accuracy is only 0.0053.

#### Tradeoffs between accuracy and efficiency

In cases involving inference memory/speed, Aug-GAM can be converted to a
dictionary of coefficients, whose size is the number of ngrams that appeared
in training (see Table 1). For a trigram model, this yields roughly a 1,000x
reduction in model size compared to the $\sim$110 million trainable parameters
in BERT, with much room for further size reduction (e.g. simply removing
coefficients for trigrams that appear only once yields another 10-fold size
reduction). Inference is nearly instantaneous, as it requires looking up
coefficients in a dictionary and then a single sum (and does not require a
GPU).

Sec A.1 explores accuracy/efficiency tradeoffs. For example, Aug-GAM
performance degrades gracefully when the model is compressed by removing its
smallest coefficients. In fact, the test accuracy of Aug-GAM models trained
with 4-grams on the Emotion and Financial phrasebank datasets actually
improves after removing 50% of the original coefficients (Fig 7A).
Additionally, one can vary the size of ngrams used at test-time without a
severe performance drop, potentially enabling compressing the model by orders
of magnitude (see Fig 7B, Fig 8). For example, when fitting a model with
4-grams and testing with 3-grams, the average performance drop is $\sim$2%.

###  3.3 Aug-Tree generalization performance

We now investigate the predictive performance of Aug-Tree, measured by the
test ROC AUC on the previous text classification datasets altered for binary
classification. Note that the performance of all tree-based methods on the
studied datasets is below the performance of the GAM methods in Sec 3.2 (see
Table 8 for a direct comparison). Nevertheless, Aug-Tree models maintain
potential advantages, such as storing far fewer parameters, clustering
important features together, and better modeling long-range interactions.

Fig 3A shows the performance for Aug-Tree as a function of tree depth compared
to decision tree baselines. Aug-Tree shows improvements that are sometimes
small (e.g. for Financial phrasebank) and sometimes relatively large (e.g. for
Emotion). Fig 3B shows the performance of a bagging ensemble of trees with
different tree methods used as the base estimator. Here, using Aug-Tree shows
a reliable and significant gain across all datasets compared to ensembles of
baseline decision-tree methods. This suggests that LLM augmentation may help
to diversify or decorrelate individual trees in the ensemble.

|  |
---|---|---
(A) Individual tree | (B) Bagging ensemble |
Figure 3: Test performance as a function of (A) tree depth and (B) number of
estimators. Values are averaged over 3 random dataset splits; error bars show
the standard error of the mean (many are within the points).

#### Varying Aug-imodels settings

We investigate many variations in the settings used for Aug-imodels. Table 7
shows variations of different hyperparameters for Aug-Tree, such as using
embeddings or dataset-specific prompts to expand keyphrases. Table 5 shows how
generalization accuracy changes when the LLM used to extract embeddings for
Aug-GAM is varied, or different layers / ngram selection techniques are used.
Across the variations, embeddings from finetuned models yield considerably
better results than embeddings from non-finetuned models.

##  4 Interpreting fitted models

In this section, we interpret fitted Aug-imodels. We first inspect an Aug-GAM
model fitted using unigram and bigram features on the SST2 dataset which
achieves 84% test accuracy. Next, we analyze the keyphrase expansions made in
fitted Aug-Tree bagging ensembles.

#### Fitted Aug-GAM coefficients match human scores

A fitted Aug-GAM model can be interpreted for a single prediction (i.e.
getting a score for each ngram in a single input, as in Fig 1) or for an
entire dataset (i.e. by inspecting its fitted coefficients). Fig 4A visualizes
the fitted Aug-GAM coefficients (i.e. the contribution to the prediction
$w^{T}\phi(x_{i})$) with the largest absolute values across the SST2 dataset.
To show a diversity of ngrams, we show every fifth ngram. The fitted
coefficients are semantically reasonable and many contain strong interactions
(e.g. not very is assigned to be negative whereas without resorting is
assigned to be positive). This form of model visualization allows a user to
audit the model with prior knowledge. Moreover, these coefficients are exact
and therefore avoid summarizing interactions, making them considerably more
faithful than post-hoc methods, such as LIME [29] and SHAP [30] (see Sec A.2
for a comparison).

Fig 4B compares the fitted Aug-GAM coefficients to human-labeled sentiment
phrase scores for unigrams/bigrams in SST (note: these continuous scores are
separate from the binary sentence labels used for training in the SST2
dataset). Both are centered, so that 0 is neutral sentiment and
positive/negative values correspond to positive/negative sentiment,
respectively. There is a strong positive correlation between the coefficients
and the human-labeled scores (Spearman rank correlation $\rho=0.63$), which
considerably improves over coefficients from a bag-of-bigrams model trained on
SST2 ($\rho=0.39$).

\begin{overpic}[width=286.19078pt,tics=10]{figs/ngram-contribution.pdf} \put(1.0,30.0){\large{A}} \end{overpic} |  \begin{overpic}[width=143.09538pt,tics=10]{figs/sst_embgam.png} \put(1.0,66.0){\large{B}} \end{overpic}
---|---
\begin{overpic}[width=286.19078pt,tics=10]{figs/trigram-unseen-contributions.pdf} \put(1.0,35.0){\large{C}} \end{overpic} |  \begin{overpic}[width=143.09538pt,tics=10]{figs/sst_auggam_tri.png} \put(1.0,66.0){\large{D}} \end{overpic}
Figure 4: Top and bottom contributing ngrams to an Aug-GAM model trained on
SST2 bigrams are (A) qualitatively semantically accurate and (B) match human-
labeled phrase sentiment scores. For the same Aug-GAM model, which is trained
only on bigrams, inferred trigrams coefficients are (C) qualitatively
semantically accurate and (D) match human-labeled phrase sentiment scores.

#### Inferred Aug-GAM coefficients for unseen ngrams match human scores

One strength of Aug-GAM is its ability to infer linear coefficients for ngrams
that were not seen during training. Whereas baseline models generally assign
each unknown ngram the same coefficient (e.g. 0), Aug-GAM can effectively
assign these new ngrams accurate coefficients. As one example, Fig 4C shows
that the Aug-GAM model trained only on bigrams in Fig 4A/B can automatically
infer coefficients for trigrams (which were not fit during training). The
inferred coefficients are semantically meaningful, even capturing three-way
interactions, such as not very amusing. To show a diversity of ngrams, we show
every 20th ngram. Fig 4D shows the coefficients compared to the human-labeled
SST phrase sentiment for all trigrams in SST. Again, there is a strong
correlation, where the Aug-GAM coefficients achieves a rank correlation
$\rho=0.71$, which even outperforms the bag-of-words model directly trained on
trigrams ($\rho=0.49$).

#### Aug-Tree augmented splits contain relevant phrases

A fitted Aug-Tree model can be easily interpreted for a single prediction
(i.e. by inspecting the ngrams that triggered relevant splits) or by
visualizing the entire tree (e.g. Fig 1C). Here, we additionally analyze how
well each ngram found by CART matches the augmented ngrams found by the LLM;
the better this match is, the easier it is to interpret a split.

Table 3 shows examples of the ngrams which were most frequently augmented when
fitting a bagging ensemble of 40 Aug-Trees to the four text-classification
datasets in Table 1. Added ngrams seem qualitatively reasonable, e.g. the
keyphrase good expands to fine, highly, solid, …, valuable. We evaluate how
well the expansions match the original CART ngram via human evaluation scores.
Human evaluators are given the original ngram and the added ngrams, then
instructed “You are given a keyphrase along with related keyphrases. On a
scale of 1 (worst) to 5 (best), how well do the related keyphrases match the
example keyphrase?”666Human evaluation scores are averaged over 3 PhD students
in machine learning not affiliated with the study and 15 random ngrams from
each dataset.. Table 3 shows that the average human score for splits in each
dataset is consistently greater than 4. This is substantially higher than the
baseline score of 1.3 assigned by human evaluators when 15 ngrams and
expansions are randomly paired and evaluated. Table 6 gives more details on
ngram expansions.

Table 3: Examples of most frequently augmented ngrams for each dataset when fitting an ensemble of 40 Aug-Trees. Human scores measure the similarity between an ngram and its expansion. They range from 1 (worst match) to 5 (best match), and the baseline score when ngrams and expansions are randomly paired and evaluated is 1.3$\pm$0.1. Error bars show standard error of the mean. Abbreviations: FPB = Financial Phrasebank, RT = Rotten tomatoes.  Dataset |  Human score |  Example CART ngram | Added ngrams
---|---|---|---
SST2 | 4.6$\pm$0.1 | good | fine, highly, solid, worthy, pleasing, satisfactory, outstanding, honorable, unwavering, valuable, …
best | most remarkable, outstanding, superb, flawless, splendid, superlative, exceptional, impeccable, …
RT | 4.4$\pm$0.1 | dull | dreary, uninteresting, lackluster, listless, lifeless, uninspired, wearisome, drab, joylessly, …
bad | unpleasant, dire, despicable, terrible, heinous, disgusting, vile, putrid, atrocious, nasty, poor, …
Emotion | 4.4$\pm$0.2 | miserable | gloomy, disillusioned, pathetic, doomed, agonized, despairing, pointless, despondent, …
sorry | embarrassed, sorrowful, remorseful, excuse, unsatisfied, guilt, regretful, forgive, apologies, …
FPB | 4.2$\pm$0.2 | increased | widened, consolidated
fell | slipped, slumped, diminished, plunged, dropped, weakened, lost ground

##  5 Analyzing fMRI data with Aug-imodels

We now explore Aug-imodels in a real-world neuroscience context. A central
challenge in neuroscience is understanding how and where semantic concepts are
represented in the brain. To meet this challenge, one line of study predicts
the response of different brain voxels (i.e. small regions in space) to
natural-language stimuli. We analyze data from a recent study in which the
authors collect functional MRI (fMRI) responses as human subjects listen to
hours of narrative stories [22]. The fMRI responses studied here contain
95,556 voxels from a single subject, with 9,461 time points used for
training/cross-validation and 291 time points used for testing. We predict the
continuous response for each voxel at each time point using the 20 words that
precede the time point.777The most recent 4 words are skipped due to a time
delay in the fMRI BOLD response. Seminal work on this task found that linear
models of word vectors could effectively predict voxel responses [31], and
more recent work shows that LLMs can further improve predictive performance
[32, 33]. Aug-GAM is well-suited to this task, as it combines low-level word
information with the contextualized information present in higher-order
ngrams, both of which have been found to contribute to fMRI representations of
text [34].

Fig 5A visualizes the voxels in the cortex which are better predicted by Aug-
GAM than BERT. The improvements are often spatially localized within well-
studied brain regions such as auditory cortex (AC). Fig 5B shows that the test
performance for Aug-GAM (measured by the Pearson correlation coefficient
$\rho$) outperforms the black-box BERT baseline. Appendix C gives further data
details and comparisons, e.g. Aug-GAM also outperforms other linear baselines.

Going beyond prediction performance, Fig 5C investigates a simple example of
how Aug-GAM could help interpret an underlying brain region. We first select
the voxel which is best-predicted by Aug-GAM (achieving a test correlation of
0.76) and then visualize the largest fitted Aug-GAM coefficients for that
voxel. These correspond to which ngrams increase the activity of the fMRI
voxel the most. Interestingly, these ngrams qualitatively correspond to
understandable concepts: questioning, e.g. “are you sure”, often combined with
disbelief/incredulity, e.g. “wow I never”. Fig 5D shows two examples of voxels
that are better predicted by Aug-Tree than Aug-GAM (Aug-Tree yields test
correlations of 0.35 and 0.36). These two voxels are both related to someone
speaking, but they seem to depend on interactions between the noun (me or you)
and the verb (says). To elicit a large response both must be present,
something which is difficult to capture in additive models, even with ngrams,
since these words may not be close together in a sentence.

This interpretation approach could be applied more rigorously to generate
hypotheses for text inputs that activate brain regions, and then testing them
with followup fMRI experiments.

Figure 5:  Aug-imodels prediction performance and interpretation for fMRI
voxels. (A) Map of the difference between the performance of Aug-GAM and BERT
for fMRI voxel prediction across the cortex. Positive values (red) show where
Aug-GAM outperforms BERT (measured by correlation on the test set). (B) Aug-
GAM outperforms BERT when averaging across all voxels (or just over the 1%/5%
with the highest test correlations). Standard errors of the mean are all less
than 0.0015. (C) Example Aug-GAM model for a single voxel (visualized with the
top Aug-GAM coefficients). (D) Example Aug-Tree model for two voxels.

##  6 Background and related work

#### GAMs

There is a large literature on additive models being used for interpretable
modeling. This includes generalized additive models (GAMs) [10], which have
achieved strong performance in various domains by modeling individual
component functions/interactions using regularized boosted decision trees [35]
and more recently using neural networks [36]. However, existing GAM methods
are limited in their ability to model the high-order feature interactions that
arise in NLP. Meanwhile, NLP has seen great success in models which build
strong word-level representations, e.g. word2vec [37, 38], GloVe [24],
FastText [39] and ELMo [40]. By replacing such models with LLM embeddings,
Aug-GAM enables easily modeling ngrams of different lengths without training a
new model. Moreover, unlike earlier methods, LLMs can incorporate information
about labels into the embeddings (e.g. by first finetuning an LLM on a
downstream prediction task).

#### Decision trees

There is a long history of greedy methods for fitting decision trees, e.g.,
CART [11], ID3 [26], and C4.5 [41]. More recent work has explored fitting
trees via global optimization rather than greedy algorithms [42, 43, 44]; this
can improve performance given a fixed tree size but incurs a high
computational cost. Other recent studies have improved trees after fitting
through regularization [45] or iterative updates [46]. Beyond trees, there are
many popular classes of rule-based models, such as rule sets [47], rule lists
[48, 49], and tree sums [15]. Aug-Tree addresses a common problem shared by
rule-based approaches: modeling the sparse, correlated features that are
common in tasks such as text classification.

Beyond fitting a single tree, tree ensembles such as Random Forest [27],
gradient-boosted trees [50], XGBoost [51], and BART [52], have all shown
strong predictive performance in diverse settings. These ensembling approaches
are compatible with Aug-Tree, as it can be used as the base estimator in any
of these approaches.

#### Interpreting/distilling neural networks

The work here is related to studies that aim to make neural networks more
interpretable. For example, models can make predictions by comparing inputs to
prototypes [53, 54], by predicting intermediate interpretable concepts [55,
56, 57], using LLMs to extract prompt-based features [58, 59], distilling a
neural network into a mostly transparent model [60] or distilling into a fully
transparent model (e.g. adaptive wavelets [13] or an additive model [61]).
Separately, many works use neural network distillation to build more efficient
(but still black-box) neural network models, e.g. [62, 63].

#### Feature and feature-interaction importances

Loosely related to this work are post-hoc methods that aim to help understand
a black-box model, i.e. by providing feature importances using methods such as
LIME [29], SHAP [64], and others [65, 66]. However, these methods lose some
information by summarizing the model and suffer from issues with summarizing
interactions [67, 68]. Slightly more related are works which aim to explain
feature interactions or transformations in neural networks [69, 70, 71], but
these works fail to explain the model as a whole and are again less reliable
than having a fully transparent model.

##  7 Discussion

Aug-imodels provide a promising direction towards future methods that reap the
benefits of both LLMs and transparent models in NLP: high accuracy along with
interpretability/efficiency. This potentially opens the door for introducing
LLM-augmented models in high-stakes domains, such as medical decision-making
and in new applications on compute-limited hardware. Aug-imodels is currently
limited to applications for which an effective LLM is available, and thus may
not work well for very esoteric NLP tasks. However, as LLMs improve, the
predictive performance of Aug-imodels should continue to improve and expand to
more diverse NLP tasks. More generally, Aug-imodels can be applied to domains
outside of NLP where effective foundation models are available (e.g. computer
vision or protein engineering).

Aug-imodels can be readily extended to new model forms beyond additive models
and trees. Other transparent models, such as rule lists, rule sets, and
prototype-based models could all potentially benefit from LLM augmentation
during training time. In all these cases, LLM augmentation could use LLM
embeddings (as is done in Aug-GAM), use LLM generations (as is done in Aug-
Tree), or use LLMs in new ways. Aug-GAM could be augmented by building on the
nonlinearity present in GAMs such as the explainable boosting machine [35], to
nonlinearly transform the embedding for each ngram with a model before summing
to obtain the final prediction. Additionally, Aug-GAM could fit long-range
interaction terms as opposed to only ngrams. Aug-Tree could leverage domain
knowledge to engineer more meaningful prompts for expanding ngrams or for
extracting relevant ngrams. Both models can be further studied to improve
their compression (potentially with LLM-guided compression techniques) or to
extend their capabilities to tasks beyond classification/regression, such as
sequence prediction or outlier detection. We hope that the introduction of
Aug-imodels can help push improved performance prediction into high-stakes
applications, improve interpretability for scientific data, and reduce
unnecessary energy/compute usage.

## References

  *   * Brown et al. [2020] Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J.D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.: Language models are few-shot learners. Advances in neural information processing systems 33, 1877–1901 (2020)
  * Bubeck et al. [2023] Bubeck, S., Chandrasekaran, V., Eldan, R., Gehrke, J., Horvitz, E., Kamar, E., Lee, P., Lee, Y.T., Li, Y., Lundberg, S., et al.: Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.12712 (2023)
  * Devlin et al. [2018] Devlin, J., Chang, M.-W., Lee, K., Toutanova, K.: Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 (2018)
  * Angermueller et al. [2016] Angermueller, C., Pärnamaa, T., Parts, L., Stegle, O.: Deep learning for computational biology. Molecular systems biology 12(7), 878 (2016)
  * Kornblith et al. [2022] Kornblith, A.E., Singh, C., Devlin, G., Addo, N., Streck, C.J., Holmes, J.F., Kuppermann, N., Grupp-Phelan, J., Fineman, J., Butte, A.J., Yu, B.: Predictability and stability testing to assess clinical decision instrument performance for children after blunt torso trauma. PLOS Digital Health (2022) https://doi.org/10.1371/journal.pdig.0000076
  * Brennan and Oliver [2013] Brennan, T., Oliver, W.L.: The emergence of machine learning techniques in criminology. Criminology & Public Policy 12(3), 551–562 (2013)
  * Dwork et al. [2012] Dwork, C., Hardt, M., Pitassi, T., Reingold, O., Zemel, R.: Fairness through awareness. In: Proceedings of the 3rd Innovations in Theoretical Computer Science Conference, pp. 214–226 (2012). ACM
  * Goodman and Flaxman [2016] Goodman, B., Flaxman, S.: European union regulations on algorithmic decision-making and a” right to explanation”. arXiv preprint arXiv:1606.08813 (2016)
  * Bommasani et al. [2023] Bommasani, R., Soylu, D., Liao, T.I., Creel, K.A., Liang, P.: Ecosystem graphs: The social footprint of foundation models. arXiv preprint arXiv:2303.15772 (2023)
  * Hastie and Tibshirani [1986] Hastie, T., Tibshirani, R.: Generalized additive models. Statistical Science 1(3), 297–318 (1986)
  * Breiman et al. [1984] Breiman, L., Friedman, J.H., Olshen, R.A., Stone, C.J.: Classification and Regression Trees. Wadsworth and Brooks, Monterey, CA (1984). https://www.routledge.com/Classification-and-Regression-Trees/Breiman-Friedman-Stone-Olshen/p/book/9780412048418
  * Rudin et al. [2021] Rudin, C., Chen, C., Chen, Z., Huang, H., Semenova, L., Zhong, C.: Interpretable machine learning: Fundamental principles and 10 grand challenges. arXiv preprint arXiv:2103.11251 (2021)
  * Ha et al. [2021] Ha, W., Singh, C., Lanusse, F., Upadhyayula, S., Yu, B.: Adaptive wavelet distillation from neural networks through interpretations. Advances in Neural Information Processing Systems 34 (2021)
  * Mignan and Broccardo [2019] Mignan, A., Broccardo, M.: One neuron versus deep learning in aftershock prediction. Nature 574(7776), 1–3 (2019)
  * Tan et al. [2022] Tan, Y.S., Singh, C., Nasseri, K., Agarwal, A., Yu, B.: Fast interpretable greedy-tree sums (figs). arXiv:2201.11931 [cs, stat] (2022). arXiv: 2201.11931
  * Wang and Komatsuzaki [2021] Wang, B., Komatsuzaki, A.: GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github.com/kingoflolz/mesh-transformer-jax (2021)
  * Honnibal and Montani [2017] Honnibal, M., Montani, I.: spaCy 2: Natural language understanding with Bloom embeddings, convolutional neural networks and incremental parsing. To appear (2017)
  * Saravia et al. [2018] Saravia, E., Liu, H.-C.T., Huang, Y.-H., Wu, J., Chen, Y.-S.: Carer: Contextualized affect representations for emotion recognition. In: Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 3687–3697 (2018)
  * Malo et al. [2014] Malo, P., Sinha, A., Korhonen, P., Wallenius, J., Takala, P.: Good debt or bad debt: Detecting semantic orientations in economic texts. Journal of the Association for Information Science and Technology 65 (2014)
  * Pang and Lee [2005] Pang, B., Lee, L.: Seeing stars: Exploiting class relationships for sentiment categorization with respect to rating scales. In: Proceedings of the ACL (2005)
  * Socher et al. [2013] Socher, R., Perelygin, A., Wu, J., Chuang, J., Manning, C.D., Ng, A., Potts, C.: Recursive deep models for semantic compositionality over a sentiment treebank. In: Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pp. 1631–1642 (2013)
  * LeBel et al. [2022] LeBel, A., Wagner, L., Jain, S., Adhikari-Desai, A., Gupta, B., Morgenthal, A., Tang, J., Xu, L., Huth, A.G.: A natural language fmri dataset for voxelwise encoding models. bioRxiv (2022)
  * Jones [1972] Jones, K.S.: A statistical interpretation of term specificity and its application in retrieval. Journal of documentation (1972)
  * Pennington et al. [2014] Pennington, J., Socher, R., Manning, C.D.: Glove: Global vectors for word representation. In: Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 1532–1543 (2014)
  * Wolf et al. [2019] Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., et al.: Huggingface’s transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771 (2019)
  * Quinlan [1986] Quinlan, J.R.: Induction of decision trees. Machine learning 1(1), 81–106 (1986)
  * Breiman [2001] Breiman, L.: Random forests. Machine Learning 45(1), 5–32 (2001) https://doi.org/10.1023/A:1010933404324
  * Bach et al. [2022] Bach, S.H., Sanh, V., Yong, Z.-X., Webson, A., Raffel, C., Nayak, N.V., Sharma, A., Kim, T., Bari, M.S., Fevry, T., et al.: Promptsource: An integrated development environment and repository for natural language prompts. arXiv preprint arXiv:2202.01279 (2022)
  * Ribeiro et al. [2016] Ribeiro, M.T., Singh, S., Guestrin, C.: Why should i trust you?: Explaining the predictions of any classifier. In: Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 1135–1144 (2016). ACM
  * Lundberg and Lee [2016] Lundberg, S., Lee, S.-I.: An unexpected unity among methods for interpreting model predictions. arXiv preprint arXiv:1611.07478 (2016)
  * Huth et al. [2016] Huth, A.G., De Heer, W.A., Griffiths, T.L., Theunissen, F.E., Gallant, J.L.: Natural speech reveals the semantic maps that tile human cerebral cortex. Nature 532(7600), 453–458 (2016)
  * Schrimpf et al. [2021] Schrimpf, M., Blank, I.A., Tuckute, G., Kauf, C., Hosseini, E.A., Kanwisher, N., Tenenbaum, J.B., Fedorenko, E.: The neural architecture of language: Integrative modeling converges on predictive processing. Proceedings of the National Academy of Sciences 118(45), 2105646118 (2021)
  * Antonello and Huth [2022] Antonello, R.J., Huth, A.: Predictive coding or just feature discovery? an alternative account of why language models fit brain data. Neurobiology of Language, 1–39 (2022)
  * Caucheteux and King [2022] Caucheteux, C., King, J.-R.: Brains and algorithms partially converge in natural language processing. Communications biology 5(1), 1–10 (2022)
  * Caruana et al. [2015] Caruana, R., Lou, Y., Gehrke, J., Koch, P., Sturm, M., Elhadad, N.: Intelligible models for healthcare: Predicting pneumonia risk and hospital 30-day readmission. In: Proceedings of the 21th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 1721–1730 (2015)
  * Agarwal et al. [2021] Agarwal, R., Melnick, L., Frosst, N., Zhang, X., Lengerich, B., Caruana, R., Hinton, G.E.: Neural additive models: Interpretable machine learning with neural nets. Advances in Neural Information Processing Systems 34, 4699–4711 (2021)
  * Mikolov et al. [2013a] Mikolov, T., Chen, K., Corrado, G., Dean, J.: Efficient estimation of word representations in vector space. arXiv preprint arXiv:1301.3781 (2013)
  * Mikolov et al. [2013b] Mikolov, T., Sutskever, I., Chen, K., Corrado, G.S., Dean, J.: Distributed representations of words and phrases and their compositionality. Advances in neural information processing systems 26 (2013)
  * Joulin et al. [2016] Joulin, A., Grave, E., Bojanowski, P., Mikolov, T.: Bag of tricks for efficient text classification. arXiv preprint arXiv:1607.01759 (2016)
  * Peters et al. [2018] Peters, M.E., Neumann, M., Iyyer, M., Gardner, M., Clark, C., Lee, K., Zettlemoyer, L.: Deep contextualized word representations. In: Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), pp. 2227–2237. Association for Computational Linguistics, New Orleans, Louisiana (2018). https://doi.org/10.18653/v1/N18-1202 . https://aclanthology.org/N18-1202
  * Quinlan [2014] Quinlan, J.R.: C4. 5: Programs for Machine Learning. Elsevier, ??? (2014)
  * Lin et al. [2020] Lin, J., Zhong, C., Hu, D., Rudin, C., Seltzer, M.: Generalized and scalable optimal sparse decision trees. In: International Conference on Machine Learning, pp. 6150–6160 (2020). PMLR
  * Hu et al. [2019] Hu, X., Rudin, C., Seltzer, M.: Optimal sparse decision trees. Advances in Neural Information Processing Systems (NeurIPS) (2019)
  * Bertsimas and Dunn [2017] Bertsimas, D., Dunn, J.: Optimal classification trees. Machine Learning 106(7), 1039–1082 (2017)
  * Agarwal et al. [2022] Agarwal, A., Tan, Y.S., Ronen, O., Singh, C., Yu, B.: Hierarchical shrinkage: improving the accuracy and interpretability of tree-based methods. arXiv:2202.00858 [cs, stat] (2022). arXiv: 2202.00858
  * Carreira-Perpinán and Tavallali [2018] Carreira-Perpinán, M.A., Tavallali, P.: Alternating optimization of decision trees, with application to learning sparse oblique trees. Advances in neural information processing systems 31 (2018)
  * Friedman and Popescu [2008] Friedman, J.H., Popescu, B.E.: Predictive learning via rule ensembles. The Annals of Applied Statistics 2(3), 916–954 (2008) https://doi.org/10.1214/07-aoas148
  * Angelino et al. [2017] Angelino, E., Larus-Stone, N., Alabi, D., Seltzer, M., Rudin, C.: Learning certifiably optimal rule lists for categorical data. arXiv preprint arXiv:1704.01701 (2017)
  * Singh et al. [2021] Singh, C., Nasseri, K., Tan, Y.S., Tang, T., Yu, B.: imodels: a python package for fitting interpretable models. Journal of Open Source Software 6(61), 3192 (2021) https://doi.org/10.21105/joss.03192
  * Freund et al. [1996] Freund, Y., Schapire, R.E., et al.: Experiments with a new boosting algorithm. In: Icml, vol. 96, pp. 148–156 (1996). Citeseer
  * Chen and Guestrin [2016] Chen, T., Guestrin, C.: Xgboost: A scalable tree boosting system. In: Proceedings of the 22nd Acm Sigkdd International Conference on Knowledge Discovery and Data Mining, pp. 785–794 (2016)
  * Chipman et al. [2010] Chipman, H.A., George, E.I., McCulloch, R.E.: Bart: Bayesian additive regression trees. The Annals of Applied Statistics 4(1), 266–298 (2010)
  * Li et al. [2018] Li, O., Liu, H., Chen, C., Rudin, C.: Deep learning for case-based reasoning through prototypes: A neural network that explains its predictions. In: Proceedings of the AAAI Conference on Artificial Intelligence, vol. 32 (2018)
  * Chen et al. [2019] Chen, C., Li, O., Tao, D., Barnett, A., Rudin, C., Su, J.K.: This looks like that: deep learning for interpretable image recognition. Advances in neural information processing systems 32 (2019)
  * Koh et al. [2020] Koh, P.W., Nguyen, T., Tang, Y.S., Mussmann, S., Pierson, E., Kim, B., Liang, P.: Concept bottleneck models. In: International Conference on Machine Learning, pp. 5338–5348 (2020). PMLR
  * Yang et al. [2022] Yang, Y., Panagopoulou, A., Zhou, S., Jin, D., Callison-Burch, C., Yatskar, M.: Language in a bottle: Language model guided concept bottlenecks for interpretable image classification. arXiv preprint arXiv:2211.11158 (2022)
  * Ghosh et al. [2023] Ghosh, S., Yu, K., Arabshahi, F., Batmanghelich, K.: Route, interpret, repeat: Blurring the line between post hoc explainability and interpretable models. arXiv preprint arXiv:2302.10289 (2023)
  * Yuksekgonul et al. [2022] Yuksekgonul, M., Wang, M., Zou, J.: Post-hoc concept bottleneck models. arXiv preprint arXiv:2205.15480 (2022)
  * McInerney et al. [2023] McInerney, D.J., Young, G., Meent, J.-W., Wallace, B.C.: Chill: Zero-shot custom interpretable feature extraction from clinical notes with large language models. arXiv preprint arXiv:2302.12343 (2023)
  * Frosst and Hinton [2017] Frosst, N., Hinton, G.: Distilling a neural network into a soft decision tree. arXiv preprint arXiv:1711.09784 (2017)
  * Tan et al. [2018] Tan, S., Caruana, R., Hooker, G., Koch, P., Gordo, A.: Learning global additive explanations for neural nets using model distillation. arXiv preprint arXiv:1801.08640 (2018)
  * Hinton et al. [2015] Hinton, G., Vinyals, O., Dean, J.: Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531 (2015)
  * Sanh et al. [2019] Sanh, V., Debut, L., Chaumond, J., Wolf, T.: Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108 (2019)
  * Lundberg et al. [2019] Lundberg, S.M., Erion, G., Chen, H., DeGrave, A., Prutkin, J.M., Nair, B., Katz, R., Himmelfarb, J., Bansal, N., Lee, S.-I.: Explainable ai for trees: From local explanations to global understanding. arXiv preprint arXiv:1905.04610 (2019)
  * Friedman [2001] Friedman, J.H.: Greedy function approximation: a gradient boosting machine. Annals of statistics, 1189–1232 (2001)
  * Devlin et al. [2019] Devlin, S., Singh, C., Murdoch, W.J., Yu, B.: Disentangled attribution curves for interpreting random forests and boosted trees. arXiv preprint arXiv:1905.07631 (2019)
  * Rudin [2018] Rudin, C.: Please stop explaining black box models for high stakes decisions. arXiv preprint arXiv:1811.10154 (2018)
  * Murdoch et al. [2019] Murdoch, W.J., Singh, C., Kumbier, K., Abbasi-Asl, R., Yu, B.: Definitions, methods, and applications in interpretable machine learning. Proceedings of the National Academy of Sciences of the United States of America 116(44), 22071–22080 (2019) https://doi.org/10.1073/pnas.1900654116
  * Janizek et al. [2021] Janizek, J.D., Sturmfels, P., Lee, S.-I.: Explaining explanations: Axiomatic feature interactions for deep networks. J. Mach. Learn. Res. 22, 104–1 (2021)
  * Singh et al. [2019] Singh, C., Murdoch, W.J., Yu, B.: Hierarchical interpretations for neural network predictions. International Conference on Learning Representations, 26 (2019)
  * Singh et al. [2020] Singh, C., Ha, W., Lanusse, F., Boehm, V., Liu, J., Yu, B.: Transformation Importance with Applications to Cosmology (2020)
  * Hazourli [2022] Hazourli, A.: Financialbert - a pretrained language model for financial text mining (2022) https://doi.org/10.13140/RG.2.2.34032.12803
  * Morris et al. [2020] Morris, J.X., Lifland, E., Yoo, J.Y., Grigsby, J., Jin, D., Qi, Y.: Textattack: A framework for adversarial attacks, data augmentation, and adversarial training in nlp. arXiv preprint arXiv:2005.05909 (2020)
  * Akl et al. [2021] Akl, H.A., Mariko, D., De Mazancourt, H.: Yseop at finsim-3 shared task 2021: Specializing financial domain learning with phrase representations. arXiv preprint arXiv:2108.09485 (2021)
  * Liu et al. [2019] Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., Stoyanov, V.: Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 (2019)
  * Pedregosa et al. [2011] Pedregosa, F., Varoquaux, G.ë.l., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., et al.: Scikit-learn: Machine learning in python. the Journal of machine Learning research 12(Oct), 2825–2830 (2011)
  * Su et al. [2022] Su, H., Kasai, J., Wang, Y., Hu, Y., Ostendorf, M., Yih, W.-t., Smith, N.A., Zettlemoyer, L., Yu, T., et al.: One embedder, any task: Instruction-finetuned text embeddings. arXiv preprint arXiv:2212.09741 (2022)
  * Raffel et al. [2020] Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., Liu, P.J., et al.: Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res. 21(140), 1–67 (2020)

##  Appendix A Aug-GAM

Table 4: Table of pre-trained models with unique huggingface identifiers. All models are used through huggingface [25], and linear/tree baselines are fit with scikit-learn [76] and imodels [49]. | BERT
---|---
Base (no finetuning) | bert-base-uncased [3]
Emotion | nateraw/bert-base-uncased-emotion
Financial phrasebank | ahmedrachid/FinancialBERT-Sentiment-Analysis [72]
Rotten tomatoes | textattack/bert-base-uncased-rotten_tomatoes [73]
SST2 | textattack/bert-base-uncased-SST-2 [73]
| DistilBERT
Base (no finetuning) | distilbert-base-uncased [63]
Emotion | aatmasidha/distilbert-base-uncased-finetuned-emotion
Financial phrasebank | yseop/distilbert-base-financial-relation-extraction [74]
Rotten tomatoes | textattack/distilbert-base-uncased-rotten-tomatoes [73]
SST2 | distilbert-base-uncased-finetuned-sst-2-english
|  RoBERTa [75]
Emotion | bhadresh-savani/roberta-base-emotion
Financial phrasebank | abhilash1910/financial_roberta
Rotten tomatoes | textattack/roberta-base-rotten-tomatoes [73]
SST2 | textattack/roberta-base-SST-2 [73]
Figure 6: Model performance decreases with increasing model uncertainty.
Cumulative validation accuracy decreases as more uncertain samples (based on
the model’s predicted probability) are added.

#### Varying Aug-GAM settings

By default (Table 5), we use the final embedding layer of the model (and
average it over the sequence length to get a fixed size vector), but Table 5
also shows results using the pooler output layer of the BERT model. The choice
of layer (i.e. final embedding layer versus pooler output) does not seem to
make a large difference in the final performance results. Table 5 also shows
one variation of the model (BERT finetuned (noun chunks)) where rather than
training on all ngrams, the model is fit to only noun-phrases extracted by
spaCy’s dependency parser [17]. This results in a performance drop across the
datasets, suggesting that these noun-phrases alone are insufficient to perform
the classification task. We also run an experiment where we extract embeddings
using Instructor ([77], hkunlp/instructor-xl), which allows giving a
contextual prompt for each dataset.

#### Evaluating zero-shot accuracy with language models

To measure generalization ability, we evaluate explanations based on accuracy
as a prompt for other models. Accuracy is computed following [1, 78]: using
exact matching with beam search, a beam width of 4, and a length penalty of
$\alpha=0.6$. For sentiment evaluation, we use each prompt with the template
Input: “${input}”{prompt}.888In initial experiments, we find that performance
drops significantly when learning a prompt that comes before the input. We use
positive and negative as positive and negative labels and require the LLM to
rank the two options. Human-written prompts are adapted to this template from
open-source prompts available through PromptSource [28].

Table 5: Generalization accuracy varies depending on the model used to extract embeddings. Finetuning the embedding model improves Aug-GAM performance, using a BERT model seems to outperform a DistilBERT model, and the layer used to extract embeddings does not have too large an effect. Top two methods are bolded in each column.  | Financial phrasebank | Rotten tomatoes | SST2 | Emotion
---|---|---|---|---
BERT finetuned |  92.8 $\pm$0.37 |  81.6 $\pm$0.05 |  86.9 $\pm$0.1 |  89.5 $\pm$0.03
BERT finetuned (pooler output) |  93.5 $\pm$0.05 |  81.3 $\pm$0.13 |  87.8 $\pm$0.21 |  89.8 $\pm$0.07
BERT finetuned (noun chunks) |  87.9 $\pm$0.08 |  79.7 $\pm$0.45 |  84.1 $\pm$0.14 |  87.1 $\pm$0.2
BERT |  84.1 $\pm$0.08 |  78.1 $\pm$0.16 |  82.8 $\pm$0.27 |  67.1 $\pm$0.06
BERT (pooler output) |  82.7 $\pm$0.28 |  78.5 $\pm$0.03 |  80.7 $\pm$0.11 |  58.0 $\pm$0.29
DistilBERT finetuned |  85.8 $\pm$0.34 |  78.5 $\pm$0.34 |  81.7 $\pm$0.07 |  68.8 $\pm$0.11
DistilBERT |  81.7 $\pm$0.34 |  79.8 $\pm$0.08 |  86.8 $\pm$0.1 |  87.5 $\pm$0.11
RoBERTa finetuned |  77.8 $\pm$0.31 |  83.6 $\pm$0.03 |  89.1 $\pm$0.24 |  88.5 $\pm$0.19
Instructor prompted | 76.1 | 80.0 | 84.7 | 71.0

###  A.1 Test-time tradeoffs between accuracy and interpretability/speed

The ability to effectively generalize to unseen tokens in Fig 4C/D raises the
question of whether one can vary the order of ngrams used at test-time, to get
a tradeoff between accuracy and interpretability (i.e. how many features are
used to make a prediction). Depending on the relative importance of accuracy
and interpretability for a given problem, one may select to use a different
number of features for testing. Fig 7 suggests that this is feasible.

Fig 7A shows the prediction performance when compressing the Aug-GAM model
(fit using 4-grams and finetuned BERT) by setting the coefficients with the
smallest magnitudee to zero. Some models require only a few coefficients to
perform well and some models (e.g. the Emotion and Financial phrasebank
models) predict more accurately when using less than 50% of the original
coefficients. Fig 7B it shows the accuracy of the same models in Fig 7A, as
the order of ngrams used only for testing is varied. As the number of features
used for testing increases, the performance tends to increase but
interpretations become more difficult.

Fig 8 characterizes the full tradeoff between the number of ngrams used for
fitting versus testing for all datasets. Generally, the best performance is
achieved when the same number of ngrams is used for training and testing (the
diagonal). Performance tends to degrade significantly when fewer ngrams are
used for testing than training (lower-left).

\begin{overpic}[width=195.12767pt,tics=10]{figs/vary_sparsity.pdf}
\put(1.0,55.0){\large{A}}
\end{overpic}\begin{overpic}[width=195.12767pt,tics=10]{figs/vary_ngrams_test/vary_ngrams_single.pdf}
\put(1.0,55.0){\large{B}} \end{overpic}

Figure 7: Aug-GAM performance when varying the ngrams used for testing. (A)
Performance when removing the smallest coefficients from an Aug-GAM model. (B)
Performance when varying the order of ngrams used for testing.

Figure 8: Varying the order of ngrams used for training and testing across
each of the five datasets in Table 1. Some models (i.e. rows) perform
reasonably well as the order of ngrams used for testing is varied, potentially
enabling a test-time tradeoff between accuracy and interpretability.
Generally, using higher-order ngrams during testing improves performance and
testing with less ngrams than used for training hurts performance
considerably.

###  A.2 Comparison with post-hoc feature importance

The coefficients learned by Aug-GAM often differ from importances assigned by
post-hoc feature-importance methods. Aug-GAM learns a single coefficient for
each ngram across the dataset, allowing for auditing/editing the model with
visualizations such as Fig 4. In contrast, popular methods for post-hoc
feature importance, such as LIME [29] and SHAP [30] yield importance scores
that vary based on the context in each input. This can be useful for debugging
complex nonlinear models, but these scores (i) are approximations, (ii) must
summarize nonlinear feature interactions, and (iii) vary across predictions,
making transparent models preferable whenever possible.

Fig 9 shows an example of the Aug-GAM coefficients for the SST2 model from Fig
4 for different ngrams when making a prediction for the phrase not very good.
While Aug-GAM yields scores for each subphrase that match human judgement (as
seen in Fig 4B/D), posthoc feature importance methods summarize the
interactions between different ngrams into individual words, potentially
making interpretation difficult. Scores are rescaled to be between -1 and 1 to
make them comparable. See Aug-GAM scores for many top-interacting phrases in
Fig 10.

Figure 9: Comparing Aug-GAM ngram coefficients (left) to word-level feature
importances from posthoc methods (right): LIME and SHAP.

#### Summing embeddings meaningfully captures interactions

One potential concern with the Aug-GAM model is that it may fail to learn
interactions since it simply sums the embeddings of individual ngrams, and the
language model extractor may not sufficiently capture interactions in its
embedding space. To investigate this concern, we first identify bigrams that
involve interaction by fitting a unigram bag-of-words model and a bigram bag-
of-ngrams model to SST2. We then use these two models to select the 10 bigrams
for which the bigram coefficient is farthest from the sum of the coefficients
for each unigram.

Fig 10 shows the resulting bigrams containing interactions. For each bigram,
it shows the Aug-GAM learned coefficient (i.e. the contribution to the
prediction $w^{T}\phi(x_{i})$) for the bigram (gray bar) along with each of
its constituent unigrams (blue and orange bars). It is clear that the bigram
coefficient is not the simple naive sum of the unigram coefficients (dashed
black bar), and the learned coefficients make intuitive sense, suggesting that
this Aug-GAM model has successfully learned interactions.

Figure 10: Aug-GAM accurately learns interactions rather than simply summing
the contributions of individual unigrams.

##  Appendix B Aug-Tree

Algorithm B2 Aug-Tree algorithm for fitting a single split.

1:Split-Aug-Tree(X, y, LLM):

2: # Add original CART keyphrase

3:keyphrase = split_CART(X, y)

4:keyphrases_expanded = LLM(“Generate similar keyphrases to ” \+ keyphrase)

5:keyphrases_running = [keyphrase]

6:impurity_decrease_best = calc_impurity_decrease(X, y, keyphrases_running)

7: # Try adding new keyphrases

8:for k in keyphrases_expanded:

9: keyphrases_running.push(k)

10: impurity_decrease_new = calc_impurity_decrease(X, y, keyphrases_running)

11: if impurity_decrease_new $<$ impurity_decrease_best then

12: keyphrases_running.pop()

13: return keyphrases_running

Table 6: Metadata on keyphrase expansions. Results are averaged over keyphrases found in the 4 text-classification datasets in Table 1 when fitting a 40-tree bagging ensemble. The LLM is queried for 100 expansion candidates, but due to imperfect LLM generations, only 91.6 candidates are generated on average. After deduplication (converting to lowercase, removing whitespaces, etc.), only 83.3 candidates remain. Screening removes almost all candidates, leaving only 0.8 candidates on average.  # Expansion candidates (Before deduplication) | # Expansion candidates |  # Expansions (After screening)
---|---|---
91.6$\pm$0.7 |  83.3$\pm$0.8 |  0.8$\pm$0.1

###  B.1 Aug-Tree variations

Table 7 explores different variations of Aug-Tree. The top row shows learning
a single tree with Aug-Tree using its default parameters, achieving the best
performance across the datasets. Table 7 shows results for different
algorithmic choices, such as replacing the generic prompt with a dataset-
specific one (Aug-Tree (Contextual prompt)), and searching for new keyphrases
using 5 CART features instead of one (Aug-Tree (5 CART features)). We also
consider preprocessing the data differently, using Stemming (with the Porter
Stemmer) or using Trigrams, rather than bigrams.

One major variation we study is using LLM embeddings to find keyphrases,
rather than querying via a prompt (Aug-Tree (Embeddings). Specifically, we
consider expanding keywords by finding the keyphrases that are closest in
embedding space (measured by euclidean distance) to the original keyphrase.
This option may be desirable computationally, as it may require a smaller LLM
to compute effective embeddings (e.g. BERT [3]) compared to a larger LLM
required to directly generate relevant keyphrases (e.g. GPT3 [1]). However,
finding closest embeddings requires making more calls to the LLM, as
embeddings must be calculated and compared across all ngrams in
$X_{\text{text}}$.

Table 7: Performance (ROC AUC) for variations of Aug-Tree. Values are averaged over 3 random dataset splits; error bars are standard error of the mean (many are within the points). | Emotion | Financial phrasebank | Rotten tomatoes | SST2
---|---|---|---|---
Aug-Tree | 0.680 $\pm$0.029 | 0.825 $\pm$0.006 | 0.622 $\pm$0.007 | 0.673 $\pm$0.008
Aug-Tree (Embeddings) |  0.599 $\pm$0.008 |  0.776 $\pm$0.018 |  0.600 $\pm$0.011 |  0.663 $\pm$0.002
Aug-Tree (Contextual prompt) |  0.667 $\pm$0.011 |  0.820 $\pm$0.004 |  0.627 $\pm$0.008 |  0.669 $\pm$0.005
Aug-Tree (5 CART features) |  0.711 $\pm$0.039 |  0.730 $\pm$0.026 |  0.608 $\pm$0.009 |  0.674 $\pm$0.003
Aug-Tree (Stemming) |  0.640 $\pm$0.019 |  0.520 $\pm$0.016 |  0.625 $\pm$0.004 |  0.679 $\pm$0.005
Aug-Tree (Trigrams) |  0.676 $\pm$0.030 |  0.826 $\pm$0.006 |  0.619 $\pm$0.010 |  0.669 $\pm$0.006
CART |  0.574 $\pm$0.002 |  0.775 $\pm$0.005 |  0.599 $\pm$0.005 |  0.636 $\pm$0.002
ID3 |  0.573 $\pm$0.004 |  0.795 $\pm$0.010 |  0.589 $\pm$0.002 |  0.638 $\pm$0.009
Table 8: Performance (Accuracy) for Aug-Treeand Aug-Tree Ensemble. Values are averaged over 3 random dataset splits; error bars are standard error of the mean (many are within the points). *Emotion and Financial phrasebank results are not directly comparable to Table 2, as they have been modified for binary classification.  | Emotion* | Financial phrasebank* | Rotten tomatoes | SST2
---|---|---|---|---
Aug-Tree |  0.637 $\pm$0.045 |  0.818 $\pm$0.014 |  0.613 $\pm$0.009 |  0.571 $\pm$0.018
Aug-Tree Ensemble |  0.800 $\pm$0.008 |  0.848 $\pm$0.006 |  0.619 $\pm$0.004 |  0.614 $\pm$0.016

##  Appendix C fMRI experiment details

This section gives more details on the fMRI experiment analyzed in Sec 5; for
more scientific details see the original study [22]. Sec 5 analyzes data from
one human subject (UTS03) in the original study, as the subject listened to
approximately hours of narrative speech from the Moth Radio Hour, which
consists of short autobiographical stories. The subject underwent fMRI
scanning as they listened, yielding an fMRI volume brain scan consisting of
95,556 voxels roughly every two seconds.

The individual voxel models described in Sec 5 are each fit to 9,461 training
points, each corresponding to a different time point (after accounting for
various preprocessing steps, such as trimming the beginning and end of the
sequence). They are evaluated on 291 volumes which come from a narrative story
that was not seen during training.

Fig 11 shows the generalization performance of the model for each voxel,
measured by the correlation between the predicted response and the measured
response. Fig 5 shows the performance difference between the Aug-GAM model and
the BERT baseline.

Figure 11: Generalization performance for individual-voxel models, measured by correlation between the predicted response and the measured response on the held-out test set. Some regions are very poorly predicted (blue), but many voxels can be predicted quite well (red). Table 9: fMRI prediction performance using different methods. Eng1000 is a linear word-embedding baseline similar to word2vec which has been used in the neuroscience literature [31]. Model | Order of ngram | $\rho$ |  $\rho$ (top 1%) |  $\rho$ (top 5%)
---|---|---|---|---
Eng1000 | 1 | 0.041 | 0.529 | 0.439
GloVe | 1 | 0.044 | 0.521 | 0.426
BERT | 5 | 0.022 | 0.386 | 0.302
BERT | 10 | 0.035 | 0.457 | 0.365
BERT | 20 | 0.053 | 0.524 | 0.429
Aug-GAM (BERT) | 5 | 0.061 | 0.583 | 0.489
Aug-GAM (BERT) | 10 | 0.062 | 0.586 | 0.489
Aug-GAM (BERT) | 20 | 0.056 | 0.566 | 0.470