---
license: apache-2.0
language:
  - en
base_model: GSAI-ML/LLaDA-8B-Base
tags:
  - diffusion-language-model
  - probing
  - fact-updating
  - wikipedia
library_name: pytorch
datasets:
  - gamemaker1/wiki-wrong
---

# DiffusEdit classifiers

DiffusEdit updates an article from evidence by editing it in place. A frozen
masked diffusion language model,
[LLaDA-8B-Base](https://huggingface.co/GSAI-ML/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](https://huggingface.co/datasets/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](https://github.com/gamemaker1/diffusedit)
under `code/pipeline`. The fine-tuned editor that the system is compared
against is at [gamemaker1/t5edit](https://huggingface.co/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](https://huggingface.co/sentence-transformers/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](https://huggingface.co/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.