GPT4MIA: Utilizing Generative Pre-trained Transformer (GPT-3) as A Plug-and-Play Transductive Model for Medical Image Analysis

A fundamental component of GPT-3 is the scaled dot-product attention. Typically, three pieces of input are fed to an attention layer: queries $Q$, keys $K$, and values $V$. The scaled dot-product attention can be described as:

where $s$ is a scaling factor. Below, we show that a special case of Eq. (1) can be viewed as a nearest neighbor (NN) classifier under a cosine distance metric system¹¹1Nearest neighbor classifiers are a typical transductive method for prediction problems..

Setup 1: Suppose the key component $K$ contains features of a set of $m$ known samples, and each feature is of a unit length. The value component $V$ contains these $m$ samples’ corresponding labels, and each label is a one-hot vector. The query component $Q$ contains a feature vector (of a unit length), which represents a new test sample whose label is yet to be determined.

Proposition 1: When the scaling factor $s$ is approaching 0 (e.g., $s$ is a very small positive number), the attention function in Eq. (1) is approaching an NN classifier in the cosine distance metric system.

The above is not difficult to show. $QK^{T}$ computes the pair-wise similarities between the test sample’s feature and the features in the keys $K$. A small $s$ would enlarge the numerical gap between similar pairs and dissimilar pairs. This then leads to a one-hot-like result after applying the $softmax$ operation. The one-hot-like result is then multiplied with the values $V$, which chooses the label of a known sample that is the most similar to the test sample.

Generative Pre-trained Transformer uses a special type of attention called “self-attention”, where the $K$, $V$, and $Q$ components are all the same. We will show that in a slightly different setup from Setup 1, the self-attention mechanism can also serve a role as an NN classifier for inferring a new sample’s label given known samples’ information.

Setup 2: For each known sample, we concatenate its feature vector with the corresponding label vector to form a feature-label vector. We repeat this process for every known sample, and put all the obtained feature-label vectors into $K$ (row by row). In addition, we construct the test sample’s feature-label vector by concatenating its feature vector with a label vector containing all zeros. We put this feature-label vector into $K$ as well. Since we are considering the self-attention mechanism, $V$ and $Q$ are constructed in the same way as for $K$.

Proposition 2: Under Setup 2, self-attention (i.e., Eq. (1) with $K=V=Q$) generates the same label vector as one that is generated from the attention in Setup 1 for the test sample. With $s$ approaching a small value, self-attention can serve as an NN classifier for inferring the test sample’s label.

Since the label vector for the test sample has all zeros at the input, the similarity measures between the test sample and known samples are influenced only by their features. This leads the inference process for the label of the test sample to be essentially the same as shown in Proposition 1. Transformer architecture used in modern large language models, including GPT-3, consists of multiple layers of self-attentions. Below we give more results on stacking self-attentions.

Proposition 3: Under Setup 2, a single layer of self-attention (Eq. (1) with $K=V=Q$) performs one iteration of clustering on feature-label vectors (including the known samples and test sample). $L$ layers of self-attention perform $L$ iterations of clustering. There exists a number $L^{*}$ for the number of layers of self-attention for which the clustering process converges.

Guided by the above theoretical analysis, below we proceed to design the prompt (input) of GPT-3 for transductive inference. We use Setup 2 to guide the prompt construction since GPT-3 uses self-attention: The features and labels of the known samples and the feature of the test sample are put together to feed to GPT-3. According to Proposition 3, stacking self-attentions is functionally more advanced than a nearest neighbor-based classifier. GPT-3 uses not only stacking self-attentions but also numerous pre-trained parameters to augment these attentions. Hence, we expect GPT-3 to be more robust than the conventional methods (e.g., KNN) for transductive inference.

A set of $m$ known samples is provided with their features $F=\{f_{1},f_{2},\dots,f_{m}\}$ and corresponding labels $Y=\{y_{1},y_{2},\dots,y_{m}\}$. A feature vector $f_{test}$ of a test sample is given. The task in this section is to construct a prompt text representation that contains information from $F$, $Y$, and $f_{test}$, which is fed to GPT-3 for inferring the label of the test sample.

Selecting and Ordering Known Samples. As a language model, the original goal of training GPT-3 was to train the model to generate output that is semantically coherent with its input. The data used for training GPT-3 implicitly imposed a prior: The later a text appears in the input prompt (the closer the text to the output text), the larger impact it would impose on the output generation process. Hence, it is essential to put the more representative feature-label texts near the end of the prompt for inferring the test sample’s label.

