Teaching AI Which Tokens Matter for Product Detection

Peter Bubenik ยท Amazon Research Publication ยท ยท Source
Image for Not all tokens earn the reward: Contrastive credit assignment for production detection

After studying this material, you should be able to:

  1. Explain why credit assignment is a problem in sequence-level reinforcement learning
  2. Describe how GRPO works and its limitation with uniform token weighting
  3. Understand what Contrastive Token Weighting (CTW) is and how it solves the credit assignment problem
  4. Distinguish between CTW-T (teacher) and CTW-S (self/teacher-free) variants
  5. Interpret why precision matters more than recall at catalog scale, and how the two-stage cascade addresses this

Step-by-Step Teaching

Step 1: The Real-World Problem First

Let's ground everything in a concrete scenario before touching any algorithms.

The Business Problem

Imagine you are shopping online for a laptop. The title says:

"UltraBook Pro โ€” 512GB SSD Storage"

But the bullet points say:

"โ€ข 128GB internal storage"

Which is true? You don't know. You might buy it expecting 512GB and receive 128GB โ†’ you return it โ†’ the marketplace loses money and trust.

This is called a production defect: a listing where the title contradicts its own bullet points.

Why Detection is Hard at Scale

Scale FactorConsequence
36 million listings to checkYou cannot do this manually
Even 1% false positive rate= 360,000 good listings wrongly flagged
Flagging good listings = seller complaints, lost revenuePrecision must be extremely high (>99%)

๐Ÿ”‘ Key Insight: At massive scale, false positives are catastrophically expensive, so precision dominates the design.


Step 2: The Two-Stage Cascade Solution

Because no single model perfectly balances speed and accuracy, the system uses two models in sequence:

All 36M Listings
       โ”‚
       โ–ผ
โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”
โ”‚  Stage 1: 4B Model  โ”‚  โ† Cheap, fast, runs on everything
โ”‚  (Detector)         โ”‚  โ† High recall: catches most defects
โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜
       โ”‚ Flags ~X candidates
       โ–ผ
โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”
โ”‚  Stage 2: DeepSeek  โ”‚  โ† Expensive, slow, very accurate
โ”‚  R1 (Verifier)      โ”‚  โ† High precision: confirms true defects
โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜
       โ”‚
       โ–ผ
  Confirmed Defects

Why This Design?

  • Stage 1 (cheap): Must have high recall โ€” better to over-flag than miss real defects
  • Stage 2 (expensive): Provides high precision โ€” filters out false positives
  • Together: >99% precision across all 36M listings

The Optimization Goal

Running Stage 2 on every listing is too expensive. The long-term goal is to make Stage 1 good enough to work alone. Until then, every improvement in Stage 1's recall = more defects caught per expensive Stage 2 call.


Step 3: How Models Are Trained โ€” Background on GRPO

To understand CTW, you first need to understand GRPO (Group Relative Policy Optimization).

What is GRPO?

GRPO is a reinforcement learning method for training language models. Here's the intuition:

  1. The model generates multiple completions (answers) for the same input
  2. Each completion gets a reward score (e.g., "correct" = 1, "incorrect" = 0)
  3. The model learns to produce more completions like the high-reward ones

The Advantage Signal

For each completion, GRPO computes an advantage:

Advantage = (reward of this completion) - (average reward of the group)
  • Positive advantage โ†’ this completion was better than average โ†’ reinforce it
  • Negative advantage โ†’ this completion was worse than average โ†’ suppress it

The Critical Problem: Uniform Credit Assignment

GRPO takes that one number (the advantage) and applies it equally to every token in the completion.

Example:

Completion: "The  listing  has  a  contradiction  because  the  title  says  512GB"
Advantage:   +0.8  +0.8  +0.8  +0.8     +0.8        +0.8   +0.8  +0.8   +0.8   +0.8

But think about this: the word "contradiction" is the pivotal token that determines correctness. The filler words "The listing has a" contribute almost nothing to whether the answer is right or wrong.

๐Ÿ”‘ Key Problem: Uniform credit assignment wastes training signal on unimportant tokens and dilutes the signal on important ones.


Step 4: Contrastive Token Weighting (CTW) โ€” The Core Innovation

The Core Idea

CTW asks: "Which tokens actually distinguish correct completions from incorrect ones?"

Those tokens should receive more credit. Filler tokens should receive less.

How CTW Works โ€” Step by Step

Step 1: Generate a group of completions

For one input, sample multiple completions:

Completion A: "...contradiction because title says 512GB but bullet says 128GB" โ†’ Reward: 1 (correct)
Completion B: "...no contradiction, both say the same storage"                  โ†’ Reward: 0 (incorrect)
Completion C: "...contradiction in the storage specification mentioned"         โ†’ Reward: 1 (correct)
Completion D: "...listing appears consistent to me"                             โ†’ Reward: 0 (incorrect)

