Should a language model like GPT-3 be directly used to perform classification?

Question

The OpenAI API enables classification by sampling from GPT-3 given a prompt. Is estimating posterior probabilities a more statistically sound approach? Below is a specification of what "estimating posterior probabilities" could mean here. I'd like feedback on how sensible it is (see the question at the bottom).

Definitions

• $\mathcal{X}$ is the set of all possible sequences of tokens. $x \in \mathcal{X}$ is a sequence of tokens. For example, a movie review could be $x = (\text{A}, \text{ thrill}, \text{ ride}, \text{ for}, \text{ the}, \text{ ages}, \text{!}, \text{ --}, \text{Peter}, \text{ K}, \text{ Rosenthal})$.

• $\Pr_{\theta}$ is a language model, i.e., a probability distribution over $\mathcal{X}$ parametrized by $\theta$.

• $Y$ is a random variable for the class an observation belongs to, where the set of classes is $\mathcal{Y}$. Let's allow $\mathcal{Y} \subset \mathcal{X}$. So if I'm classifying movie reviews, $\mathcal{Y} = \{ \text{Just another superhero movie}, \text{Generic hype} \}$ represents a classification task.

• The function $\text{prompt}: \mathcal{X} \rightarrow \mathcal{X}$ takes a sequence of tokens (which must include $x$, and optionally $\mathcal{Y}$) and returns a new sequence of tokens. A prompt asks a language model to perform the task.

  • See this set of documentation for more examples and detail. (You need to pay to hit OpenAI API endpoints.)

Proposed method

GPT-3's estimates, $\Pr_\theta(x_i \: | \: x_{1:i-1})$, are great. So simply selecting the completion from $\mathcal{Y}$ with the highest posterior probability should be more direct than sampling. For example, we'd query the language model for its probability of this prompt and completion:

This movie review
"""
A thrill ride for the ages! --Peter K Rosenthal
"""
is best categorized as {completion}.

where {completion} is replaced w/ "Just another superhero movie", and then replaced w/ "Generic hype".

I'll start w/ a false start to formalizing the meaning of posterior probability here, and then adjust it. I'm dubious about the adjustment though, so I'd like feedback.

First some new definitions: let $c \in \mathcal{Y}$ be a completion containing $n$ tokens, and $p = \text{prompt}(x)$ be the prompt. In the example above, the prompt is the full sequence of tokens up until {completion}.

False start: it's tempting to simply do $\arg\max_{c \in \mathcal{Y}} \Pr_\theta(c \: | \: p)$. This is wrong b/c longer completions trivially result in lower probabilities. It also asks the language model to do too much work, as it doesn't incorporate the prior $\Pr(y)$ from your specific distribution. It implicitly comes from GPT-3's training distribution as $\Pr_\theta(c)$. (But this prior can be manually set for single-token completions.)

Now for the adjustments: start w/ framing the classifier as

$$ \begin{equation} \arg\max_{y \in \mathcal{Y}} \Pr(y \: | \: p, c) = \arg\max_{y \in \mathcal{Y}} \text{Pr}_\theta(p, c \: | \: y) \Pr(y). \end{equation} $$

