CreativeInstruct: Scalably Teaching LLMs to Balance Quality, Creativity, and Diversity

Problem§

Post-training (e.g., instruction tuning, RLHF) improves LLM quality but suppresses output diversity, harming both explicitly creative tasks (story generation) and implicitly creative ones (RL exploration). CreativeInstruct aims to teach a single model to dynamically balance base-model creativity with post-trained quality via explicit control tokens.

Method: CreativeInstruct§

CreativeInstruct constructs supervision data by combining outputs from a base model (creative, diverse) and a post-trained model (high-quality, less diverse). For each instruction $x$, a target sequence is formed:

$$ y = y_{post} \oplus [\text{StartCreativity}] \oplus y_{creative} \oplus [\text{EndCreativity}], $$

where $y_{post}$ is a high-quality prefix and $y_{creative}$ is a creative suffix sampled from the base model. The special token [StartCreativity] is inserted just before the creative portion. The model is trained with standard causal language modeling loss:

$$ \mathcal{L} = -\sum_{t} \log p_\theta(y_t \mid y_{<t}, x). $$

During inference, the model learns to decide when to emit [StartCreativity] and shift to a more diverse generation mode, effectively learning a learnable "creativity switch." This requires only one model at inference time, unlike mixture-of-experts or multi-model voting baselines.

Structural Diversity Metric§

To evaluate narrative-level diversity beyond lexical/embedding similarity, the paper proposes representing each generation as a directed graph $G=(V,E)$ where nodes are events/entities and edges are relations (e.g., causal, temporal). Diversity between two stories $G_1$ and $G_2$ is quantified via graph edit distance (GED):

$$ \text{GED}(G_1, G_2) = \min_{\{e_i\}} \sum_i c(e_i), $$

where $e_i$ are edit operations (node/edge insertion, deletion, substitution) and $c(e_i)$ is the edit cost. Normalized GED provides a robust structural diversity score that captures story-level variation.

RL Applications§

The authors show that CreativeInstruct checkpoints serve as better substrates for downstream RL. Applying GRPO on a CreativeInstruct-tuned model improves AMC accuracy by ~4% and MATH by ~5 points over applying the same RL to the post-trained checkpoint. The increased genrative diversity provides richer exploration during RL training.

Implementation Sketch§

The following PyTorch snippet illustrates the core supervised fine-tuning loop on constructed CreativeInstruct data:

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("llama-2-7b")
tok = AutoTokenizer.from_pretrained("llama-2-7b")

# Add special tokens to vocabulary
creative_tokens = ["[StartCreativity]", "[EndCreativity]"]
tok.add_tokens(creative_tokens)
model.resize_token_embeddings(len(tok))

def train_step(batch):
    # batch[instruction], batch[target] where target includes creative spans
    texts = [inst + " " + tgt for inst, tgt in zip(batch["instruction"], batch["target"])]
    enc = tok(texts, return_tensors="pt", padding=True, truncation=True)
    labels = enc.input_ids.masked_fill(enc.attention_mask == 0, -100)
    outputs = model(enc.input_ids, attention_mask=enc.attention_mask, labels=labels)
    loss = outputs.loss
    loss.backward()
    return loss

The key is that the model is trained to generate the [StartCreativity] token at the appropriate point, enabling controllable creativity during decoding without auxiliary models or complex pipelines.