Proto-lm: A Prototypical Network-Based Framework for Built-in Interpretability in Large Language Models

Let $x_{i}$ represent a single input text and $y_{i}$ its associated label. We denote $f(x_{i})$ as the encoding of $x_{i}$ by an underlying LLM, $f$. We pass $f(x_{i})$ through a token-level attention layer that learns to emphasize different parts of the text, allowing us to identify important words within each sample. We first calculate $\upsilon_{t}$, the output of a fully-connected layer with bias ($\psi$ ), applied to $h_{t}$, which the embedding of each token $t$ in $f(x_{i})$, followed by a $\tanh$ activation function ($\sigma$). We then compute $\nu_{t}$, the dot product of $\upsilon_{t}$, and a token-level weight vector $W_{\nu}$.

$$\upsilon_{t}=\sigma(\psi(h_{t}))\;\;\;\;\;\;\;\;\nu_{t}=\upsilon_{t}\cdot W_{\nu}$$

(1)

We calculate the attention score $\alpha_{t}$ for each token $t$ using the softmax function. The attended encoding of $f(x_{i})$, $S_{i}$, is then obtained by taking the weighted sum of the token embeddings, where the weight for each token is determined by its attention score.

$$\alpha_{t}=\frac{\exp(\nu_{t})}{{\sum_{t\in f(x_{i})}}\exp(\nu_{t})}\;\;\;\;\;S_{i}=\sum_{t\in f(x_{i})}\alpha_{t}h_{t}$$

(2)

In the prototypical layer $P$, we create $N$ prototype vectors and assign $n$ prototypes to each class $c\in C$, where $C$ represents the set of all classes in the dataset. To ensure an equal number of prototype vectors are allocated to capture representative features for each class, we set $n=\frac{N}{|C|}$, where $|C|$ denotes the total number of classes in the dataset. The input $S_{i}$ is then fed into $P$, where we calculate the similarity between $S_{i}$ and every prototype vector $p\in P$ by inverting the $L_{2}$-distance. We then concatenate all $N$ similarities to obtain the vector $M_{i}$, which serves as the embedding of input $x_{i}$ in the prototypical space. Each dimension of $M_{i}$ can be interpreted as the similarity between $M_{i}$ and a prototype. Subsequently, $M_{i}$ is fed through a fully-connected layer $W_{h}$ of dimension $N\times|C|$ to obtain logits for each class. Let $W_{c}$ denote the set of weights in $W_{h}$ that connect similarities in $M_{i}$ to the logit of class $c$. Our model’s output probabilities are computed as follows:

$$Pr(\hat{y}=c\;|\;x_{i})=\frac{\exp(M_{i}w_{c})}{{\sum_{c\in C}}\exp(M_{i}w_{c})}$$

(3)

To create more representative prototypes and shape the clusters of prototypes in the prototypical embedding space, we introduce a cohesion loss term $\mathscr{L}_{coh}$ and a separation loss term $\mathscr{L}_{sep}$ into our loss function in addition to the cross entropy loss $\mathcal{L}_{ce}$. Let $K$ be an integer hyperparameter, $p_{j}$ denote a singular prototype, and $P_{y_{i}}$ represent all prototypes in $P$ that belong to class $y_{i}$, our loss terms are defined as follows:

$$\mathscr{L}_{coh}=\frac{1}{K}\cdot\;{\sum}_{\forall j:p_{j}\in P_{y_{i}}}{\max_{K}}\lVert{S_{i}-p_{j}}\rVert^{2}_{2}$$

(4)

$$\mathscr{L}_{sep}=-\frac{1}{K}\cdot\;{\sum}_{\forall j:p_{j}\not\in P_{y_{i}}}{\min_{K}}\lVert{S_{i}-p_{j}}\rVert^{2}_{2}$$

(5)

