Teaching a Post-Trained Model to Flip a Creativity Switch

Paper: 2608.07460 Authors: Ananya Sahu, Mohit Bansal, Elias Stengel-Eskin Categories: cs.CL, cs.AI

The Gap

Everyone in the field knows the tax: SFT plus preference optimization makes a model helpful, safe, and formatted — and also makes it write the same story ten times in a row. Mode collapse under RLHF/DPO is well documented, and the existing fixes fall into three families, each with a real cost.

Decoding-time tricks (higher temperature, nucleus/typical/min-p sampling) buy diversity by pushing the model into its own tail, which means quality degrades roughly monotonically with variety. Multi-model inference (contrastive decoding against the base model, proxy tuning, logit interpolation, ensembles of differently-tuned checkpoints) genuinely works but requires holding two or more models in memory at inference, and the diversity knob is global rather than local. Data-side approaches (persona conditioning, seed-word injection, diversity-aware DPO) keep single-model inference but bake the diversity in as a fixed style rather than something the model can dial up exactly where it matters.

There is also a measurement gap. The field mostly reports distinct-n / self-BLEU (lexical) and embedding cosine spread (semantic). Both are blind to the thing readers actually notice in stories: whether the plot shape differs. Ten stories can use disjoint vocabulary and still be “orphan discovers power, mentor dies, orphan wins.”

This paper’s bet: make creativity a controllable, locally-scoped, learned behavior of a single model, and measure it at the level of narrative structure.

PROBLEM: post-training raises quality but collapses diversity
    |
    +-- prior fix A: hotter decoding ...... quality falls with diversity
    +-- prior fix B: base+instruct at test . 2x memory, global knob only
    +-- prior fix C: diverse SFT data ...... fixed style, no control
    |
    v
ASSUMPTION: creativity is a *span-level* mode the model can be
            taught to enter and exit on its own
    |
    v
METHOD: build tagged data where creative segments are wrapped in
        [StartCreativity]; SFT the model to emit the tag itself
    |
    v
EVIDENCE: (1) matches/beats multi-model + distilled baselines on
              lexical + semantic + new structural diversity
          (2) human raters prefer it as more creative 70.3% of time
          (3) GRPO from this checkpoint: +~4% AMC, +~5pt MATH
    |
    v
CONCLUSION: a single model can hold both distributions, and the
            creative one is a better substrate for exploration/RL

The Increment

One sentence: Before, getting base-model variety out of an aligned model meant paying at inference time (two models, or a temperature that hurts quality); after, the aligned model carries a learned in-band switch that turns variety on for the spans where it helps — and that switch turns out to also make RL work better.

Core Mechanism

Three pieces. First, data construction. You need supervision that says “here is a span where base-model-like unpredictability is the right move, and here is a span where post-trained competence is the right move.” The scalable way to get this is to generate from both distributions on the same prompt, then splice: keep the structural/quality scaffolding from the post-trained model and substitute creative segments sourced from the base-model-like generations, wrapping those segments in a [StartCreativity] marker. The scaling story matters here — this is the difference between the method and a hand-curated diversity dataset. No human writes the tags.

Second, training. Plain instruction tuning on the tagged corpus. Critically, the tag is in the vocabulary and in the loss, so the model learns two coupled things at once: when to open a creative span (a planning decision) and how to write inside one (a distributional shift). At inference you sample normally from a single model; the tags appear on their own and are stripped from the final output. You can also force or suppress the tag if you want the knob explicitly.

Third, measurement. The structural diversity metric abstracts each generation into a graph — entities and events as nodes, relations between them as edges — and measures pairwise graph edit distance across a sample set. Graph edit distance is the minimum number of node/edge insertions, deletions, and relabelings to morph one graph into another. High average distance means the stories differ in shape, not just wording.

DATA BUILD
  prompt --+--> [ post-trained model ] --> fluent, safe, generic draft
           |                                        |
           +--> [ base-like model    ] --> odd, varied, uneven text
                                                    |
                          v                         v
                    [ splice + tag ]  <-------------+
                          |
                          v
        "... she opened the door. [StartCreativity] the hallway
         was inventory, not architecture ... "

