abhishek.it
Back to writing
·
#llm#machine-learning#diffusion#mamba#state-space-models#training-from-scratch#ai-research#code-generation

Blackwood Diffusion: Language Models That Revise

A technical note on Blackwood Diffusion: a masked-denoising language model with a 49,152-token agentic tokenizer, staged context growth, and revision-native generation.

Blackwood Diffusion is a training track for a language model that does not generate only left to right.

The core idea is masked denoising for language: corrupt a sequence, train the model to reconstruct it, then use the same ability at inference time to revise a draft over several passes.

Most useful code work is not pure continuation. It is patching, filling, deleting, moving, and repairing. A test is broken. A function has a missing branch. A refactor touched three files and left one import wrong. A tool result invalidated an earlier assumption.

Autoregressive models can do all of that, and they are very strong. But they still emit the answer one token at a time. Diffusion-style language models open a different path: update many positions, revise uncertain parts, and use several denoising passes instead of one long left-to-right decode.

If a model can revise many tokens at once, the unit of generation stops being "the next word" and becomes "the next better draft."

This is not replacing the original Blackwood autoregressive LLM track. It is a parallel pretraining track beside it.

The normal LLM path stays focused on the general-purpose base: broad language ability, instruction following, long-form reasoning, and standard chat/code generation. Blackwood Diffusion focuses on a sharper question: agentic tool workflows that revise structured drafts instead of only streaming tokens.

That matters for business use cases where the output is rarely just prose. Finance, accounting, operations, procurement, compliance, internal analytics, and similar domains are full of structured state changes:

  • read records
  • call a tool
  • reconcile fields
  • update a draft
  • check constraints
  • repair one wrong span
  • produce an auditable final answer

Those workflows look more like iterative state repair than creative writing. That is why diffusion-style generation is interesting here.

Business angle
Most business agents are revision systems.
They read structured data, call tools, compare fields, repair mistakes, and produce an answer that needs to be auditable. That is closer to denoising a changing workspace than writing a story from left to right.
Blackwood Diffusion training loop learn to reconstruct corrupted text, then use the same loop for iterative generation clean tokens code, text, tools 49,152 vocab mask schedule corrupted draft visible tokens + <mask> positions hybrid backbone Mamba-style mixing periodic attention token logits everywhere reconstruct loss on masked and sampled spans inference loop keep confident tokens remask uncertain spans not just completion revision as the primitive product target code repair + agents

The Training Shape

The training objective is discrete masked denoising.

For a clean token sequence (x_0), sample a mask ratio, corrupt positions into <mask>, and train the model to predict the original tokens:

L(θ)=Ex0,m[i:mi=1logpθ(x0,imask(x0,m),m)]\mathcal{L}(\theta)= \mathbb{E}_{x_0,m} \left[ -\sum_{i:m_i=1} \log p_\theta(x_{0,i}\mid \text{mask}(x_0,m), m) \right]

That looks close to BERT at first glance, but the target is different. BERT-style training is usually for representation learning. Here the target is generation: start with a blank or partial draft, fill likely tokens, keep high-confidence regions, remask uncertain regions, and repeat.

The promising backbone so far is a Mamba-heavy hybrid, not a pure Transformer and not pure Mamba.

The current candidate is a K5-style hybrid: Mamba-style sequence layers do most of the token mixing, and attention appears periodically to give the network global routing. The diffusion objective sits on top of that.

The reason for this combination is practical. Pure attention is powerful, but expensive as context grows. Pure state-space models are efficient, but language still has global dependencies. A hybrid gives me a place to test the middle: enough global routing to write real text, enough efficient sequence processing to make the generation path interesting.

That is the part that feels most promising: diffusion changes the generation algorithm, while Mamba changes the cost shape of sequence processing. If both hold, the result is not just "a different loss." It is a different inference profile.

Mamba-heavy blocks      -> cheap sequence mixing
periodic attention     -> global routing where it matters
masked denoising       -> parallel token repair
confidence remasking   -> revise only uncertain spans

The goal is to keep quality while reducing the number of sequential decisions needed at inference time.

In a normal decoder, the model gets one chance at each position as it moves forward. In a diffusion model, the answer can be revisited. That is the interesting part. The model can put down a rough structure first, then use later passes to clean up names, syntax, imports, tool arguments, and long-range consistency.

That maps naturally to code.

def normalize_user(record):
    name = <mask>
    email = <mask>
    if <mask>:
        return None
    return {
        "name": name,
        "email": email,
    }

