Relative Rewards for Typed Extraction Decisions

Why People Wanted It

A typed extraction model chooses among queried types and a null class meaning no extraction. For a span such as “Paris,” nearby score settings might select the correct city type, an incorrect person type, or null. A span pair makes the analogous choice among relation types and no relation.

Likelihood supervision teaches annotated choices. The additional goal here is to strengthen correct decisions against alternatives the model can also produce. So the real question is not just how to score the gold label: how can the consequences of nearby discrete decisions supply a training gradient without a separate critic?

The Core Move

Independently perturb a candidate’s scores, reward the resulting decisions relative to their peers, and use a score-function gradient to make higher-reward score samples more likely.

Let $\mathbf z_c$ be candidate $c$’s logits, including null. Draw $S>1$ independent samples with fixed noise scale $\sigma>0$:

$$\boldsymbol\epsilon_c^{(k)}\sim\mathcal N(\mathbf 0,\mathbf I),\qquad \mathbf x_c^{(k)}=\mathbf z_c+\sigma\boldsymbol\epsilon_c^{(k)},\qquad a_c^{(k)}=\arg\max_m x_{c,m}^{(k)}.$$

Thus the sampled score vector has distribution $\pi(\cdot\mid\mathbf z_c)=\mathcal N(\mathbf z_c,\sigma^2\mathbf I)$. The noise scale is applied once: standard-normal noise multiplied by $\sigma$, not noise already scaled by $\sigma$ multiplied again.

Why Score Perturbations Create Competing Decisions

Here is the part that gets technical: the stochastic object is the entire Gaussian score vector, while the extracted class is a deterministic function of that vector.

Write the class set as $\{0\}\cup\mathcal K$, where $0$ means null and $\mathcal K$ contains the queried types included in the training decision. Types with unknown annotation status are excluded. Candidates ignored because their missing labels are ambiguous must also be excluded from the loss.

For “Paris,” independently perturbing the city, person, and null logits can move the winner across class boundaries. Sampling a categorical distribution from normalized logits would be a different mechanism. Likewise, generating one noise vector and expanding it along a sample dimension would repeat a decision rather than generate independent alternatives.

Given the candidate’s gold class set $G_c$, assign the decision reward

$$R_c^{(k)}= \begin{cases} 0 & a_c^{(k)}=0,\\ 1 & a_c^{(k)}\in G_c\setminus\{0\},\\ -1 & \text{otherwise}. \end{cases}$$

The null case takes precedence even when $G_c=\{0\}$. A false-positive type on such a candidate receives negative reward; choosing null receives zero. How can these rewards become a useful comparison rather than isolated scores?

Why the Leave-One-Out Baseline Is Relative

Now the version for readers who want more of the machinery: each sample’s baseline excludes that sample, so it depends on independent peer draws rather than on its own sampled score vector.

Define

$$b_c^{(-k)}=\frac{1}{S-1}\sum_{j\ne k}R_c^{(j)},\qquad \delta_c^{(k)}=R_c^{(k)}-b_c^{(-k)}.$$

For sampled decisions city, person, and null on a gold-city candidate, the rewards are $1,-1,0$. Their advantages are respectively $3/2,-3/2,0$. These are consequences of the reward rule, not empirical performance measurements.

An advantage is relative, not an additional label. Null can have positive advantage when its peers choose wrong types, or negative advantage when they choose correct types. If every reward agrees, all advantages vanish.

Conditional on the current logits, the peer baseline is independent of sample $k$. Subtracting it does not change the expected score-function gradient because the expected gradient of a normalized distribution’s log density is zero. It supplies a baseline without a learned critic. But neither a reward nor an argmax provides an ordinary differentiable path—where does the actual gradient enter?

Why Likelihood and Reward Stay in One Objective

Here is the final technical step: differentiate the Gaussian log density of each fixed sampled vector, not the discrete winning class and not the path used to generate that vector.

A per-candidate surrogate for minimizing negative expected reward is

$$\widehat L_{\mathrm{reward},c} =-\frac{1}{S}\sum_k \operatorname{stopgrad}(\delta_c^{(k)}) \log\pi\!\left(\operatorname{stopgrad}(\mathbf x_c^{(k)})\mid\mathbf z_c\right).$$

Here $\operatorname{stopgrad}$ means keep the value but block differentiation through its computation. For fixed $\sigma$, the Gaussian score is

$$\nabla_{\mathbf z_c}\log\pi(\mathbf x\mid\mathbf z_c) =\frac{\mathbf x-\mathbf z_c}{\sigma^2}.$$

Consequently,