TRAIN
  tagged corpus --> SFT --> single model M  (tag token in loss)

INFERENCE
  prompt --> M --> tokens ... [StartCreativity] ... tokens
                     |                |
                quality mode     creative mode
                     |                |
                     +------ strip tags -----> output

EVALUATE
  N samples --> parse to graphs --> pairwise graph edit distance
      G1: (hero)-fights->(dragon)-guards->(gold)
      G2: (clerk)-audits->(dragon)-owes->(bank)
      GED large ==> structurally diverse
      lexical/semantic metrics would also fire here, but they
      also fire on G1 vs a re-worded G1; GED does not

The metaphor: a jazz chart. A big-band arrangement is fully written out — every note fixed, reliably good, and identical every night. That is the post-trained model. A free improv session is different every night and sometimes falls apart. That is the base model. What working charts actually do is write the heads and the ensemble passages, then print a marked box that says *solo here, 16 bars. The [StartCreativity] tag is that printed box.

Now walk the method through the chart. The data-construction pipeline is the arranger deciding which bars should be written and which should be solo space — and it makes that decision by listening to both a written version and an improvised version of the same tune and stitching. Training is the band rehearsing until players internalize not just how to solo but when the solo is supposed to start — nobody has to cue them. Inference is the performance: one band, one chart, no second orchestra standing by in the wings (that is the multi-model baseline, which is like hiring a separate improv ensemble and cross-fading between stages). Cranking the temperature, by contrast, is telling the whole band to play loose all night, including during the tight ensemble hits — which is exactly why it costs quality.

And the RL result fits the metaphor too: a band that knows how to solo has more material to discover when you tell it to go find something new. A band that can only play the chart has nothing to search over.

Key Concepts

• Mode collapse from post-training: Imagine asking 100 people to name a fruit. Raw human answers scatter across apple, mango, persimmon, durian. Now train them by rewarding whichever answer a judge liked most — everyone converges on “apple.” Nothing is wrong with apple; it is just that you have destroyed the distribution. Preference optimization does this at the sequence level: it sharpens probability mass onto whatever the reward model rated highest, so a model that could write 50 story openings now writes one opening 50 ways. The pathology is not lower quality, it is lower *entropy over good options.

• In-band control tokens: Instead of controlling a model from outside (change the sampler, change the prompt, mix in another model’s logits), you put the control signal *inside the text as a special token and train on data where that token predicts a real distributional shift. The model then generates its own control signal. This is the same trick as control codes in conditional generation and reward-conditioned tokens in offline RL, applied to creativity. The elegant part is that the model chooses where to place it, so control becomes span-local instead of global.

• Graph edit distance for narrative structure: Take two Lego builds. You can compare them by counting how many bricks of each color they use (lexical diversity) or by asking people whether they feel similar (semantic diversity). Or you can count the minimum number of brick add/remove/swap operations to turn one into the other — that is edit distance, and it is sensitive to *assembly, not inventory. Applied to stories: turn each story into a graph of who did what to whom, then ask how many surgical edits map one plot onto another. Two stories about completely different characters that both run “protagonist loses mentor, then avenges them” will have small graph edit distance even though every noun differs. That is the signal the paper claims prior metrics were missing.

Framework Shift

Before (mainstream approach):          After (this paper):

 [ base model ]   [ RLHF model ]        [ one CreativeInstruct model ]
      |                 |                          |
      +--- mix logits --+                   emits its own tag
              |                                    |
              v                          +---------+---------+
        one global knob                  |                   |
        (temp / alpha / lambda)      normal span      [StartCreativity]
              |                          |                   |
              v                          v                   v
     diversity <==> quality        quality mode       creative mode
       (tradeoff dial)                    |                   |
              |                           +-------+-----------+
              v                                   v
     2 models at inference                one model, span-level control
                                                  |
                                                  v
                                          better RL starting point


 measurement before:                    measurement after:
   distinct-n, self-BLEU  (words)         + graph edit distance
   embedding cosine       (meaning)         over event graphs (shape)

One sentence: From *tuning a global dial between two models at inference to teaching one model to open and close creative spans by itself, the core shift is moving diversity control from the sampler into the model’s own token stream.

