DiffusEdit classifiers

DiffusEdit updates an article from evidence by editing it in place. A frozen masked diffusion language model, LLaDA-8B-Base, reads the evidence and the article once, with no token masked, and the small classifiers in this repository read its hidden states. One classifier labels each sentence as kept, edited or dropped. A second one scores each token of an edited sentence by how likely it is to be out of date. A third one scores pairs of sentences and says whether one sentence has to change because another one did. The diffusion model then fills in the tokens that were scored as stale, and a reranker chooses among the fills. The model itself is never trained. The classifiers hold 12.9M parameters between them.

Every classifier was fit on the train split of gamemaker1/wiki-wrong, tuned on its validation split and scored once on its test split. The code that trains them and runs the full system is at github.com/gamemaker1/diffusedit under code/pipeline. The fine-tuned editor that the system is compared against is at gamemaker1/t5edit.

Layout

sentence/    the sentence classifier, KEEP / EDIT / DROP
staleness/   the token classifier, is this token out of date
cascade/     the pair classifier, must sentence k change because sentence j did
reranker/    the gradient-boosted tree that ranks the candidate fills

Each directory holds a classifier.pt file and the metrics.json that was written when it was trained, with its validation and test scores. A token classifier directory also holds sweep.json, the grid of mask thresholds it was tuned over, and the attention classifier holds eval.json, its scores under the evidence controls described below.

Input

LLaDA-8B-Base at revision 0f2787f2d87eac5eed8a087d5ecd24277e6255b2 runs in bfloat16 over the sequence [BOS] evidence \n\n article, where the evidence passages are written one per line as title: text and the article follows after a blank line. The forward pass is stopped after the deepest block that a classifier needs, and the hidden state of every token is kept at blocks 24 and 32 of the 32 blocks. Each state has 4,096 values. The sentence classifiers read the mean of the token states of a sentence. The token classifiers read the state of each token of an edited sentence, and the attention classifier also reads the states of every evidence token.

The classifiers the system uses

Directory	What it reads	Architecture	Parameters
sentence/linear-block24	the mean token state of the sentence at block 24	a linear layer with three outputs	12.3K
staleness/attention-block32	the token state at block 32 and the states of all evidence tokens	four attention heads over the evidence tokens plus one null slot, then a two-layer MLP	4.5M
cascade/consequence	the mean states at block 32 of the repaired sentence j and the candidate sentence k	an MLP with one hidden layer of 512 units on the concatenation of the two states, their product and their absolute difference	8.4M
reranker/rerank.pkl	the logged scores and text features of every candidate fill	a scikit-learn HistGradientBoostingClassifier with 200 iterations and 15 leaves per tree

Block 24 was chosen for the sentence classifier because its validation macro-F1 is higher than that of block 32. The thresholds are stored inside the checkpoints: the drop gate tau_d of the sentence classifier, the mask threshold tau of the token classifiers and the cascade threshold tau_c of the pair classifier. Each was chosen on the validation split.

The evidence controls test whether a classifier reads the evidence or only the sentence. The sentence classifier is trained and scored a second time on hidden states of the article alone. The attention token classifier is scored a second time with its evidence slots emptied, and a third time with each article given the evidence of another article. The scores of the controls are in the metrics files, and all of them fall well below the scores with the right evidence.

The other classifiers

Directory	What it is
sentence/linear-block32	the sentence classifier on block 32, the final block
sentence/evidence-mean-block24, sentence/evidence-nearest-block24, sentence/evidence-attention-block24, and the same at block 32	the sentence classifier with a summary of the evidence concatenated to the sentence state. The summary is the mean of the evidence passage states, the state of the nearest passage, or a learned attention over the passages.
sentence/minilm-baseline	the same linear classifier over embeddings of the sentence and the evidence from all-MiniLM-L6-v2, with no diffusion model
staleness/linear-block32	a linear classifier on the token state alone, the form of the staleness head of Detect, Remask, Repair

The reranker uses two scores that the system logs for every candidate fill. One is the mean log-probability that LLaDA-8B-Base gives each token of the candidate sentence when that token alone is masked. The other is the probability that the fact verifier tals/albert-base-vitaminc-mnli gives to the best evidence passage supporting the candidate sentence.

Checkpoint format

classifier.pt is a dictionary saved with torch.save. A linear classifier holds state_dict, the per-feature mean and std that normalise the hidden states, the layer it reads and its tuned threshold. The attention classifier also holds ev_mean and ev_std for the evidence states, a config dictionary with d, dim, heads, hidden and dropout, and tau_match, the threshold at which the share of masked tokens equals the share of stale tokens. The pair classifier holds state_dict, d, hidden, layer, seed and the paths it was trained from. The reranker is a pickle of {"model", "features", "run"} written with scikit-learn 1.9.1, and it loads only with a compatible scikit-learn.

In the code, classifier.train defines the sentence classifier, staleness.train and staleness.attention the token classifiers, cascading.train the pair classifier and infilling.rerank the reranker. staleness.attention.Scorer and infilling.run show how the classifiers are applied to the hidden states of a new article.

Licence

Apache 2.0. LLaDA-8B-Base keeps its own licence.