GIST Selects Which Answers—and Which Tokens—to Use in Offline AI Training

Not every part of a language model’s answer is equally useful for training. A new preprint proposes GIST, a method that selects both whole responses and individual tokens from a fixed pool, and reports gains on math and coding benchmarks.

A more selective training set

Offline reinforcement learning with verifiable rewards starts with a fixed collection of model-generated answers. An automatic checker labels each answer correct or incorrect, and the model is trained on the resulting data. A common approach applies a training loss to every token in each retained answer, even though some responses—and some steps within them—may be more informative than others.

In a preprint, Xuliang Huang and colleagues propose Gibbs-Targeted Hierarchical Data Selection, or GIST. Its central idea is to make two decisions before training: which answers to keep, and which tokens within those answers should receive a training signal. The method uses the checker’s outcome, the sampling temperature used to generate an answer, and token probabilities assigned by the initial model.

The reasoning behind the selection

The authors analyze an idealized target policy for reinforcement learning with a constraint against departing too far from the initial model. This Gibbs target increases the relative probability of correct answers while retaining a connection to the original policy. The paper examines the mismatch between the initial model and this target in both directions using KL divergence, a measure of difference between probability distributions.

The analysis distinguishes two useful cases. Correct answers can contain steps that the initial model gives too little probability. Incorrect answers that the model generates frequently can reveal behavior to which it assigns too much probability. At the token level, the proposed signals are low-probability tokens in correct answers and high-probability tokens in incorrect ones.

These are selection principles, not perfect labels of individual reasoning steps. In particular, the paper notes that a low token probability alone does not establish a large contribution to the mismatch. Its mathematical selection guarantee concerns mismatch captured at a fixed policy, not a guaranteed improvement after training. GIST approximates the theoretical quantities using signals available before training rather than estimating the ideal target directly.

How GIST builds its dataset

For each prompt, the initial model generates 16 responses at temperature 1.0 and 16 at temperature 0.3. The lower temperature favors responses nearer the model’s most likely outputs. GIST takes correct answers from the first group and incorrect answers from the second. It keeps a prompt only if it has at least four of each, then samples four correct and four incorrect responses.

Within each retained response, the initial model scores its tokens. For a correct answer, GIST applies the training loss to tokens below that answer’s mean token probability; for an incorrect answer, it selects tokens above the mean. Other tokens remain in the context but receive no loss. The selections and masks stay fixed during training.

The researchers combine GIST with two objectives: BOLT, which assigns reward-based weights, and supervised fine-tuning, or SFT, which gives each retained response equal weight. Notably, SFT with GIST imitates selected tokens from incorrect answers as well as correct ones.

What the experiments found

The preprint tests Qwen3-0.6B, Qwen3.5-2B and Gemma-4-E2B-it on GSM8K and MATH-500 math tasks and the MBPP-full coding benchmark. Across both objectives and all three models, GIST raises the average of flexible math accuracy and coding pass@1. It improves 29 of the 30 individual benchmark-and-metric comparisons; the exception is Qwen3.5-2B on MBPP-full with BOLT, where pass@1 falls from 42.40% to 41.80%.

For example, Qwen3-0.6B with BOLT rises from a three-benchmark average of 41.49% to 44.04% with GIST. Its GSM8K strict accuracy rises from 43.29% to 45.79%. Strict math scoring requires the prescribed answer format, while flexible scoring can credit answers extracted from other formats. That distinction matters: for Gemma with SFT, MATH-500 strict accuracy rises from 12.40% to 41.80%, while flexible accuracy rises from 42.00% to 53.00%, indicating that much of the strict-score gain reflects formatting.

In a separate controlled Qwen3-0.6B/GSM8K comparison, all selectors use the same candidate pool and 1,593 eligible prompts. GIST’s 45.79% strict accuracy exceeds seven selection baselines and the 43.82% achieved by training on all candidates from those prompts. It supervises 4.55 million token positions through its selected checkpoint, versus 13.60 million for all candidates—a 66.6% reduction—and uses 2,394 rather than 4,779 optimizer steps. These counts include repeated tokens across epochs; masked tokens still participate in forward and backward passes, and selection requires an additional scoring pass.

Boundaries on the findings

Each reported result comes from a single training run. Paired tests on GSM8K support some comparisons, but not all: against training on all candidates in the controlled pool, the strict-accuracy difference has a 95% bootstrap interval of −0.08 to 4.02 percentage points. The main end-to-end comparison also changes which prompts are used: GIST excludes prompts that do not meet its quotas, including the hardest ones, while the all-candidate baselines train on every training prompt.

The results therefore support GIST as a promising selection strategy within the tested offline setup, not a general guarantee of better or cheaper training. Its performance outside these models, benchmarks, verifier rules and fixed pre-training selections remains untested in the preprint.