Expert Assessment

A caveat up front: I am working from the abstract and the shape of the surrounding literature, not a full read of the data-construction section — so treat my method reconstruction as a plausible reading, not a verified one, and weight the experimental critiques accordingly.

Problem choice: Real gap, and well-positioned. Diversity collapse from post-training is one of the field’s genuinely unsolved practical problems, and the framing that it hurts *implicitly creative tasks — RL exploration especially — is the most valuable sentence in the abstract. There is a growing body of evidence that RL from a low-entropy policy just re-ranks what the policy already does, so “creativity as exploration substrate” is exactly where the field’s attention is heading. The paper is riding a real wave, not manufacturing one.

Method maturity: Clever, but the cleverness is in the packaging rather than the primitive. In-band control tokens are old (control codes, reward tokens, quality-conditioned pretraining); the contribution is applying them at *span granularity to creativity and building the tagged data automatically. That is a good idea, and the automatic-splicing pipeline is where the real risk lives: splice quality determines everything, and I would want to see how brittle the “which spans should be creative” heuristic is. Simpler baselines that deserve explicit refutation: min-p or entropy-targeted sampling (much stronger than plain temperature, and free), and simply distilling base-model outputs with light quality filtering. The abstract says distilled variants were compared, which is the right instinct.

Experimental integrity: Two things I would push on at a poster. First, the paper introduces its own headline diversity metric and then wins on it. That is not disqualifying — someone has to introduce new metrics — but a self-authored metric that favors the authors’ method needs independent validation, and the human eval should be doing that work. The 70.3% creativity preference is a decent signal, but *creativity preference alone does not settle the quality question; I want the paired quality judgment at the same time, not a separate automatic quality score. Second, the RL numbers are the flashiest claim and the thinnest evidence: AMC is a tiny benchmark (tens of problems), so “+4%” can be one or two items, and single-seed GRPO comparisons are notoriously noisy. Multiple seeds with variance, or this result gets downgraded to “suggestive.”

Writing quality: Based on the abstract’s emphasis, the likely soft spot is the data-construction description — it is the load-bearing part and the easiest to compress into “we splice and tag.” The section I would most want rewritten is the RL experiment: if the authors expanded it with more benchmarks, seeds, and an entropy analysis showing *why the creative checkpoint explores better, that section alone would move the paper from “nice method paper” to “result the field cites.” Right now it reads as a promising appendix promoted to the abstract.

Verdict: weak accept — the mechanism is sound and clearly useful, but the two most exciting claims (the new structural metric and the RL transfer) each rest on evidence one notch thinner than the claim.

Takeaways

Things you can lift and use elsewhere:

Span-level control tokens as a general pattern. Any time you want a model to behave differently in *part of its output — more terse here, more rigorous there, more speculative in this paragraph — you can build tagged data by splicing generations from two different sources and letting the model learn both the placement and the behavior. This generalizes far beyond creativity: verbosity control, citation-required spans, “I am guessing here” spans, code sections that should be defensive versus fast.

Make the model emit its own control signal. The reason this beats a prompt or a sampler flag is that the *placement decision is itself learned from data. If you have been solving control problems with system prompts, ask whether the decision of when to apply the control is something you actually want the model to learn.

Edit distance on structured abstractions as a diversity metric. If your outputs have structure — API call sequences, SQL query plans, test suites, agent trajectories, molecules — abstract to a graph and use graph edit distance instead of embedding cosine. Embeddings tell you if two outputs *feel different; edit distance tells you if they are built differently. This one is immediately portable and cheap.

Check the entropy of your RL starting point. The most actionable claim in the paper: if you are running GRPO/PPO and plateauing, the problem may be that your SFT checkpoint has already collapsed the distribution you are trying to search. Before tuning the RL recipe, measure the entropy and sample diversity of the initialization.

One honest caveat to carry: the multi-model baselines this replaces are not expensive in absolute terms. If you have the memory, base+instruct contrastive decoding is a strong, training-free baseline. The case for CreativeInstruct is strongest when you need single-model deployment or when you want the creative checkpoint as an RL starting point.