This reframing fixes the issue with the prior. (Edit: it seems to not be enough when there's a uniform prior. See the edit in the updated answer.) To fix the completion length issue, let's replace $\Pr_\theta(p, c \: | \: y)$ with an average per completion token. To construct that, we'll first express $\Pr_\theta(p, c \: | \: y)$ as

$$ \begin{align} \text{Pr}_\theta(p, c \: | \: y) &= \text{Pr}_\theta(p \: | \: y) \text{Pr}_\theta(c \: | \: p, y) && \text{probability chain rule} \\ &= \text{Pr}_\theta(p) \text{Pr}_\theta(c \: | \: p, y) && P=p \perp Y=y \\ &= \text{Pr}_\theta(p) \prod_{i=1}^{n} \text{Pr}_\theta(c_i \: | \: p, y, c_{1:i-1}) && \text{probability chain rule} \\ &= \text{Pr}_\theta(p) \exp \Bigg\{\sum_{i=1}^{n} \log\text{Pr}_\theta(c_i \: | \: p, y, c_{1:i-1}) \Bigg\}. \end{align} $$

Note that we can ignore $\Pr_\theta(p)$ as it doesn't depend on $y$. The completion does depend on the class b/c it's set to the class.

Define the average likelihood (i.e., the average inverse perplexity^1,2) as

$$ \begin{equation*} \bar{\text{Pr}}_\theta(c \: | \: p, y) = \exp \Bigg\{\frac{1}{n} \sum_{i=1}^{n} \log\text{Pr}_\theta(c_i \: | \: p, y, c_{1:i-1}) \Bigg\}. \end{equation*} $$

Finally, estimate the posterior probability of a class, AKA the Completion After Prompt Probability, as

$$ \begin{equation*} \text{CAPPr}(y \: | \: p, c) = \bar{\text{Pr}}_\theta(c \: | \: p, y) \Pr(y) / Z \end{equation*} $$

where $Z$ is the normalizer.

Example

Let's perform the estimation for the $x$ and $\mathcal{Y}$ movie review example. $p = \text{prompt}(x)$ is this sequence of tokens:

This movie review
"""
A thrill ride for the ages! --Peter K Rosenthal
"""
is best categorized as 

Example computation for the second class:

$$ \begin{equation*} \bar{\text{Pr}}_\theta(c \: | \: p, \text{Generic hype}) = \exp \Bigg\{ \frac{1}{2} \Big[ \log \text{Pr}_\theta(\text{Generic} \: | \: p) + \log \text{Pr}_\theta(\text{hype} \: | \: p, \text{Generic}) \Big] \Bigg\}. \end{equation*} $$

Say I magically know the class distribution for my data:

$$ \begin{align*} \Pr(\text{Just another superhero movie}) = 1/3 \\ \Pr(\text{Generic hype}) = 2/3. \end{align*} $$

Then the normalizer is

$$ \begin{equation*} Z = \frac{1}{3} \bar{\text{Pr}}_\theta(c \: | \: p, \text{Just another superhero movie}) + \frac{2}{3} \bar{\text{Pr}}_\theta(c \: | \: p, \text{Generic hype}). \end{equation*} $$

Finally, the posterior probability estimates are

$$ \begin{align*} \text{CAPPr}(\text{Just another superhero movie} \: | \: p, c) &= \frac{1}{3} \bar{\text{Pr}}_\theta(c \: | \: p, \text{Just another superhero movie}) / Z \\ \text{CAPPr}(\text{Generic hype} \: | \: p, c) &= \frac{2}{3} \bar{\text{Pr}}_\theta(c \: | \: p, \text{Generic hype}) / Z. \end{align*} $$

Question

Did it make sense to reframe the classifier from

$$ \arg\max_{c \in \mathcal{Y}} \text{Pr}_\theta(c \: | \: p) \\ \text{to} \\ \arg\max_{y \in \mathcal{Y}} \text{CAPPr}( y \: | \: c, p) ? $$

References

1. Wikipedia page on perplexity per word
2. huggingface's perplexity documentation

Answer 1

I evaluated the proposed method on the Choice of Plausible Alternatives (COPA) task. I picked this task from the SuperGLUE NLP benchmark b/c each $c \in \mathcal{Y}$ is always a few tokens long.

See the evaluation notebook here.

I still need to empirically evaluate calibration, and evaluate on more classification tasks. But a partial, empirical answer to the question in the title and the "Did it make sense..." question is: yes. Zero-shot accuracy on COPA train + validation is 0.92-0.95 depending on the prompt, which is comparable to text generation + post-processing. (Majority class gives 0.5 accuracy.)

Edit: after some more experimentation, it seems like the issue with the prior coming from $\Pr_\theta(c)$ is not fully resolved. In other words, the average likelihood of completions which are really common—regardless of the prompt—is often too high. I hypothesized that this would be an issue in this comment.

The solution I'm rolling with for now is to do something along the lines of replacing the marginal class probability in Bayes' rule. (I did a similar thing with some success in this separate work. See the section A simple fix.)

We can introduce a "discount" hyperparameter which controls how much $\Pr_\theta(c)$ counts toward the average likelihood. Modify

$$ \begin{equation*} \text{Pr}_\theta(c \: | \: p) \propto \text{Pr}_\theta(p \: | \: c) \text{Pr}_\theta(c) \end{equation*} $$

to

$$ \begin{equation*} \log\text{Pr}_\theta(p \: | \: c) \text{Pr}_\theta(c)^\gamma = \log\text{Pr}_\theta(p \: | \: c) + \gamma\log\text{Pr}_\theta(c). \end{equation*} $$

Skipping some steps (and ignoring the unease with the ill-defined $\text{Pr}_\theta(p \: | \: c)$ term), this fudged expression is incorporated into the average likelihood as:

$$ \begin{equation*} \bar{\text{Pr}}_\theta(c \: | \: p, y) = \exp \Bigg\{\frac{1}{n} \sum_{i=1}^{n} \log\text{Pr}_\theta(c_i \: | \: p, y, c_{1:i-1}) + \gamma \log\text{Pr}_{\theta}(c_i \: | \: y, c_{1:i-1}) \Bigg\}. \end{equation*} $$

Setting $\gamma = 0$ obviously applies no discount. Setting $0 < \gamma < 1$ applies some discount. Setting $\gamma = 1$ is equivalent to dropping $\Pr_\theta(c)$, which should directly counteract the issue. Increasing $\gamma > 1$ further penalizes completions with higher priors, which doesn't seem sensible.