Surprise is inherently sequential. Conditioned on the past, a new piece of information is surprising if it “contains information that we didn’t expect”. This is almost a tautology. But it seems that this notion of time, or at the very least sequencing, is crucial to the concept of surprise.
Intuitively, if we use a distribution $p$ to measure our own beliefs in the world, we can use $p(x\given c)$ to represent the likelihood of some outcome $x$ based on what we’ve seen before, the context $c$, i.e. the past. For example, $x$ could be the current state $x_t$ and $c$ could be the entire history of the past $x_{\lt t}$. A reveal is surprising only if it updates belief at the time of reveal.
In information theory, we can define a quantity called “surprisal”:
$$S(x) = -\log p(x\given c)$$where $c$ is the context, i.e. the past. If there is no notion of the past, then there is no conditioning, and hence there is no notion of surprise. Notably, if $p(x\given c)=1$ so that we know to expect it completely, then $x$ really is not surprising at all, so $S(x)=0$.
This is famously thought of as “the number of bits an optimal coding needs to encode an outcome $x$. Here an optimal coding is one that minimizes the expected number of bits needed to transmit an outcome under distribution $p$.
The intuition here is that if an outcome were very likely (high probability), then it makes sense to assign a very short encoding to the outcome. So the higher the surprisal of the outcome, the more bits the coding scheme would use to encode it. Thus in this sense, a surprising $x$ contains a lot of bits of information.
Training a model to be surprising
The easiest way to train a model to be surprising is to train it to maximize the average surprise over the whole dataset (i.e. cross-entropy).
$$\text{maximize}_{\theta}\ \ev[x\sim D]{-\log \pi_\theta(x_t \given x_{\lt t})}$$where $D$ is the data distribution. But this will just give us gibberish, since we’ll get a uniform distribution over all the tokens in the data.
Let’s try something a bit smarter. Suppose we have a pre-trained base language model $\pi_0$. It was trained by minimizing the next-token surprise of all the text that it has seen. Thus near convergence, it should not be surprised by anything that is “in-distribution”, e.g. all of the information on the internet. Now let’s say it’s also been further post-trained to be good at specific tasks, e.g. coding, tool-use, writing, judging, and this results in a teacher model $\pi_T$.
We want to now train a model $\pi_\theta$ to generate something surprising relative to the our very-capable teacher model. So naively, maybe you can try to have it learn to maximize surprise:
$$\text{maximize}_{\theta}\ \ev[x\sim \pi_\theta]{-\log \pi_T(x_t \given x_{\lt t})}.$$In other words, train $\pi_\theta$ to generate data that is surprising to the base model $\pi_0$.
But still, it might quickly learn to generate gibberish, since a fully uniform distribution is likely quite “surprising” to the teach distribution $\pi_T$.
Instead, we want something more constrained, like:
$$\text{maximize (surprise) subject to (data sampled from a manifold of realistic data}).$$Constraining the surprise
Warm up
So how do we constrain our model to only generate data from to a manifold of reasonable-looking sequences? We can use the distribution we learned!
$$\text{maximize}_\theta\ \ev[x\sim \pi_\theta]{-\log \pi_T(x_t\given x_{\lt t})} - \beta \kldiv{\pi_\theta}{\pi_0} $$where we maximize our generator $\pi_\theta$ to maximize the surprise of the teacher $\pi_T$, in a way that keeps the generated data close to the manifold of some reference model $\pi_0$ to prevent it from generating gibberish.
With this new objective, we’re actually doing reinforcement-learning where the reward function is the surprise of the teacher $\pi_T$!
The optimal policy that maximizes surprise subject to is:
$$\pi^* (x_t\given x_{\lt t}) \propto \pi_0(x_t\given x_{\lt t})\pi_T(x_t\given x_{\lt t})^{-1/\beta} = \frac{\pi_0(x_t\given x_{\lt t})}{\pi_T(x_t\given x_{\lt t})^{1/\beta}}$$A few observations:
- If $\beta\to\infty$ , then we just sample from the base model $\pi_0$. Remember that $\beta$ is how strong we’re sticking to the base model, so this makes sense.
- If $\beta\to1$, we sample proportional to the likelihood ratio between the base and the teacher.
If the base and the teacher models are different, sampling from this distribution will give us base-plausible but teacher-undervalued outputs. In other words, samples remain on the data manifold, but live in the teacher’s blind spots.
Open Q: What happens if we sample from this model?
Aside: what if the teacher were the base model?
If the base model is just the teacher, then $\pi_0 = \pi_T$, and we see that
$$\pi^* (x_t\given x_{\lt t}) \propto \pi_T(x_t\given x_{\lt t})^{1-1/\beta}$$which is equivalent to using a higher temperature in sampling when $\pi_T$ is the softmax distribution.
Now optimize the surprise of the entire sequence
Now, actually we probably want to maximize the surprise of the entire sequence, rather than just the conditional at every step:
$$\text{maximize}_\theta\ \ev[x\sim \pi_\theta]{-\log \pi_T(x)} - \beta \kldiv{\pi_\theta}{\pi_0} $$This also has a closed form solution if we optimize over all distributions on full sequences, and it’s the Gibbs/Donsker-Varadhan variational form with optimal policy
$$\pi^*(x) \propto \pi_0(x) \pi_T(x)^{-1/\beta}.$$If we instead maximize this objective over our neural network architecture, then we will end up with something different.
In the case where $\pi_0 = \pi_T$, we end up with a power distribution, much like what Karan & Du, 2025 did. In their paper, they sharpened the distribution of the base model $\pi_0$. Here, we wonder what would happen if we flatten the post-trained model $\pi_T$.