CytoDiffusion: A Generative Classifier for Blood Cell Morphology

Two ways to teach a machine to classify

Almost every image classifier you have heard of works the same way underneath. It is discriminative, meaning it learns to draw boundaries between categories. Show it enough pictures of cats and dogs and it learns the dividing line, the features that separate one from the other, and it puts any new picture on one side or the other. This works beautifully when your categories are clear and your examples are plentiful.

Blood is neither. A blood smear contains many cell types that shade into one another, some abnormal cells are vanishingly rare, and the important ones, the blasts that signal leukaemia or a bone marrow disorder, are precisely the ones a model sees least during training. A discriminative model, trained mostly on normal cells, learns a boundary tuned to the common cases and can be caught flat footed by the rare, dangerous ones that fall outside everything it has learned. The site has covered plenty of the discriminative approach to cells, from multi headed models for nuclei detection to graph based epithelial cell classification, and both are powerful within their trained world.

CytoDiffusion is generative instead, and that word is the whole story. Rather than learning where the line between cell types sits, it learns what each cell type looks like, the full range of normal appearances for a neutrophil, a lymphocyte, a blast, and so on. Given a new cell, it does not ask which side of a boundary the image falls on. It asks which cell type it could best have generated this image as, and answers with that type. Learning the whole picture rather than just the dividing line is what gives the method its unusual mix of strengths.

Classifying by generating

The engine underneath is a diffusion model, the same family that powers modern image generators. A diffusion model learns to turn random noise into a realistic image step by step, and in doing so it learns the distribution of what realistic images look like. The clever twist here is using that generative skill backward, to classify.

The logic rests on a piece of probability every statistics student meets, Bayes’ theorem. If you can estimate how likely an image is under each possible class, the best guess is simply the class that makes the image most likely. Assuming no class is favoured in advance, classification becomes a search for the cell type that best explains the picture.

The trouble is that a diffusion model does not hand you that probability directly. What it gives you is a related quantity, how well it can predict the noise that was added to an image when it is told to treat that image as a particular class. If the model is told a lymphocyte image is a lymphocyte, it predicts the added noise accurately. Told the same image is a neutrophil, its prediction is worse. So the class the model reconstructs with the smallest error is the class it believes the image belongs to.

There is a practical catch. Estimating that error well means adding noise many times over and averaging, and doing it for every candidate class for every image is slow. The team borrowed a trick from the mathematics of decision making, a successive elimination scheme that behaves like a smart tournament. It samples a few noisy versions, throws out the classes that are clearly losing, and spends its remaining effort only on the classes still in contention, stopping early once a winner is clear. On average it settled each cell in about forty two rounds, roughly one and a half seconds, focusing computation where the decision is genuinely hard.

What learning the whole picture buys you

The point of all this machinery is not a tiny bump on a leaderboard. On plain accuracy, CytoDiffusion matched or narrowly beat the best conventional models across several datasets, which is table stakes. The real payoff shows up on the harder properties that matter in a clinic, and this is where the generative approach pulls ahead.

Take anomaly detection first, since it is the most clinically loaded. Because CytoDiffusion has learned what normal cells look like, a cell it cannot recreate well is, by definition, unusual, and that makes spotting rare abnormal cells a natural consequence of the design rather than a separate feature to bolt on. On one dataset it caught abnormal blast cells with a sensitivity of about ninety percent while keeping specificity high, where a strong vision transformer managed a sensitivity of only twenty eight percent, low enough to be dangerous in a screening setting where a missed blast is a missed leukaemia. Detecting the abnormal is exactly what a discriminative model, trained to sort the normal, is not built to do.

The domain shift result matters just as much, because it is the quiet reason so many medical AI tools fail in the real world. A model trained on images from one hospital’s microscope and staining protocol often stumbles on another’s, and hospitals differ constantly. CytoDiffusion held its accuracy far better across these shifts than the conventional models, which suggests it learned something about the underlying biology of a cell rather than the surface quirks of one lab’s imaging. And in the low data regime, where a rare disorder might leave you only a handful of example cells, it stayed accurate where discriminative models faltered, because learning to generate a class needs less data than learning to sharply separate it from every other. The same appetite for squeezing signal from limited medical images runs through our look at how AI is moving into early diagnosis.