For every $S_{i}$ in the input batch, $\mathscr{L}_{coh}$ penalizes the average distance between $S_{i}$ and the $K$ most distant prototypes of its class ($y_{i}$) while $\mathscr{L}_{sep}$ penalizes the average distance between $S_{i}$ and the $K$ most similar prototypes that do not belong to $y_{i}$. Intuitively, for every $S_{i}$, $\mathscr{L}_{coh}$ “pulls” $K$ prototypes of the correct class close while $\mathscr{L}_{sep}$ “pushes” $K$ prototypes of the incorrect class away. We then add cross-entropy loss $\mathscr{L}_{ce}$ and weight each loss term with a hyperparameter $\lambda$ such that $\lambda_{0}+\lambda_{1}+\lambda_{2}=1$ to obtain the following loss function:

$${\mathscr{L}_{total}}=\lambda_{0}\cdot\mathscr{L}_{ce}+\lambda_{1}\cdot\mathscr{L}_{coh}+\lambda_{2}\cdot\mathscr{L}_{sep}$$

(6)

To understand each prototype vector $p_{j}$ in natural language, we project each prototype onto the nearest token-level attended encoding ($S_{j}$) of a sample text ($x_{j})$ in the training data that belongs to the same class as $p_{j}$ and assign the token-attended text of that prototype. Let $D$ be the training dataset. We formally denote the projection as in eq.7: $\forall(x_{j},y_{j})\in D\text{ such that }y_{j}=c$:

$$\text{Text of}\;p_{j}\leftarrow\underset{S_{j}}{{{\operatorname*{argmin}}}}\;{\lVert{S_{j}-p_{j}}\rVert^{2}_{2}}$$

(7)

In other words, for each prototype $p_{j}$, we find the nearest $S_{j}$ in the training dataset that belongs to the same class as $p_{j}$ and assign the corresponding token-attended text to $p_{j}$ (Details in App. D & E).

Compared to black-box models such as BERT/RoBERTa/BART, which have uninterpretable embedding dimensions, proto-lm offers the added benefit of interpretability in conjunction with competitive performance. We show an example of a 2-dimensional prototypical space in Fig. 3, where proto-lm provides insight into the model’s decision-making process by allowing us to visualize examples in prototypical space, where each dimension represents the normalized similarity $\in[0,1]$ between an example and one prototype.

In Fig. 3, we select one positive and one negative prototype from the prototypical layer of the model to create a 2-dimensional space in which we place examples #1-#4. The vertical axis represents similarity to the negative prototype, and the horizontal axis represents similarity to the positive prototype. From this, we can see that example #1 is much more similar to the negative prototype than to the positive prototype, and example #3 is much more similar to the positive prototype than to the negative prototype, and both are correctly classified as a result.

From the prototypical space, we can see clearly that the model’s decision to misclassify example #2 as positive is due to elements in the example that make it similar to the positive prototype. Similarly, the prototypical space of proto-lm reveals that example #4 was misclassified because, while it contains elements that are similar to both the positive and negative prototypes, it is closer to the negative prototype. The interpretable prototypical space of proto-lm provides an explanation for why difficult cases such as examples #2 and #4 are misclassified. It should be noted that while similar analyses and explanations can be obtained through post-hoc techniques such as generating sentence embeddings Reimers and Gurevych (2019); Gao et al. (2021), proto-lm has this capability built-in.

We investigate the effect of different weightings of the loss terms in our loss function. We train proto-lm, holding all other hyperparameters constant, under 11 different values of $\lambda_{0}$ evenly distributed on $[0,1]$, and report our results in Figure 4. For these experiments, we place equal weight on $\mathscr{L}_{coh}$ and $\mathscr{L}_{sep}$ such that $\lambda_{1}=\lambda_{2}=\frac{1-\lambda_{0}}{2}$ (additional details in App. C).

Because prototypes not only help interpret the model but also capture the most representative features from the data, placing an excessive emphasis on $\mathcal{L}_{ce}$ actually achieves the adverse effect. As shown in Figure 4, increasing $\lambda_{0}$ comes at the cost of learning representative prototypes (which $\mathscr{L}_{coh}$ and $\mathscr{L}_{sep}$ do) and this is reflected in decreasing classification accuracy on the downstream task. We observe that the optimal accuracy is achieved when $\lambda_{0}=0.3$. As we increase $\lambda_{0}$ from 0.3, we see not only a decline in accuracy but also a decline in the quality of the prototypes associated with the input. On the other, placing sole emphasis on $\mathscr{L}_{coh}$ and $\mathscr{L}_{sep}$ without any emphasis on $\mathscr{L}_{ce}$ (as in the case of $\lambda_{0}=0$) leads to prototypes that do not help with the downstream task.

