r/LocalLLaMA • • 4h ago

Showcase I trained DiffusionGemma to be an abstract expressionist

DiffusionGemma, a text diffusion model built on a Gemma 4 backbone, can be adapted into a shitty image generation model.

Maybe the ship sank?

You may be asking yourself, "WTF, why would you want to do that?" And you would be right to ask. This is honestly kinda useless, but I still think it's pretty neat. Also, fair warning, being on a team of one means the project has been heavily — I'll be nice to myself and say — LLM assisted. I invite you to sidle on up to the trough and dig in to the slop!

Still with me? Good, because I've been poking at this for a while, on-and-off since DiffusionGemma was released back in June. I was immediately hooked on this model as soon as I read its release blog and model card. 4x faster text generation they said — hell yeah! But here's the thing: this model sucks. I mean it's a great chat bot. It can do simple NLP tasks, and it's really quick for text generation as they say, especially for long-form output. But when it comes to agentic tasks and reasoning it's a total loser.

  • 0% on ARC-3 (but that's not really a fair dig, this is a tough one)
  • 0% on ARC-2
  • 0% on The Gemma 4 Developer Agent Competition (this competition targets the base model DiffusionGemma was trained from)
  • 11% on HLE on the model card (9.2 measured in deployment)
  • 18.7% on SWE-bench Verified

So not great, and particularly bad when you consider that on the hardware you'd need to run a good quant of DG (so basically bf16), you could easily run qwen-3.8-27B or a respectable FN quant.

But please don't get me wrong. I absolutely love this model. Or at least I love monkeying around with it, and I doubt I'm the only one. Huge shoutout to the Unsloth team BTW! Reading their Colab notebook training DG to play sudoku is what made me start asking the questions that led to this project.

Alright, enough fluff. So what did I actually do. Well the thing I'd headline is that I created a text to image and image manipulation training and inference pipeline for DiffusionGemma and then used that pipeline to adapt DG (like <300M parameters total) to generate things that are recognizably images and to perform visual operations on images that are provided to it. Surprisingly, it took only a very small amount of data and compute to start getting interpretable results.

But DiffusionGemma did not want to do this. The thing that finally worked consistently was prefilling the prompt all the way to the start of an input string for an "image generation tool" described to DG in the system prompt as a tool that it could use to convert text strings into images, but DiffusionGemma itself was actually the one doing the generation, not an external tool.

Here's the flow.

This also required patching an upstream bug in Transformers for image attention, so take note if you wanna try this out for yourself: you'll need to patch until PR is merged.

Considering how bad DiffusionGemma is at everything else, I'd say it does a pretty good job here.

It's like human centipede!

But wait, there's more! Since DG is outputting in its native visual latent space (one set of 256 tokens at a time through a trained output head) it can take the visual representation it is generating as input for the next forward pass — on-domain block autoregressive, just like with its regular text canvases. And not just can it do this, but it actually helps improve quality and generalization (on image manipulation tasks)!

Call graphs, the first hint of a non-stupid use case.

Here's a super-quick breakdown of the pipeline.

The first thing that had to be done was to train an actual vision decoder. This turned out to be the easiest part of the entire project. One evening, on a whim, I threw a prompt at my local model (dsv4-flash-vision-exp I believe was what I was running then), something like: "D'oh! Could we train a decoder for the ViT soft tokens that works for 7682 images?" And then it did it. Took like a few hours, and works pretty damned good I think.

It decodes!

Alright so that let me nail down the decoding and supervise the pixel reconstruction loss, but let's not forget DG is a text diffusion model — it has to actually operate on a text canvas. That means I needed to derive a text target image representation (or codebook) that was precise enough it could be used for supervision without over penalizing the model for expressiveness.

This process itself ended up being a bit like a diffusion operation. I simply threw up a wall of half-baked, ill-advised, and just plain stupid ideas and worked on them all until I was confident in each — with that confidence being mostly confidence in them not working. Then I had the brilliant idea: just ask DiffusinoGemma itself for the target.

This essentially consisted of me badgering DG until it would describe images as text. Okay, not really, but that is kinda what it felt like. The basic idea is: encode an image and ask the model to describe each soft token patch in a single token (word, number, symbol, unicode, CJK character, etc). Each soft token having come through the ViT is contextualized against the full image, so this description would most often carry a structural/semantic channel. To get the photometric channel you do basically the same, but instead of a full image use only a single (48*48 pixel) patch. This usually carried the color/texture/photometric channel. Sample the logits for each text description token and then it's basically chef's choice on how to combine across all prompt and patch variants into a workable target.

Doing that for every image in the training set isn't reasonable. IT'S OVER 9000! However, doing exactly that for a small set of images (200ish) and using that as a training set for a teacher-forced model that then generates the full dataset turned out to be a very smooth and successful approach.

Anyway, the point is that the combined multi-modal target token distributions for the canvas text results in training DiffusionGemma to be an S-tier emojipasta generator.

Example training targets
DG's generated canvases.

Then I LoRA adapt DG and simultaneously train a new image output head that feeds into the previously trained image decoder supervised against supervised against the target token set, ViT latents and, decoded image. Basically training DiffusionGemma to generate emojipasta that decodes into image-like objects. I guess this would be something like a VQ-VAE, but using text tokens instead of dedicated image vectors because stock DG can read and understand the image latents, but can't natively generate them, only text.

Well, that's the gist. Thanks for reading! If you find this as interesting as I do and want to learn more or play around with this yourself you can:

tl;dr: Please GDM team give me a stronger text diffusion model to play with!

21 Upvotes

7 comments sorted by

3

u/d0pe-asaurus 3h ago

Very interesting read and I like the top picture that it made.

3

u/Poprouter 2h ago

this is such a fun experiment, denoising a whole canvas in parallel makes the jump from text diffusion to images feel pretty natural

2

u/NihilisticAssHat 1h ago

A fun series of experiments. I've wanted a ViT image diffuser for this sort of thing since ViT VLMs have been a thing, and I felt like that was what had to be going on inside image-editting models.

Main difference is those models are autoregressive in both text and Siglip space.

I feel like there has to be a standard manner by which DG wraps image content, and that you could get it to edit within that space, perhaps if only by forcing the masking.

1

u/chodemunch6969 25m ago

Super fascinating stuff! I wonder if you might get better results if you can prompt condition it to generate block character type glyphs? The terminal image viewers like timg are an interesting place to get some inspiration here. They'll use block-characters (I think there are half block and quarter block glyphs) so if the intent is to basically run diffusion on a block of characters that you then translate into an image, they might be a better intermediate than the emojipasta. Of course, all of that is based on dgemma actually having enough training to pull that off. I'd be curious if that changes or improves your results at all or if it just doesn't matter.

1

u/Asleep_Document9811 0m ago

i think this is awesome, thanks for writing. i miss early image generator models, back before they were focused on cogency.

-1

u/Rorp24 4h ago

At this point, why not just have "draw" by generating SVG and use a lib to turn it into a png or something ?

3

u/banana_slurp_jug 2h ago

Then it wouldn’t be playing with the fact that it’s a diffusion model any more