A machine that knows what it does not know

Two further results stand out because they touch trust rather than raw accuracy. The first is uncertainty. A good clinical tool should not just be right often, it should know when it is unsure, so that hard cases can be flagged for a human rather than confidently mislabelled. The team measured this with a technique from psychophysics, the science of how well an observer’s confidence tracks the actual difficulty of what they are judging, and found that CytoDiffusion’s sense of its own uncertainty was remarkably well calibrated, close to the theoretical ideal.

Strikingly, when they ran the same analysis on ten expert haematologists, the model’s self knowledge came out better than the humans’. That does not mean it is a better haematologist, but it does mean its confidence is a more reliable signal, which is exactly what you want for triage. Cases it is sure about can be cleared automatically, and cases it flags as uncertain can be sent to a person, and the flag can be trusted.

The second trust result is the authenticity test, and it is the kind of finding that makes you pause. The team had CytoDiffusion generate synthetic blood cell images, then asked ten experienced haematologists to tell the real cells from the fake ones. The experts scored about fifty two percent, barely better than a coin flip. The synthetic cells were, for practical purposes, indistinguishable from real ones, which is strong evidence that the model has genuinely learned the distribution of blood cell appearances rather than memorizing a shortcut.

CytoDiffusion generates synthetic blood cell images that expert haematologists cannot distinguish from real ones, demonstrating good command of the underlying distribution. Deltadahl and colleagues, Nature Machine Intelligence, 2025

Explaining itself without an afterthought

Interpretability is often bolted onto a medical model as an afterthought, a separate technique that guesses at what the model was looking at. CytoDiffusion’s explanations fall out of the generative process itself. Because it can imagine a cell as any type, it can answer a genuinely useful question, if this cell were a neutrophil instead of an eosinophil, what would have to change?

The answer is a counterfactual heat map, an image that highlights exactly the regions that would need to differ for the cell to belong to another class. When the team looked at these maps, they lined up with real biology, focusing on the granularity of the cytoplasm and the shape of the nucleus, the very features a haematologist uses. Because the explanation comes from the same machinery that does the classifying, it reflects the model’s actual reasoning rather than a separate approximation of it.

Clinical translation gap

This is where enthusiasm has to yield to care, because a strong benchmark and a bedside tool are separated by a great deal, and a few things about this study mark that distance plainly.

The results are retrospective. They come from evaluating the model on curated collections of images, some public, one built by the team, rather than from a prospective trial in a working laboratory where the model reads fresh slides and its decisions are checked against what actually happens to patients. Developing a model this way is correct, and it is not the same as proving it safe and useful in routine care, which the authors do not claim. The synthetic images, impressive as they are, also raise their own careful questions about how such generated data should and should not be used, since a tool that can fabricate convincing blood cells could mislead as easily as it could augment.

The stakes sit on the hardest cases. A blood cell classifier earns its keep on exactly the rare, ambiguous, abnormal cells that a benchmark, however large, samples thinly, and a model that shines on average can still err on the one blast that matters. CytoDiffusion’s strength on anomaly detection is genuinely encouraging here, but encouraging on a research dataset is not the same as validated across the messy diversity of real patients, laboratories, and diseases. Turning this into a clinical tool would require prospective validation, regulatory clearance, and careful study of how it fails, not just how it succeeds.

One more thing belongs in an honest read. Several authors disclose commercial interests in healthcare AI companies, including a co founder of a generative AI firm and connections to a medical imaging company. That does not undermine peer reviewed work whose code and data are released, but a diagnostic claim from a team with a commercial stake is precisely the kind that independent validation exists to check. The paper offers a foundation and an evaluation framework, not a cleared product.