As noted in §4.2, prototypes serve to capture features from data that are useful for making predictions in addition to interpretability. As a result, the number of prototypes ($N$) in the prototypical layer directly affects the expressiveness of our model. We conduct experiments varying the number of prototypes $N$ using RoBERTa-Large as the base and show our results in Fig. 5. We observe a general increase in accuracy up until $N=1000$, after which accuracy plateaus for some tasks while increasing only slightly for others. We reason that increasing $N$ can only improve the model’s expressiveness until the point when $N$ reaches the embedding dimension of the underlying LLM, after which more prototypes in the prototypical space no longer aid in capturing more useful features from the training data. In Fig. 5’s case, as RoBERTa-large’s output dimension is 1024, we see the increasing trend of accuracy stop at around 1000 prototypes.

The interpretability of proto-lm stems from the retrieval of the most informative training examples. If all prototypes are predominantly associated with a limited number of training examples, this reduces their utility from an interpretability perspective. Hence, it is beneficial to encourage prototype segregation, that is, a broader projection onto the dataset and a more diverse representation of different samples. Besides normalizing prototype distances, which has been shown to influence prototype segregation (Das et al., 2022), proto-lm introduces an additional hyperparameter, $K$. This parameter controls the number of prototypes that each training example associates and disassociates with during training. As changing $K$ also influences the decision-making process of the model by altering the number of samples the models compare for each input, we examine the impact of $K$ on both prototype segregation and model performance. In our experiment, we employ proto-lm with RoBERTa-large as our base on SST2 and IMDb. We set $N$, the total number of prototypes, to be 1000, and $n=1000/2$, the number of prototypes per class, to be 500. We train models under seven different $K$ values, keeping all other hyperparameters constant. We report the accuracy of the models and the percentage of unique prototypes in $P$ under each different $K$ in Fig.6.

A prototype is considered unique if no other prototype in $P$ projects onto the same sample in the dataset. We observe that the best prototype segregation (highest number of unique prototypes) occurs when $K=1$, and the number of unique prototypes significantly drops as we increase $K$. Intuitively, if more prototypes are drawn close to each sample during training (eq.4), it becomes challenging to create unique prototypes. It is important to note that eq. 5 is less related to the uniqueness of prototypes as prototypes are projected onto samples from the same class. We also witness a slight rise in model accuracy as we increase $K$. We conjecture that while unique prototypes contribute to interpretability, the model doesn’t necessarily require prototypes to be unique to make accurate decisions. Thus, we observe a minor trade-off between interpretability and model performance.

Recently, the concept of faithfulness, which measures the extent to which an explanation accurately reflects the true decision-making process of a model, has garnered attention as a criterion for evaluating explainability methods (Jacovi and Goldberg, 2020, 2021; Chan et al., 2022; Dasgupta et al., 2022). For textual inputs, faithfulness is concretely evaluated using the metrics of comprehensiveness (Comp) and sufficiency (Suff), as defined by DeYoung et al. (2019). Specifically, Comp and Suff quantify the reduction in the model’s confidence for its output when salient features are removed and retained, respectively.

We extend the application of Comp and Suff to prototypical networks to assess the faithfulness of proto-lm. For an input $x_{i}$, we initially identify the top $k\%$ of most similar prototypes and the bottom $k\%$ of least similar prototypes. Let $p^{k}_{i}$ denote the $k\%$ prototypes identified, we compute Comp as the percentage difference in model confidence when $p^{k}_{i}$ are removed (by setting $p^{k}_{i}$’s weights in $W_{h}$ are set to 0). Specifically, let $\hat{y}_{i}$ be the prediction of a model $\mathcal{M}$ on $x_{i}$, let $pr_{\hat{y}_{i}}$ be the output logit of $M$ for $\hat{y}_{i}$, and let $pr_{\hat{y}_{i}}(P\setminus p^{k}_{i})$ be the output logit when $p^{k}_{i}$ prototypes are removed from the prototypical layer, we calculate Comp as follows:

