Mercity/creative-writing-llm

Diversity-aware post-training for creative story generation (Qwen3-4B)

Research artifacts for a controlled study of *whether RL post-training can raise the semantic diversity of short-story generation without paying for it in quality*.

Everything here is measured, not asserted: the repo carries the LoRA adapters, the scored generation pool they were selected from, the per-step reward telemetry, the judge-calibration evidence, the experiment reports, and the source that produced all of it.

⚠️ Work in progress. Only E0 (quality-only baseline) has finished its full 300-step run. E1 was still training when this snapshot was uploaded — its checkpoints stop at step 200 of 300. E2, E3 and E4 have not been trained yet; their configs are present, their adapters are not. No cross-arm comparison exists yet, and none should be read into this repo.

The finding this study is built around

Sampling the base policy 16 times per prompt over 1,000 prompts (16,000 scored stories, outputs/pool_4b/) gives a mean effective rank of 2.006 out of a ceiling of 16. Sixteen independent samples of the same prompt span roughly two effective directions in embedding space; mean pairwise distance is 0.132, i.e. same-prompt stories sit at about 0.87 cosine similarity.

Quality-only RL does not fix this. E0 ran 300 steps of GRPO on the judge's quality score alone and moved judge quality +0.30 (6.557 → 6.860, first vs. last quarter of the run) while the diversity statistics stayed flat (mean pairwise deviation 0.1354 → 0.1391; mean group log-det −12.81 → −12.64). The held-out checkpoint study (outputs/ckpt_study/, 10 prompts × 6 samples per checkpoint) tells the same story:

Craft improves; the model still writes the same story six times. That gap is what the diversity arms exist to close.

Method

Base policy. Qwen/Qwen3-4B-Instruct-2507, LoRA r=32, alpha=64, dropout 0.0 on all attention and MLP projections (q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj).

Online arms. GRPO via TRL 1.10 with multi_objective_aggregation="normalize_then_sum", which is a line-for-line implementation of GDPO (arXiv 2601.05242): each reward channel is normalized within its prompt group before the channels are summed and the advantage is normalized batch-wise. Generation is vLLM colocated in the training process. G=8 samples per prompt, generation batch 16 (2 prompts × 8), max_completion_length=1024, lr 3e-5 with constant_with_warmup, beta (KL) 0.02, 300 steps. Everything ran on a single RTX 5090 (32 GB) — the 32 GB budget is what forced the 4B policy, since colocated GRPO needs two resident copies of the weights.

Offline arms. DPO over preference pairs mined from the same 16k pool (beta 0.1, sigmoid loss, lr 5e-6, epochs chosen for optimizer-step parity across arms).

Reward channels.

• —quality — judge score, per story.
• —deviation d_i — mean embedding distance from a sample to the rest of its group (BAAI/bge-base-en-v1.5), credited only if judge quality ≥ tau=5.
• —marginal m_i — the sample's marginal contribution to the group's log-determinant (a set-level diversity volume), z-scored within group.

Every diversity credit is per-sample by construction. A single set-level scalar shared by a whole group has zero within-group variance and therefore contributes exactly nothing to a GRPO advantage — the failure mode diagnosed in prior_run/ and written up in logs/experiments/00_prior_run_autopsy.md.

Judge. deepseek-v4-flash-0731 via OpenRouter, scoring one story at a time on an absolute rubric (not batched relative scoring), with reasoning disabled. Calibrated before any training against four cells with known ordering: detects incoherence (Δ 2.25), detects repetition (Δ 4.67), not saturated (sd 0.76). See logs/experiments/01_judge_calibration.md.

Programmatic gates. Completeness, 150–600 word window, 4-gram / line loop detection and finish_reason == "length" are checked deterministically before a story is ever judged, so truncated or degenerate text cannot earn reward through any channel.

The arms

Repository layout

outputs/
  E0-baseline/            LoRA adapters: checkpoint-{50..300} + final
                          reward_history.json     per-step reward telemetry (300 steps)
                          trl_log_history.json    full TRL log history
                          judge_cost.json         judge call/token/USD accounting
  E1-div-individual/      same structure, checkpoints 50-200 (run in progress)
  pool_4b/                pool_train.jsonl        16,000 scored base-policy stories
                          emb_train.npy           their bge-base-en-v1.5 embeddings
                          summary_train.json      pool-level summary stats
  pairs_4b/               divpo_emb_train.jsonl, divpo_prob_train.jsonl,
                          multipos_train.jsonl, pair_stats_train.json
  ckpt_study/E0-baseline/ stories.md, raw.json, metrics.csv  (per-checkpoint story dumps)
logs/
  experiments/            the written reports (start here)
  figures/                plots referenced by the reports
configs/                  one YAML per arm, with the reasoning for every knob in comments
src/                      all Python: training, rewards, diversity, judge, gates, eval, tests
data/                     train/eval prompt splits + split metadata
prior_run/                the previous, failed run's artifacts (see autopsy)

Read the reports in this order: 00_prior_run_autopsy.md → 01_judge_calibration.md → 02_setup_and_deviations.md → 03_pool_4b_baseline.md → 04_E0_checkpoint_story_study.md → E0-baseline.md.

Data schemas

outputs/pool_4b/pool_train.jsonl — one story per line: prompt_id, prompt, idx, text, n_words, n_tokens, finish_reason, gate_passed, gate_reasons, ends_cleanly, quality, novelty, deviation, marginal, group_logdet, mean_logprob, cumlogprob. Row i of emb_train.npy corresponds to line i of the JSONL.