Where the method falls short

Beyond the clinical gap, the authors are candid about technical limits.

The biggest is cost. Classifying by generation is expensive, since the model has to run its noise prediction many times for each candidate class, and that cost grows with the number of classes. The successive elimination trick helps a great deal, and the authors argue the expense is more tolerable in medicine, where datasets have fewer classes than a giant general image benchmark, but the method is far heavier than a single forward pass through a conventional network. Pushing it to very fine grained cell subtypes, where the number of classes multiplies, would test that tolerance.

The evaluation, though unusually broad, is still bounded. It covers particular datasets, particular cell types, and a specific set of comparison models, and the authors note that a truly granular division of blood cells into many subtypes, which is where a lot of clinical subtlety lives, remains future work. The model was also built on a general image generator adapted to this task, and how much of its strength comes from that foundation versus the medical fine tuning is not fully separated.

There is also a quieter caveat about what the model has and has not seen. Its power to flag anomalies rests on having learned a faithful picture of normal cells, and if the normal cells in a new setting differ from those it trained on, in a population or a staining style it never encountered, that safety net could fray exactly where it is needed. The authors point to conditioning on patient characteristics and augmenting rare classes as ways to make the approach fairer and more robust, but these are directions, not delivered features.

Why it matters

Step back and the contribution is a shift in how we might build medical image classifiers. The dominant recipe, learn a boundary between classes, is efficient and often accurate, and it is also brittle in the exact ways clinical practice cannot afford, blind to the rare, fooled by a new microscope, overconfident when it should hesitate. By learning the full distribution of what cells look like, CytoDiffusion turns anomaly detection, uncertainty, robustness, and explanation from separate add ons into natural properties of one model. The same generative philosophy is spreading across medical imaging, as our coverage of a graph diffusion approach to image classification shows from a neighbouring angle.

Just as valuable is the evaluation framework the paper insists on. Rather than reporting accuracy and stopping, the authors judge a model across the dimensions that actually decide whether it is safe to use, robustness, anomaly detection, calibrated uncertainty, data efficiency, and interpretability. That multidimensional bar is a better standard for medical AI than the single number the field too often chases, and adopting it would make a lot of published results more honest about their readiness. As diffusion models move from making art to reading blood, the ones that stay interpretable, well calibrated, and honest about the clinical gap are the ones worth building on. Our broader coverage of AI in pathology and cancer traces the same current.

A reference implementation

The code below is a compact, runnable version of the paper’s core idea, classification by diffusion. It defines a small conditional denoiser, trains it briefly to predict the noise added to class conditioned samples, then classifies a new sample by asking which class it can best reconstruct, using the successive elimination scheme to stop early. A full system would use a latent diffusion model on real images, which this sketch leaves out for clarity. The authors full code is linked under the block.

# CytoDiffusion style classify by generation, compact reference
# A conditional denoiser predicts added noise per candidate class, and the class
# with the lowest reconstruction error wins, with early elimination of losers.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)


class CondDenoiser(nn.Module):
    # predicts the noise added to x at timestep t, conditioned on a class label
    def __init__(self, dim, n_classes, hidden=128):
        super().__init__()
        self.cls_embed = nn.Embedding(n_classes, dim)
        self.net = nn.Sequential(
            nn.Linear(dim * 2 + 1, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, dim),
        )

    def forward(self, x_t, t, c):
        cond = self.cls_embed(c)
        tt = t.view(-1, 1).float()
        return self.net(torch.cat([x_t, cond, tt], dim=-1))


def add_noise(x0, t, n_steps=100):
    # simple variance preserving corruption, alpha_bar decreasing with t
    alpha_bar = torch.cos((t.float() / n_steps) * 1.5).clamp(min=0.05) ** 2
    ab = alpha_bar.view(-1, 1)
    noise = torch.randn_like(x0)
    return ab.sqrt() * x0 + (1 - ab).sqrt() * noise, noise