$$\text{Comp}=\frac{pr_{\hat{y}_{i}}-pr_{\hat{y}_{i}}(P\setminus p^{k}_{i})}{pr_{\hat{y}_{i}}}$$

(8)

We analogously calculate Suff as the percentage difference in model confidence when the $k\%$ of prototypes are retained:

$$\text{Suff}=\frac{pr_{\hat{y}_{i}}-pr_{\hat{y}_{i}}(p^{k}_{i})}{pr_{\hat{y}_{i}}}$$

(9)

We compute Comp and Suff $k\in{1,5,10,20,50}$ and our present mean results across SST2, QNLI, MNLI and SST5 in Fig 7. Our formulation of prototype Comp and Suff are inspired by DeYoung et al. (2019). We note here that under this formulation, a lower sufficiency is better i.e., a smaller amount of reduction in model confidence when only salient features are retained is better. As we remove more top $k\%$ prototypes, we observe a general increase in Comp. We also note a general decrease in Suff (a lower Suff is preferable) as we retain more top $k\%$ prototypes. Moreover, when the bottom $k\%$ of prototypes are removed, there are relatively small changes in Comp and large changes in Suff when only the bottom $k\%$ of prototypes are retained. These trends underscore the influence of the learned prototypes on the model’s decision-making process and their faithfulness.

Simulatability refers to the capacity of a human to mimic or replicate the decision-making process of a machine learning model Doshi-Velez and Kim (2017); Pruthi et al. (2022). It is a desirable attribute as it inherently aligns with the goal of transparently communicating the model’s underlying behavior to human users Fernandes et al. (2022). We assess the simulatability of proto-lm by providing human annotators with various explanations of model outputs on SST2/QNLI and recording the percentage of model outputs that the annotators can replicate. We provide explanations under the following four settings:

• •

  No Explanations (NE): Annotators are presented with only the sample data, without any explanations, and asked to make a decision.

• •

  Random Explanations (Random): Each sample in the dataset is presented along with prototypes from proto-lm chosen randomly as explanations. This setting serves as our baseline.

• •

  Prototypical Explanations (PE): Each sample in the dataset is presented along with the top 5 prototypes most similar to the sample when proto-lm made its decision.

• •

  PE + Output: In addition to the prototypical explanations, the model decision for each sample is also presented.

We employ workers from Amazon Mechanical Turk to crowdsource our evaluations and presented 3 workers with 50 examples each from the SST2 and QNLI datasets, along with explanations for the model’s decisions in the settings mentioned in §5.2. We use the best performing models for SST2 and QNLI (those presented in Table 1), since previous studies found that the utility of PE’s are reliant on model accuracy (Das et al., 2022). Additionally, we assess the accuracy of the human annotators in relation to the ground truth labels. We present results for both in Fig. 8. The results show that the case-based reasoning explanations offered by PEs are more effective in assisting annotators in simulating model predictions than the random baseline. We also notice a minor drop in accuracy when we provide PE + output compared to just PE. We attribute this to the fact that the models are not entirely accurate themselves, and in instances where the model is inaccurate, presenting the output leads annotators to replicate inaccurate decisions.

Furthermore, we compare proto-lm’s simulatability against 3 other well-known interpretability methods: LIME, Integrated Gradients, and Guided Backprop. We employed three workers for each example and reported the mean percentage of examples where the workers correctly replicated the model’s decision. For an example of the a simulatability questionnaire with PE, see Fig.12. For an example of an accuracy questionnaire with Random explanations, see Fig.13. For an example of the questionnaire for LIME/IG/GB explanations, see Fig.14. For results of proto-lm against LIME, IG and GB, please see Table 2. Our addditional results further indicate that proto-lm allow the workers to replicate the model’s decisions better than all the other interpretability methods on both tasks.