This is not a "what comes next?" problem. The missing pieces depend on the whole function. A good editor should look at all the holes together, infer the intended shape, and revise them as a group.

That is the behavior this run is built to test.

The sampler is confidence-aware:

1. start with prompt + masked draft region
2. predict all masked positions in parallel
3. keep high-confidence tokens
4. remask low-confidence or verifier-marked spans
5. repeat for a small fixed number of passes

The important detail is step 4. Revising only uncertain parts makes generation feel closer to editing than streaming.

Why Businesses Should Care

The business use case is where this becomes more than a modeling curiosity.

Most enterprise workflows are not open-ended chat. They are structured repair loops. The model receives partial state, calls tools, observes results, and updates only the parts that changed.

In finance, that might mean reading a portfolio summary, checking current prices, updating risk commentary, and correcting one stale number without rewriting the whole report.

In accounting, it might mean reconciling invoice fields against a purchase order, marking uncertain rows, calling an ERP tool, then repairing only the mismatched fields.

In operations, it might mean triaging a support ticket, pulling customer metadata, checking policy, and revising the response after one tool result changes the decision.

In compliance, it might mean reviewing a draft, marking policy-sensitive spans, retrieving supporting evidence, and rewriting only the risky parts.

Those are revision-native tasks:

business state
  -> tool call
  -> observation
  -> mark uncertain fields
  -> repair only those fields
  -> verifier / audit trail

A general chat model can serve these workflows. A model trained around structured denoising is aimed at something narrower: fast, repeatable, auditable repair.

The practical product starts as a vertical agent that is excellent at one workflow: finance commentary updates, accounting reconciliation, compliance review, CRM hygiene, procurement checks, or internal analytics summaries.

Public Surface

The tokenizer is a 49,152-token BPE built for this run. It has atomic structure tokens for masking, thinking boundaries, tool calls, tool results, function names, function arguments, search, memory, file operations, planning, reflection, and termination.

That does not make a model good at tools by itself. It gives the model clean structural handles, so supervised training and reinforcement learning are not fighting fragile string boundaries.

Public links:

Papers Behind The Bet

This is not coming out of nowhere. The useful thread is not "image diffusion, but for words." It is a line of work around discrete and masked diffusion for language.

Diffusion-LM was one of the early papers that made this direction feel plausible. It showed that intermediate denoising states can be useful for controllable text generation.

DiffuSeq matters because a lot of useful work is really sequence-to-sequence: take an input state, produce a repaired or rewritten output. That is true for code, but also for accounting rows, compliance notes, and operational reports.

SEDD and MDLM are closer to the token-level problem. They treat language diffusion as discrete modeling, not just a direct copy of continuous image diffusion.

LLaDA, Dream, and Dream-Coder are the larger proof points. LLaDA showed the pretraining/SFT path. Dream pushed open diffusion LLM quality further. Dream-Coder is especially relevant for code, because arbitrary-order generation, infilling, and verifier-driven repair are exactly the shape of the work.

Mercury and Mercury 2 are the commercial proof points. Inception Labs showed that diffusion LLMs can be productized for code and reasoning, with parallel refinement instead of token-by-token decoding. Mercury 2 is especially relevant for the "thinking" direction: reasoning becomes a loop of refining a whole draft over a small number of steps, not just streaming a hidden chain left to right.

DDIM still matters, but only as sampler intuition. The text mechanics are different, but the question is the same: how few denoising passes can work before quality breaks?

DDPD is the most directly relevant paper for planned denoising. It separates generation into a planner and a denoiser. The planner decides which positions should be denoised next.

That maps cleanly to code and business workflows. Some spans are easy: punctuation, field names, repeated identifiers. Some spans need more evidence: conditions, API arguments, invoice amounts, policy-dependent branches. A planner/denoiser split is one way to make the sampler less blind.

DDIP is only a loose analogy. The useful pattern is simple: diffusion can act as a repair prior under constraints. For Blackwood, the constraints are code context, tool observations, tests, type errors, accounting rows, policy rules, and the existing draft.

Evaluation Plan

Loss is only the first signal.

For this kind of model, the useful questions are more specific:

  • Can it repair lightly corrupted text?
  • Can it fill larger missing regions without losing structure?
  • Does it become more confident in the right tokens over denoising steps?
  • Do samples improve with more passes, or do they drift?
  • How many passes are needed before the output is useful?
  • Is the final latency actually better for editing-style workloads?

Those are the measurements that matter. Token recovery is useful, but the real test is whether iterative revision produces useful code, tool outputs, and long-form text under a practical latency budget.

