Information Projection for Finetuning

Why People Wanted It

An expert trace can contain valuable task information while still looking unlike anything the base model would naturally produce. Training directly on that trace asks the model to imitate both its useful content and its particular wording or path. The question is whether we can retain the task-relevant information while choosing examples that are as compatible with the base model as possible.

The key is to define “same information” as a constraint on trajectories, rather than requiring the exact expert sequence. For a fixed query, let $C$ be the set of trajectories considered equivalent to the expert example: for a reasoning task, this could mean using its information to reach a correct answer; for fact learning, retaining the same factual content.

The Core Move

Restrict the base model’s distribution to the trajectories in $C$. In other words, keep its relative preferences among acceptable trajectories, and assign zero probability to trajectories that do not preserve the required meaning:

$p_C(x) = \frac{p(x)\mathbf{1}[x\in C]}{\sum_{x'\in C}p(x')}.$

The resulting distribution is the information projection target: the closest distribution to the base model among distributions that put all their probability on equivalent trajectories.

Quiz

For a fixed query, what distribution is the information projection target?

A distribution that assigns equal probability to every equivalent trajectory.

The base model distribution with all trajectories retained but their scores shifted toward the expert wording.

The expert trace copied exactly, regardless of the base model.

The base model distribution restricted to equivalent trajectories and renormalized.

$ check

Why the Equivalence Constraint Carries the Information

Here is the part that gets technical: the constraint is not “stay near the expert wording,” but “put no probability outside the set of trajectories that preserve the expert’s task-relevant content.” Let $\mathcal{P}_C$ be all distributions $\pi$ with $\pi(x)=0$ whenever $x\notin C$. This is the formal version of treating correctness or semantic equivalence as the thing that must survive.

The concrete example is a fact-learning trace: if the expert example states a fact, then $C$ contains trajectories judged to retain that same factual content. The exact phrasing may change, but a trajectory that drops or changes the fact is excluded. For reasoning, the criterion can instead admit different trajectories that use the expert information and reach a correct final answer. The definition depends on the task; it does not itself provide a universal semantic judge.

Why Restriction Is the Closest Valid Distribution

Now the derivation: among distributions supported only on $C$, minimizing $\mathrm{KL}(\pi\|p)$ is equivalent to matching the base model’s normalized distribution within $C$. Define $Z=\sum_{x\in C}p(x)$, so $p(x)=Zp_C(x)$ for $x\in C$. Then

$\mathrm{KL}(\pi\|p)=\mathrm{KL}(\pi\|p_C)-\log Z.$

The constant $-\log Z$ does not depend on $\pi$, so the minimum occurs at $\pi=p_C$. In plain terms, once unacceptable trajectories are removed, the closest choice does not invent a new preference among the remaining ones: it preserves the base model’s relative probabilities and renormalizes them.

This makes the direction of the projection important: it minimizes $\mathrm{KL}(\pi\|p)$ over distributions constrained to $C$. It is not a generic claim that any notion of distance, or the reverse KL direction, gives the same answer.

Why the Target Is a Distribution, Not a Rewritten Expert Trace

Here is the deeper implication: the target is a distribution over all equivalent trajectories, not a single canonical rewrite of the expert example. The expert trajectory supplies the information constraint, while the base model supplies the relative weighting inside that constraint. For a dataset, the construction is applied per query and the conditionals can be combined across examples.

The normalization $Z$ is the base model’s total probability mass on $C$. If that mass is tiny, the restricted distribution is still mathematically defined when it is nonzero, but drawing from it may be difficult. That distinction matters: specifying the closest valid distribution says what the target is, not that it is easy to obtain samples from it. So what can turn this target into usable training examples?

Minimal Causal Code

This loop computes the restricted distribution for a finite set of candidate trajectories. It makes the mask-and-renormalize operation explicit; it is not a way to discover semantic equivalence for arbitrary text.

import torch

# One score per candidate trajectory, from the base model.
log_p = torch.tensor([-1.2, -0.4, -2.0, -0.8])
# A task-dependent judge marks which candidates preserve the expert meaning.
equivalent = torch.tensor([True, False, True, True])

masked_log_p = log_p.masked_fill(~equivalent, float("-inf"))
p_C = torch.softmax(masked_log_p, dim=0)

What You Gain, and What You Lose

The gain is a precise target for “keep the expert information, but stay as close as possible to the base model.” It allows multiple trajectories to count as valid and preserves the base model’s relative preferences among them, rather than insisting on exact imitation of one off-policy trace.

The constraint is also the hard part. The set $C$ depends on a task-specific notion of correctness or semantic equivalence; the projection does not decide that notion for you. And when the base model assigns little mass to acceptable trajectories, the target can be difficult to sample from. The projection defines the destination, not a guarantee that reaching it is easy. A separate module covers methods for sampling toward this target.

Quiz

Which statement best describes a limitation of the information projection target?

It chooses a universal semantic-equivalence rule that works for every task.

It requires every equivalent trajectory to match the expert wording exactly.

It defines the target but does not by itself make equivalent trajectories easy to sample.

It removes the base model’s preferences among acceptable trajectories.

$ check