outputs/pairs_4b/*.jsonl — DPO rows: prompt_id, prompt, chosen, rejected (+ per-row weight for the multi-positive set).

data/ — 1,000 train and 50 eval prompts (seed 42), filtered to 10–60 words from 272,600 scanned; split_meta.json also carries the exact system prompt used for every generation in the study.

Using an adapter

Each checkpoint directory is a standard PEFT adapter (adapter_model.safetensors + adapter_config.json, ~253 MB each).

from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

path = snapshot_download(
    "Mercity/creative-writing-llm",
    allow_patterns="outputs/E0-baseline/final/*",
)

tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Instruct-2507")
base = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B-Instruct-2507", dtype="auto", device_map="auto"
)
model = PeftModel.from_pretrained(base, f"{path}/outputs/E0-baseline/final")

system = (
    "You are a fiction writer. Write a complete short story of 200-500 words responding "
    "to the writing prompt.\nWrite only the story: no title, no preamble, no commentary, "
    "no author's note.\nFinish inside the word budget. The story must reach a real ending, "
    "not stop mid-scene."
)
msgs = [{"role": "system", "content": system},
        {"role": "user", "content": "A lighthouse keeper receives a letter addressed to the sea."}]
ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(model.device)
out = model.generate(ids, max_new_tokens=1024, temperature=1.0, top_p=1.0, do_sample=True)
print(tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True))

Swap outputs/E0-baseline/final for outputs/E0-baseline/checkpoint-150, outputs/E1-div-individual/checkpoint-200, etc. The study sampled at T=1.0 during training and T=0.9 / top_p=0.95 for the checkpoint story study.

Selected numbers

Base-policy pool (logs/experiments/03_pool_4b_baseline.md, n=16,000):

Quality and diversity are largely independent across prompts (r = −0.108), so a method that raises diversity without lowering quality is exploiting existing slack rather than defying a tradeoff.

E0 run (logs/experiments/E0-baseline.md, 300 steps, 4,800 stories judged, verdict healthy):

No entropy collapse, and the reward-hacking trip conditions (diversity up while quality or validity falls) were not tripped. Judge cost for the whole E0 run: $0.84 over 5,109 calls.

Qualitative (logs/experiments/04_E0_checkpoint_story_study.md): the collapse is a tonal monoculture, not a lexical one — 92–95% of stories carry solemn/elegiac vocabulary against ~15% comic, even on explicitly comic prompts, and verbatim opening duplication rises with quality-only training. This predicts that n-gram metrics (distinct-4, self-BLEU) will separate the arms far less than effective rank does.

Environment

torch 2.13.0+cu130 · vLLM 0.27.1 · transformers 5.15.0 · TRL 1.10.0 · PEFT 0.20.0 · sentence-transformers 5.7.0 (BAAI/bge-base-en-v1.5) · single RTX 5090, 32.6 GB, sm120. FlashInfer is disabled (it misdetects sm120); Liger is disabled because it suppresses TRL's entropy metric.

Caveats and honest notes

• —E1 is mid-flight. Its checkpoints are a snapshot at step 200/300 and its reports do not exist yet. E2/E3/E4 are unrun. Do not read arm comparisons out of this repo.
• —`04_E0_checkpoint_story_study.md` documents an earlier, partial 163-step E0 run (at the pre-correction lr 3e-6). outputs/ckpt_study/E0-baseline/metrics.csv is the regenerated version covering all checkpoints of the final 300-step run; where they disagree, trust the CSV.
• —Learning rate deviates from the plan (3e-5, not 3e-6): at 3e-6 the adapter was effectively frozen — KL pinned near 8e-4 for 171 steps with every metric inside its noise band. The evidence is recorded in the config comments and 02_setup_and_deviations.md.
• —Optimizer/scheduler/RNG state is not included for E1's checkpoints. It is resume-only state, is rotated away by save_total_limit, and none of it is needed to load or evaluate an adapter.
• —`prior_run/train.py` is a redacted copy — the original contained a hard-coded OpenRouter API key, replaced here with a placeholder. Nothing else in the file was changed.
• —Judge is not perfect. It fails the truncation-sensitivity check (A−B = 1.33 against a 2.0 bar), which is tolerated only because truncation is caught deterministically by the gates before any story reaches the judge. Pool-scoring judge failure rate was 0.45%.
• —Every deviation from the original plan is enumerated in logs/experiments/02_setup_and_deviations.md rather than discovered in a footnote.

References

• —GDPO — group-wise reward normalization for multi-objective RL, arXiv 2601.05242 (Liu et al., NVIDIA)
• —DivPO — Lanchantin et al., 2025 (diverse preference optimization)
• —Effective rank — Roy & Vetterli, the exp-entropy of the Gram spectrum, used here as the primary continuous mode-count metric (k-means + silhouette was measured to be unable to separate collapse from spread and was demoted)

License

Apache-2.0, matching the Qwen/Qwen3-4B-Instruct-2507 base model. The prompt splits derive from a public writing-prompts corpus; stories in outputs/pool_4b/ are model generations.

Status (updated — E0 and E1 both complete, 300 steps each)

Held-out evaluation (30 prompts × 16 samples = 480 stories/model)

E1 wins on both axes: vs base it gains 2.9× E0's effective-rank improvement and 1.6× E0's quality improvement. Not a diversity-for-quality trade.

Key methodological finding: embedding metrics separate the arms by 2.9×; n-gram metrics (distinct-4, self-BLEU) do not separate them at all — distinct-4 actually rates E0 higher. The collapse is tonal/structural, not lexical.

Key finding on LLM judges: across 15,870 scored stories, corr(judge "novelty", a story's actual embedding deviation) = +0.057, while corr(judge "novelty", judge quality) = +0.799. An LLM judge's novelty score is ~80% a restatement of quality and carries almost no information about semantic distinctiveness.

See `REPORT.pdf` for the full write-up (metrics, figures, story examples, qualitative read) and `NOTES.md` for the complete engineering log including every bug and every retracted claim.