Staging Plan

The run is staged and branched. One trunk learns the core denoising distribution. Separate branches test generation behavior, longer context, and tool/code specialization without mixing every risk into one run.

Staged training and branching keep the base objective stable, then fork targeted branches for generation, context, and agents base denoiser short context stable CE + recovery strong trunk mixed masks keep broad behavior generation branch higher mask ratios + sampler sweeps test NFE, remasking, drift context branch 4k -> 8k -> 16k -> 32k length curriculum + replay agent/code branch tool format + verified edits preserve general ability only merge lessons that survive evals

The context branch starts small and grows deliberately. A small-context denoiser proves the objective, but a useful product model needs long files, long conversations, and multi-tool traces. The plan is staged context extension: train the base behavior first, then extend through 4k, 8k, 16k, and 32k with mixed-length batches and replay from shorter contexts.

The model has to keep short-context quality while learning long-context behavior. The context branch has to keep both.

The generation branch is separate because denoising loss and open-ended generation are not the same thing. A model can recover masked tokens well and still need sampler work before its free-form generations are good. That branch is where higher mask ratios, confidence remasking, step-count sweeps, and verifier-guided repair belong.

Why This Could Be Fast

The exciting part is not just "diffusion for text." That phrase by itself does not mean much.

The exciting part is the shape of the work.

Autoregressive decoding has a long dependency chain. If the model needs to emit a long answer, every token depends on the previous token being sampled first. You can optimize the kernel, cache keys and values, batch requests, quantize weights, and do speculative decoding, but the core loop is still sequential.

A denoising model has a different loop. One pass can update many token positions. If the model needs a small number of good passes rather than one pass per output token, there is a path to lower latency for editing-heavy tasks.

The sampler has to be good, the model has to avoid drifting, and the number of passes has to stay low. If those pieces work, the speed story becomes very different from a normal decoder.

The product version would feel less like watching a model type and more like watching it sharpen a draft.

pass 1: fill the obvious structure
pass 2: fix names and local syntax
pass 3: repair long-range consistency
pass 4: only touch low-confidence spans

For code, the target workflow is simple: touch less, revise better, stop when the verifier is happy.

Why Code Is The First Target

Code gives better feedback than normal prose.

You can run tests. You can typecheck. You can execute a script. You can compare a patch against a verifier. That makes code a good place to train and evaluate a model whose main job is revision.

The product shape is not "another chat box." It is closer to an editing engine:

  • fill missing code spans
  • repair broken tests
  • update several related snippets at once
  • revise a tool-call plan after observing a result
  • regenerate only the uncertain parts of an answer

The interesting part is not that the model is diffusion-flavored. The interesting part is that the primitive changes from "append the next token" to "improve the current draft."

That is how software actually gets written: start rough, run feedback, revise the parts that are wrong, repeat.

The Agent Angle

Agents are not just chat models with tools attached. They are loops:

plan -> act -> observe -> revise -> act again

A left-to-right model can run that loop. But the model's output is still produced as a stream. If the observation invalidates something earlier in the answer, the model usually has to start over or patch around it.

A revision-native model gives the runtime a cleaner interface: mark uncertain spans, tool-dependent spans, or test-error spans, then repair only those regions while preserving the rest.

That could make tool use less brittle:

  • keep the plan stable while revising one tool argument
  • repair a malformed function call without regenerating the whole response
  • update the final answer after a tool result arrives
  • preserve correct code while changing only the failing branch

Atomic tool tokens matter because the model needs clean boundaries for actions, observations, and revision targets. The architecture is one part of the story. The runtime protocol matters just as much.

What Would Make This Real

The bar is concrete.

The model has to denoise normal text and code reliably. The sampler has to produce coherent generations, not just good token-recovery numbers. It has to beat a strong left-to-right baseline on at least one practical editing workload. And the latency has to be good enough that the user feels the difference.

The first winning use case can be narrow:

  • masked code completion
  • patch repair after failing tests
  • multi-span refactors
  • structured tool-call repair
  • draft improvement for long answers

That is already useful.

The first small-model results are the next update: denoising, generation, code infill, sampler steps, and latency. If the training/eval schedule holds, that should land within the next two weeks.

Current Status

This is still research, but the bar is straightforward: measure sampler quality, latency, code infill, and tool-format behavior together. Then compare against strong left-to-right baselines on practical editing workloads.

The public thesis is narrower:

Language models for code should not only be trained to continue. They should also be trained to revise.

That is what Blackwood Diffusion is testing.