We compute pair-wise similarities between the features in the set $F$ of the known samples and obtain an affinity matrix S, in which each entry $\textrm{S}_{i,j}$ describes the similarity between samples $i$ and $j$ and is computed as $sim(f_{i},f_{j})$. A cosine similarity function is the default choice for $sim(.,.)$.

For a feature vector $f_{i}\in F$, we define a simple measure of how well $f_{i}$ represents the other known samples: $\textrm{rep}_{i}={\sum_{j=1}^{m}\textrm{S}_{i,j}}$. To select the top $k$ representative samples, one can compute $\textrm{rep}_{i}$ for each $i=1,2,\dots,m$, and choose the largest $k$ representative samples: $\textbf{index}=argsort(\textrm{rep}_{1},\dots,\textrm{rep}_{m},``descend")$, and index is represented as $\textbf{index}[1,2,\dots,k]$. The order of the samples in the prompt for GPT-3 should be in the reverse order of that in the index list, where the most representative sample ($f_{index[1]}$) should be put at the end of the prompt in order to give more influence on the output generation. When dealing with imbalanced classification problems, we perform the above process for the samples in each class, and join them in an interleaved fashion.

Converting Features and Labels to Feature-label Texts. For all the feature vectors $f_{i}$ where $i$ is in the index list computed above, we convert these features to texts in an array-like format. For each feature text thus obtained, we put its corresponding label together with the feature text to form a feature-label text. We then put these feature-label texts together into a long text. More details can be found in the Python-like pseudo-code in Listing 1.1 below.

In this section, we propose two use cases for improving an already-trained vision-based classification model with our proposed GPT4MIA method. The main workflow is illustrated in Fig. 2.

Use Case #1: Detecting Prediction Errors. The first use case of utilizing GPT-3 as a transductive inference method is for detecting prediction errors by a trained vision-based classifier. Conventionally, a validation set is commonly used for comparing and selecting models. Here, we utilize a validation set to provide known samples for transductive inference. Feature vectors in $F$ are obtained from the output probabilities of the vision-based classification model, and labels in $Y$ are obtained by checking whether the classification model gives the correct prediction on each validation sample.

Use Case #2: Improving Classification Accuracy. The second use case aims to improve an already-trained classifier by directly adjusting its predictions. This is a more challenging scenario in which the method not only seeks to detect wrong predictions but also acts to convert them into correct ones. Feature vectors in $F$ are obtained from the output probabilities of the vision-based classification model, and labels in $Y$ are obtained from the validation set for each validation sample.

We utilize the RetinaMNIST and FractureMNIST3D datasets from the MedMNIST dataset [11] for these experiments. We apply a ResNet-50 model trained with the training set as the vision-based classifier, for which the weights can be obtained from the official release.³³3The model weights are obtained from https://github.com/MedMNIST/experiments. We then collect the model’s output probabilities for each validation sample and label it based on whether the prediction is correct. An error detection method is then built based on the information from the validations for classifying the predictions into two classes (being correct or incorrect). The error detection model is then evaluated using the test set working with the same prediction model which was used on the validation (ResNet-50 in this case). We compare our proposed GPT4MIA method on this task with a set of well established inductive methods and transductive methods. From Table 1, one can see that GPT4MIA significantly outperforms the known competing methods for detecting prediction errors from a CNN-based classifier.

We utilize the RetinaMNIST and FractureMNIST3D datasets from MedMNIST [11] for these experiments. ResNet-50 is used as the trained vision-based classification model. The model weights are obtained from the MedMNIST official release. In Table 2, we observe that GPT4MIA performs similarly when comparing with the state-of-the-art transductive inference method Underbagging KNN in balanced accuracy. In Tabel 3, we observe that GPT4MIA performs considerably better in balanced accuracy.

We validate the effect of performing sample selection and ordering described in Section 2.2. In Table 4 and Table 5, we show the performances for the setting without the step of sample selection and ordering. From these results, it is clear that sample selection and ordering is important for better performance when utilizing GPT-3 as a transductive inference tool.

In Table 6, we give additional results of GPT4MIA and other well-known classification models on toy datasets from the scikit-learn package. This experiment serves as a sanity check for our proposed GPT4MIA. In Fig. 3, we visualize the classification results (test sample points) for GPT4MIA and the ground truth. In addition, in Fig. 4, we give visualizations of the test sample points from the experiments for Use Case #1 on the FractureMNIST3D dataset.