Step 2: The reward partitions completions into two groups

Correct group:   {A, C}
Incorrect group: {B, D}

Step 3: Run a frozen scorer on each token position

A frozen scorer (a model that doesn't update during training) reads both groups and asks:

"At this token position, can I tell whether I'm reading a correct or incorrect completion?"

  • If yes โ†’ this position is discriminative โ†’ assign high weight
  • If no โ†’ this position is non-discriminative โ†’ assign low weight

Step 4: Reweight the gradient update

Standard GRPO:
Token:    "The"  "listing"  "has"  "contradiction"  "because"  "512GB"
Weight:    1.0     1.0       1.0       1.0             1.0        1.0

CTW:
Token:    "The"  "listing"  "has"  "contradiction"  "because"  "512GB"
Weight:    0.2     0.3       0.2       2.1             0.8        1.9

The model now learns more from the tokens that actually matter.


Step 5: Two Variants of CTW

CTW-T (Teacher variant)

โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”         โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”
โ”‚  Student Model  โ”‚ โ”€โ”€โ”€โ”€โ”€โ”€โ–บ โ”‚  Frozen Teacher Model โ”‚
โ”‚  (being trained)โ”‚         โ”‚  (stronger, separate) โ”‚
โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜         โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜
                                      โ”‚
                              Provides token weights
  • Uses a separate, stronger model as the scorer
  • More powerful signal, but requires having a stronger model available

CTW-S (Self/Teacher-Free variant)

โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”         โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”
โ”‚  Student Model  โ”‚ โ”€โ”€โ”€โ”€โ”€โ”€โ–บ โ”‚  Frozen Copy of Student  โ”‚
โ”‚  (being trained)โ”‚         โ”‚  (same model, not updated)โ”‚
โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜         โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜
                                      โ”‚
                              Provides token weights
  • Uses a frozen copy of itself as the scorer
  • No external model needed
  • This is the one actually deployed in production

๐Ÿ”‘ Why CTW-S works: Even the student model, when frozen and used contrastively, can identify which tokens separate correct from incorrect completions. The key signal comes from the reward-induced partition, not from the scorer being powerful.


Step 6: What Actually Drives the Gain?

The paper runs a controlled experiment on math problems to isolate causes.

The Ablation Question

Is the improvement from:

  • (A) The reweighting mechanism itself? OR
  • (B) The fact that completions are split into correct/incorrect groups (the partition)?

The Finding

The reward-induced partition is what drives the gain, not the reweighting itself.

This means the key innovation is using the correct/incorrect split as a contrastive signal to identify important tokens โ€” the specific mathematical form of reweighting matters less.

What matters:     Correct group  vs.  Incorrect group
                       โ†“                    โ†“
              "These tokens appear      "These tokens appear
               in correct answers"       in wrong answers"
                       โ†“
              Contrastive difference = credit signal

Step 7: Results and What They Mean

MetricSFT BaselineCTW-SImprovement
F1 Score83.785.9+2.2
Recall~82.5~85.9+3.4
Precisionmaintainedmaintained~0

Interpreting the Results

  • Recall improved by +3.4: The detector now catches more real defects
  • Precision maintained: We didn't introduce more false positives
  • Gain concentrated in recall: This is exactly what we want โ€” Stage 1 should over-catch, Stage 2 filters

๐Ÿ”‘ Business translation: For every 1,000 verifier (Stage 2) calls, CTW-S finds more true defects than the baseline, making the expensive stage more efficient.


Summary: The Full Picture

PROBLEM:
  Title vs. bullet contradictions mislead shoppers at 36M listing scale
  
CONSTRAINT:
  Precision must be >99% (false positives are too costly)
  
SOLUTION ARCHITECTURE:
  Two-stage cascade: cheap detector โ†’ expensive verifier
  
OPTIMIZATION TARGET:
  Improve detector recall (more defects caught per verifier call)
  
METHOD:
  CTW = Contrastive Token Weighting for GRPO
  
HOW CTW WORKS:
  1. Generate group of completions
  2. Reward splits them: correct vs. incorrect
  3. Frozen scorer finds tokens that discriminate the two groups
  4. Upweight those tokens in gradient update
  
KEY INSIGHT:
  The partition (correct vs. incorrect) is the real signal
  Not all tokens deserve equal credit for a sequence-level reward
  
RESULT:
  +2.2 F1, +3.4 Recall, precision maintained
  Teacher-free variant (CTW-S) works in production

Self-Check Questions

  1. Why does a 1% false positive rate become catastrophic at 36M listings?
  2. What is the fundamental limitation of GRPO's uniform credit assignment?
  3. How does CTW use the reward signal to assign per-token credit?
  4. What is the difference between CTW-T and CTW-S, and why is CTW-S preferred in production?
  5. Why is recall the most important metric for Stage 1 of the cascade?
  6. What does the controlled math experiment tell us about why CTW works?

More to study