Incremental Prediction for Tractable Information Estimation

Why People Wanted It

The problem is the price of asking a generative model for a large target's likelihood. A pixel may be easy to score, but estimating the likelihood of an image with a diffusion model can require intractably many Monte Carlo estimates. If the representation objective needs a pointwise mutual-information (PMI) target for the whole view, the target itself can become the bottleneck.

So the real question was never just how to define shared information. It was how to estimate it without scoring the entire target in one shot.

The Core Move

Break the target into successive pieces and train against each piece's conditional information increment; when those increments are added, they recover the whole-target PMI:

$$\operatorname{PMI}(x;y_{1:k})=\sum_{i=1}^{k}\left[\log p(y_i\mid x,y_{<i})-\log p(y_i\mid y_{<i})\right].$$

Each term asks how much better the next piece is predicted with $x$ than without it, given the pieces already seen.

Quiz

Why can successive conditional PMI increments recover the whole-view PMI?

The target pieces are assumed to be statistically independent.

The conditional likelihood ratios multiply to the full likelihood ratio by the chain rule.

The method removes the need to estimate any conditional likelihoods.

Each increment estimates a different unrelated PMI and their average approximates the whole.

$ check

Why Whole-View Likelihood Is Hard to Estimate

Here is the part that gets technical: a large target makes a likelihood-ratio estimate costly because the model must account for the target's full information content. The canonical example is an image: scoring one pixel can be simple, while estimating the image likelihood with a diffusion model may require many Monte Carlo estimates. The issue is not that a full-view PMI is conceptually invalid; it is that calculating its likelihood terms can be impractical.

For a target $y=(y_1,\ldots,y_k)$, the whole-view PMI compares $p(y\mid x)$ with $p(y)$. Factoring both probabilities in the same token order gives

$$\log\frac{p(y\mid x)}{p(y)}=\sum_{i=1}^{k}\log\frac{p(y_i\mid x,y_{<i})}{p(y_i\mid y_{<i})}.$$

In words, each step compares the next-token prediction with and without the conditioning view, while holding the earlier target pieces fixed. This turns one large likelihood ratio into local conditional ratios. What remains difficult is choosing how those local terms are estimated.

Why the Increments Telescope

Now the version for readers who want more of the machinery: the increments telescope because each conditional probability is a factor in the chain rule, not because the pieces are independent. Multiplying the conditional ratios cancels the shared prefix conditionals in the numerator and denominator, leaving exactly $p(y\mid x)/p(y)$; taking logs yields the sum above.

That distinction matters: the target pieces may depend strongly on one another. The factorization accounts for that dependence by conditioning each next piece on the prefix $y_{<i}$. Bootstrapping supervises these incremental PMI contributions rather than asking for the full-view PMI estimate at once. The exact sequence of pieces therefore defines the sequence of conditional prediction problems; it does not assume an independent-token target.

Why Timestep Weighting Changes the Target

Here is a deeper layer: the information being estimated can also be selected by corruption scale, so the objective need not emphasize only exact, fine-grained agreement. At corruption level $s$, the scale-specific PMI measures dependencies that remain predictive after that corruption; increasing $s$ destroys progressively finer information.

A weighted target averages these scale-specific quantities:

$$\operatorname{PMI}_{q}(x,y)=\mathbb{E}_{s\sim q(s)}[\operatorname{PMI}(x,y_s)].$$

In words, $q(s)$ chooses which corruption scales contribute more to the representation target. Each scale-specific PMI can be estimated through a diffusion ELBO, integrating a conditional denoising advantage over another timestep $t$. Since the corruption at $s$ and the diffusion corruption at $t$ are both Gaussian, they can be composed into an effective timestep weighting, described by timestep sampling $f(t)$ and timestep weighting $g(t)$. Combining that weighting with bootstrapped increments gives a target that is local in the predicted pieces and adjustable in the information scale it emphasizes. What remains open is how a particular tokenization and weighting should be chosen for a given task.

Minimal Causal Code

This loop makes the state visible: each step adds the next conditional log-likelihood ratio, using the already observed target prefix. It is a schematic objective calculation, not a diffusion estimator.

import torch

def incremental_pmi(x, y_tokens, log_prob):
    # log_prob(target_piece, context) returns a scalar log probability.
    total = torch.zeros((), device=x.device)
    prefix = []
    increments = []

    for piece in y_tokens:
        with_x = log_prob(piece, (x, prefix))
        without_x = log_prob(piece, (None, prefix))
        increment = with_x - without_x
        total = total + increment
        increments.append(increment)
        prefix.append(piece)

    return total, increments

What You Gain, and What You Lose

You gain a likelihood-ratio target made from successive local predictions instead of one full-view estimate that can be intractable. The increments still recover the full PMI when summed, and timestep weighting lets the objective favor information that survives selected corruption scales.

You do not get a free estimate: each conditional prediction still has to be modeled, and the decomposition depends on how the target is ordered or tokenized. Timestep weighting also changes which information is emphasized rather than simply making the original target cheaper. The method trades one difficult global likelihood calculation for a structured family of conditional ones.

Quiz

What does timestep weighting control in this objective?

The number of target tokens in the sequence.

Which corruption scales, and thus which surviving information, receive more weight.

Whether the target pieces are conditionally independent.

The order in which target tokens are generated, without changing the objective.

$ check