def train(model, data, labels, n_classes, steps=1500, lr=1e-3):
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    for _ in range(steps):
        idx = torch.randint(0, len(data), (64,))
        x0, c = data[idx], labels[idx]
        t = torch.randint(1, 100, (64,))
        x_t, noise = add_noise(x0, t)
        pred = model(x_t, t, c)
        loss = F.mse_loss(pred, noise)
        opt.zero_grad()
        loss.backward()
        opt.step()
    return model


def class_error(model, x0, c, n_draws=16):
    # average reconstruction error for a candidate class over sampled noise
    errs = []
    for _ in range(n_draws):
        t = torch.randint(1, 100, (1,))
        x_t, noise = add_noise(x0.unsqueeze(0), t)
        pred = model(x_t, t, torch.tensor([c]))
        errs.append(F.mse_loss(pred, noise).item())
    return sum(errs) / len(errs)


def classify(model, x0, n_classes, rounds=8):
    # successive elimination, drop the worst class each round until one remains
    candidates = list(range(n_classes))
    while len(candidates) > 1 and rounds > 0:
        scores = {c: class_error(model, x0, c) for c in candidates}
        worst = max(scores, key=scores.get)
        candidates.remove(worst)                       # eliminate the losing class
        rounds -= 1
    return candidates[0]


if __name__ == "__main__":
    dim, n_classes, per_class = 8, 3, 200
    # three synthetic cell types as Gaussian blobs with distinct means
    means = torch.randn(n_classes, dim) * 3
    data = torch.cat([means[c] + torch.randn(per_class, dim) * 0.6 for c in range(n_classes)])
    labels = torch.cat([torch.full((per_class,), c) for c in range(n_classes)])

    model = train(CondDenoiser(dim, n_classes), data, labels, n_classes)

    correct = 0
    for _ in range(200):
        c = torch.randint(0, n_classes, (1,)).item()
        x = means[c] + torch.randn(dim) * 0.6
        correct += (classify(model, x, n_classes) == c)
    print(f"classify by generation accuracy {correct / 200:.3f}")

Conclusion

The core achievement here is a blood cell classifier that works by imagining cells rather than dividing them. By learning the full distribution of what each cell type looks like and classifying an image as whichever type it can best recreate, CytoDiffusion turns a set of separate clinical wishes, catching the abnormal, surviving a new microscope, knowing when it is unsure, learning from little data, and explaining itself, into natural consequences of one generative design. It matched conventional models on plain accuracy and clearly beat them on the harder properties that decide whether a tool is trustworthy.

The conceptual shift worth remembering is generative over discriminative. Most classifiers learn a boundary, which is efficient and brittle in the ways clinical work cannot afford. Learning the whole picture costs more computation but yields a model that flags the unfamiliar because it cannot recreate it, hesitates when it should, and can show what it was looking at, all from the same machinery. That its synthetic cells fooled expert haematologists is the clearest sign it truly learned the distribution rather than a shortcut.

The approach travels because its logic is general. Any medical imaging task where the rare, abnormal, and unfamiliar cases carry the highest stakes could benefit from a model built to recognize normality rather than to sort known classes, and the multidimensional evaluation framework the paper insists on is a better standard for the whole field than the single accuracy number it usually reports.

The honest limits keep it grounded. The method is computationally heavy, the results are retrospective rather than from a prospective trial, the hardest rare cases remain the true test, and the commercial interests behind the work are a reason for independent validation, not against it. This is a foundation and a framework, not a cleared clinical tool.

Future directions follow naturally. Bring the cost down, push the model to the fine grained cell subtypes where clinical subtlety lives, condition it on patient characteristics for fairness, and, above all, validate it prospectively on fresh slides in real laboratories against real outcomes. If those steps hold up, the quiet, expert act of finding the one wrong cell in a thousand could gain a tireless, well calibrated second pair of eyes, and the haematologist bent over the microscope could spend their judgement where it matters most.

Frequently asked questions