$$\nabla_{\mathbf z_c}\widehat L_{\mathrm{reward},c} =-\frac{1}{S\sigma^2}\sum_k \delta_c^{(k)}(\mathbf x_c^{(k)}-\mathbf z_c).$$

Under gradient descent, a positive advantage pulls the mean logits toward that sampled vector; a negative advantage pushes them away. This changes the distribution of future decisions without differentiating through argmax. Leaving the sample attached would instead make $\mathbf x-\mathbf z=\sigma\boldsymbol\epsilon$ constant with respect to the mean, cancelling the intended gradient.

Supervised learning remains present. With $p_c$ the normalized class distribution, its candidate loss is

$$\ell_c=-\log\sum_{g\in G_c}p_c(g).$$

This allows multiple gold types while each sampled decision still selects one class. Let $w_c$ be the candidate’s supervision weight, with null candidates down-weighted. The combined head loss is

$$L_{\mathrm{head}}= \frac{\sum_c w_c\left(\ell_c+\rho\widehat L_{\mathrm{reward},c}\right)}{\sum_c w_c}.$$

Entity and relation losses use this construction, with their respective null weights, and combine as $L=L_{\mathrm{ent}}+\lambda L_{\mathrm{rel}}$. Likelihood can still contribute when sampled rewards tie. What does this look like with the sampling and gradient boundaries made explicit?

Minimal Causal Code

This is one optimizer update for one extraction head and one instance. model(features) returns logits shaped [C, M]; class zero is null. gold_mask has the same shape and marks all supervised gold classes, including null for null candidates. Unknown type columns and ignored candidates have already been removed. weights contains the supervision weights. Relation and entity head losses can be combined before the same backward pass.

import torch

assert samples > 1 and sigma > 0
optimizer.zero_grad()
logits = model(features)                    # [C, M]
assert gold_mask.shape == logits.shape
assert gold_mask.dtype == torch.bool
assert gold_mask.any(dim=-1).all()
assert weights.sum() > 0

# Supervision: total probability of all gold classes.
log_p = logits.log_softmax(dim=-1)
label_loss = -torch.logsumexp(
    log_p.masked_fill(~gold_mask, -torch.inf), dim=-1
)                                          # [C]

policy = torch.distributions.Normal(logits, sigma)
reward_rows, log_density_rows = [], []
for k in range(samples):
    # Fresh independent noise for every sample, candidate, and class.
    # Keep the observation fixed when differentiating its log density.
    with torch.no_grad():
        score_sample = logits + sigma * torch.randn_like(logits)
        decision = score_sample.argmax(dim=-1)  # [C]
        is_gold = gold_mask.gather(
            dim=1, index=decision[:, None]
        ).squeeze(1)                            # [C]
        reward = torch.where(
            decision == 0,
            torch.zeros_like(decision, dtype=logits.dtype),
            torch.where(is_gold, 1.0, -1.0),
        )

    # Gaussian vector log density, NOT categorical class log probability.
    log_density = policy.log_prob(score_sample.detach()).sum(dim=-1)
    reward_rows.append(reward)
    log_density_rows.append(log_density)

rewards = torch.stack(reward_rows)          # [S, C]
log_densities = torch.stack(log_density_rows)
baselines = (rewards.sum(dim=0, keepdim=True) - rewards) / (samples - 1)
advantages = (rewards - baselines).detach()
reward_loss = -(advantages * log_densities).mean(dim=0)  # [C]

loss = (weights * (label_loss + rho * reward_loss)).sum() / weights.sum()
loss.backward()
optimizer.step()                           # Update the score-producing model.

The gather operates separately for each candidate: each row looks up its own sampled class. Fresh noise inside the loop prevents accidental reuse across samples. The detached observations and advantages leave the intended gradient through the Gaussian mean logits into the model.

What You Gain, and What You Lose

You gain a decision-level training signal that distinguishes correct types, incorrect types, and no extraction. The samples themselves provide the reference reward, without a separate critic. Likelihood supervision continues to anchor training to annotated answers.

You pay for repeated sampling and an additional noisy gradient estimate. The perturbation scale controls which alternatives are explored, while the reward weight controls their influence alongside supervision. Identical sampled rewards contribute no relative-reward gradient, even when every sampled decision is wrong.

Null’s zero reward is not a guarantee of neutral training pressure: its advantage depends on its peers. Nor does this objective resolve unknown annotations or replace decoding thresholds and calibration. It supplements supervised extraction rather than making the same objective simply faster or automatically calibrated.

Quiz

What does the leave-one-out baseline compare for a sampled candidate decision?

Its reward against the reward of the gold class only

Its reward against a separately trained critic’s prediction

Its reward against the average reward of the other samples for that candidate

Its score against the average score of all candidates in the batch

$ check