| --- |
| 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. |
|
|
