last Tuesday at 8:31 PM
Hi HN, author here!
For context, we're a 2-person lab training generative video models. Goal is a new set of controllable, animation tools (you can read more about that here https://www.linum.ai/about if you're curious).
The biggest bottleneck for our last text-to-video model in terms of training and inference cost is attention. Video models are incredibly token dense (e.g. 110K tokens for a several second clip). If we can condense that context window more aggressively, we can train bigger models for a lot less $$ and offer them to prosumers at reasonable price points (unlike the big models today like Seedance, which cost an arm and a leg to run).
Traditionally, image and video models have two disjoint components: VAE (Variational Autoencoder) and Diffusion Transformer (DiT). They're trained separately, and empirically VAEs seems to struggle to get past 16x16 token reduction.
Here, we're switching to pixel-space, throwing away the VAE, and achieving 32x32 token reduction (4x smaller context windows) while learning a better overall model in a fraction of the training samples.
The central thesis is "simpler is better". If we can put the compression problem into the more powerful Diffusion Transformer (DiT) would should be able to learn a "latent space" optimized for generation and get better compression without hurting generation quality.
I'll be checking this post off and on the next couple of hours, so feel free to drop questions below. And I'll try to answer them to the best of my ability.
P.S. The model checkpoints from this blog are Apache 2.0, so feel to try playing with it yourself on a GPU!
last Tuesday at 10:33 PM
I know it goes a tiny a bit against the spirit of what y'all are doing, but applying a few layers of pixel-wise local attention (so 1x1 "patch", with 3x3 or 5x5 attention window, basically treating it as a dynamic conv) has worked way better than both linear and conv unpatching in my recent experiments.
Diagram for reference https://x.com/1rreverant/status/2107546198093730287 (In my most recent recent experiment I actually removed the MLP and just used a linear project on the pixel's hidden states)
last Wednesday at 2:10 AM
That sounds like an interesting idea!
Can you confirm I'm understanding correctly?
1) Linear unpatchify as usual to go from hidden states to pixel space
2) Attention within a local window (e.g. 3x3, 5x5) to "blend" pixel space data and come up with a better image (as an alternative to MLP or Convolution)
And follow up questions:
1) How do you handle boundaries between your "attention windows"? Do you move the window just like a convolution does or are the "attention windows" all mutually exclusive from one another?
2) How much faster/slower is this operation vs. a linear layer + MLP?
last Wednesday at 2:35 AM
1) yes, but with more than channels than 3 (and no correspondence to color, its "pixel space" in the sense of position, not value)
1D grid example for clarity (obviously meant to be done in 2D though)
So embed -> [N dim, N dim, N dim ... , N dim] instead of embed -> [RGB, RGB, RGB ... , RGB]
2) / [followup 1)] Stride of 1, so we place a window at each pixel (so lots of overlapping)
Also I wouldn't phrase it as "come up with a better image", the point is to give less spatial decoding pressure to the patch tokens so that they can almost completely focus on feature learning instead. There is no reason to have the model learn spatial decoding when the structure prior of images is comically strong (especially compared to text), its a waste of training time and parameters
[followup 2)] I haven't measured but it was passable is all I can say (my experiments setup are horrendous right now lol)
(Also don't mind the phrasing, I just wanted to be 100% clear)
last Wednesday at 2:40 AM
I appreciate the clarifications!
If we have extra compute lying around in the coming weeks, we’ll try this out and report back. It’s a good idea :)
last Wednesday at 2:43 AM
I loved to hear! Just ping me on twitter (same account I sent the link with)
today at 8:00 PM
Neat. What are your thoughts on the recently released Iris-3B, a pixel-space model which also bypasses the need for a traditional VAE?
https://arxiv.org/pdf/2610.09450
https://huggingface.co/speridlabs/iris-3b
today at 8:10 PM
Thanks for links! I hadn’t seen this yet.
Very related; our models are quite literally cousins, as we’re both interesting on the work of JiT from last fall.
Architecture is roughly the same up to some minor differences (they use MM-DiT from FLUX, we use single stream from Z-Image).
The biggest delta is that we’re going with 32x32 token reduction and they go for 16x16. Our goal here was to cut attention windows 4x to speed up training inference. And have been publishing work (this included) on how to match models with less compression to than us.
In this blog specifically there are 3 things we stack on top of each other to get good results:
1) transformer blocks per modality before they enter shared single stream transformer blocks
2) very different noise schedules during training than those proposed in prior work
3) predicting the output at k different resolutions along the depth of the diffusion transformer to provide gradient signal earlier in the network to learn better (and faster)
The other work does a great job scaling up the JiT to 3B parameter; this is purely an experimental release highlighting how we make up the gap.
today at 7:29 PM
I wonder how much of that is due to the use of a suboptimal VAE? The same optimisations could be applied to it, and my intuition tells me that your total compute spend would be even more optimal if you retrained the VAE yourself with a better approach (esp. ensuring translation and rotation invariance, + ability to rescale/blur the latents).
You could also have the same advantages of a draft in low resolution latent space, with an easier to learn data distribution that's more robust to perturbation, but instead of 512x512, you could get a 4096x4096 output for the same compute (assuming a 8x VAE).
There's also the advantage of being able to use a high number of diffusion/flow matching steps for the main LDM while the VAE can be single step and much smaller, since it does not have to handle language or significant scene understanding, just perceptual compression. This sounds especially important for a video model where I would be extremely hesitant to train a generative model without relying on interframe compression.
today at 8:08 PM
If you think about it there's overlap in what the VAE encodes and main model encodes. Objects at a distance resemble texture and textures zoomed in gain structure. The VAE makes textures more efficiently representable at the cost of reducing the representable space of pixels. So things like tiny text become nonsense scribbles. Working in pixel space, especially with something with recursive or cascaded structure, opens the possibility of using the learnt structure of real writing at a higher level to perfect tiny details that actually cannot be approximated without being obviously wrong.
Time is "just" another dimension. There's temporal continuity between frames, a video VAE would be learning and representing those temporal shifts, but there's nothing to say that e.g. a recursively applied generative model at the pixel level also doesn't learn and represent those things.
As ever, figuring out how to train the thing is the hard bit I expect.
(Handwaving over "textures" here, VAEs encode more like somewhat macro blocks of image whose content is also conditioned on surrounding blocks, rather than tiny patches of patterned pixels.)
(And yes I'm a total imposter layman here, I just see VAEs as seeming to be a crutch that reduce data size - super super helpful of course - but being strictly speaking redundant and inhibiting correct fine detail.)
today at 7:56 PM
This is a totally fair point and definitely worth exploring!
The jumping point for this no-VAE work was three-fold:
1) Our goal here is to get 32x32 token reduction to make video training and inference downstream cheaper. To date, the best open-weight Image VAEs like Flux-2 seem to cap out at 16x16 token reduction (8x8 VAE + 2x2 linear patchification). Others like H3 have pushed to 32x32 reduction but requires them swapping out the small VAE decoder with a 2B parameter decoder. So, this is a foray to get 32x32 compression without compromising quality.
2) We believe that end-to-end trained networks will tend to perform better than modularly trained networks (e.g. VAE + DiT). This hypothesis comes from work like REPA-E, where authors are able to get much better results by backpropagating through the pretrained VAE. The latent space for perception / reconstruction seems to have a different "optimal" configuration than a latent space specifically built for generation. That's why we liked the idea of trained E2E here.
3) There's work with VAEs that show that providing additional modalities (e.g. text captions) can help the VAEs improve as well. That's natural to this construction, so we thought language might actually help with the compression, not hurt. To be honest, this this is the most hypothetical of three ideas; and definitely warrants specific ablations.
today at 8:01 PM
Also i failed to clarify that the VAE for Linum v2 was the Wan2 VAE
today at 5:44 PM
Is there a way to use this in ComfyUI today?
today at 6:59 PM
We haven't built a ComfyUI integration; but all the weights and inference code are available to download. So, it should be pretty easy to point Codex/Claude to it and have it re-arrange stuff to be compatible.
From our POV, this quite an experimental checkpoint. That's why we've dubbed it a research release. Our goal is to share our findings (and checkpoints) along the way to a legit image and video release for Linum v3.
Image release should be out in the coming weeks, so stay tuned for that. Video will take a little longer (we're doing on-going research atm on pixel-space video training).