BenjMurrell

A technical thread on loss scaling in diffusion and flow matching models (related to a new preprint): Since the dawn of time, people have been messing with (or dropping entirely) these pesky time-dependent loss scaling terms, mostly because the models train better without them.

arxiv.orgEdit Flows: Flow Matching with Edit OperationsAutoregressive generative models naturally generate variable-length sequences, while non-autoregressive models struggle, often imposing rigid, token-wise structures. We propose Edit Flows, a non-autor...

(We also show this for Edit Flows, which lies just outside of Generator Matching) arxiv.org/abs/2506.09018

We're still newcomers in this space, and the literature is vast, so if you think we've missed something, or have any key citations to suggest, please let us know and we'll be happy to update the preprint!

arxiv.orgGenerator Matching: Generative modeling with arbitrary Markov processesWe introduce Generator Matching, a modality-agnostic framework for generative modeling using arbitrary Markov processes. Generators characterize the infinitesimal evolution of a Markov process, which ...

For this, we closely follow the incredible Generator Matching framework, which subsumes nearly all flow matching and diffusion models. arxiv.org/abs/2410.20587

This means that if you're training a flow matching or diffusion model, you can happily rescale your loss with any (positive almost everywhere and integrable) scaling function, and the theoretical guarantees from Generator Matching will still hold.

We want things to be a bit less handwavey than this though, and we needed it for Branching Flows. bsky.app/profile/benj...

BenjMurrell@benjmurrell.bsky.social

We figured out flow matching over states that change dimension. With "Branching Flows", the model decides how big things must be! This works wherever flow matching works, with discrete, continuous, and manifold states. We think this will unlock some genuinely new capabilities.

arxiv.orgTime dependent loss reweighting for flow matching and diffusion models is theoretically justifiedThis brief note clarifies that, in Generator Matching (which subsumes a large family of flow matching and diffusion models over continuous, manifold, and discrete spaces), both the Bregman divergence ...

But when we looked, we couldn't find a sufficiently general justification. So our preprint, driven by @lukasbillera.bsky.social with assists from @hedwignordlinder.bsky.social, formalizes this, and extends it a little in ways that are trickier to heuristically reason about: arxiv.org/abs/2511.16599

Intuition suggests that it should be safe to re-weight your loss differently across time. E.g. If you trained a different model for each timestep, a time-dependent scaling would just be a constant for each model, and wouldn't change what is learned.

With my wonderful lab, who mostly aren't on here (except @lukasbillera.bsky.social and @antonoresten.bsky.social ?) we've been tinkering in this space since the end of the summer, but we think this is just too cool to sit on any longer. The manuscript should be up by tomorrow and I'll drop a link.

Ah @hedwignordlinder.bsky.social is on here too!

Oh and as usual, these visualizations (which are, to us, key to understanding these processes) were made by @antonoresten.bsky.social in @makie.org (@simi.bsky.social).

Eg. here is a protein example where the model designs two domains with an intervening linker (light blue chain in vid). A regular flow model would need to know, early on, exactly how many AAs are needed in the linker, but Branching Flows can decide on-the-fly and grow or shrink it.

github.comGitHub - MurrellGroup/BranchingFlows.jlContribute to MurrellGroup/BranchingFlows.jl development by creating an account on GitHub.

We've got a flexible implementation in Julia (github.com/MurrellGroup...) that uses our Flowfusion.jl package so you can compose families of base flows with the branching and deletion process. If there is community interest, we can set up a python implementation as well.

Given that all of the deletions, trees, anchors, etc are in Z, which doesn't interact with the theory, it is easy to manipulate the trajectories the model is learning. Autocorrelated insertions? Change how you build the trees! Same for the "anchors" which control the process on internal branches.

We dislike ad hoc padding in discrete diffusion models, and the situation is even worse in continuous domains. Branching Flows removes this blemish. We also expect that this makes the trajectories easier to learn.

The process can be seen clearly with a QM9-trained model (continuous atom positions, discrete atom types), starting from a single atom.

You can see the branching pattern in the static plots "with trails", and if you stare at them you can see where lineages delete too.

The trick that makes it tractable to learn, for eg. continuous states, is that branching events are not generic insertions into space (which, we think, you need for TD jumps (@arnauddoucet.bsky.social ?). Instead, they duplicate the state that is branching.

Side note: this is heavily inspired by the processes from phylogenetics (the other thing my lab works on).

With Branching Flows, we draw X1 from the data. Then we draw X0 (one element? many elements? it all works). Then we add "to-be-deleted" nodes into X1. Then we draw a forest of trees, one per X0 element, that maps (one to many) the X0 elements to all the X1s (plus to-be-deleteds).

Then on the internal nodes of the trees, we place "anchor" states (same space as the X1 elements). We put all that in Z. Then, given this Z, Xt evolves over the trees, sampling when (but not which) branching and deletion events occur, all constructed to terminate at X1.

How does this work? First, you need to understand the amazing Generator Matching (arxiv.org/abs/2410.20587). In GM you first sample Z, and then construct a stochastic process Xt that, conditioned on Z, terminates at the data distribution.

Typically Z will include X0 from an easy-to-sample distribution and X1 from the data distribution, and if you're feeling fancy you might couple them ( @alextong.bsky.social ). But what we learned while working on Branching Flows is that Z is a *playground*.

And this is "infix sampling", but where you let the model figure out how many amino acids are needed to span the gaps.

Another infix, this time a nanobody CDR3.

For example, here is what "binder design" looks like when you start from a single amino acid.

Another binder, which threads a tail down a groove.

And here is a small molecule example.

This is our "graphical abstract":

21 June 2025

One day remaining before this closes! bsky.app/profile/benj...

BenjMurrell@benjmurrell.bsky.social

My lab, at Karolinska, in Stockholm, is looking for a PhD student with a computational/quantitative background to work on probabilistic/generative models of proteins (structure and sequence). The research will involve methods development, and applications in vaccine design.

We tried to set up a simple demo/tutorial model for the protein design ecosystem we've been developing, and it turned out a bit more interesting than we expected. 🧵 This was a team effort from a few people in my lab, including @antonoresten.bsky.social and others (not sure who is on this app)

We will have a preprint up about this soon. This is just a "base model" so you can't really ask it to do anything specific (but stay tuned). We've put the code and weights up, so give it a spin: github.com/MurrellGroup... If you have a GPU, it is pretty fast.

Shoutout to the @julialang.org community, where we rely on Flux.jl, CUDA.jl, Zygote.jl, Manifolds.jl, and others as well as @makie.org for the visualizations!

n the LLM context, "induction heads" allow copying via recognition of repeating tokens. This is a bit like that, but SE(3) and non-AR? Tangentially, is anyone aware of what induction heads look like in text diffusion models?

We can mask that head's ability to attend to non-self AAs, and this prevents the model from generating symmetric backbones, without damaging the in silico refoldability of the structures themselves.

Why? Most other methods apparently train on single chains (hat tip: @sokrypton.org ), and then use eg. AA pos offsets (sometimes with multi-chain finetuning?), to handle multiple chains. Chroma trains on multi-chain structures, but maybe symmetry is less visible to random-edge GNNs?

We dug in a bit, and it turns out that, among all of the layers in this network, there is a single (!!) attention head that attends to the matching residue of any symmetric copies of a chain (and also to within-chain structural repeats).

After checking that we weren't just mistakenly looking at the training PDBs, we realized: these are, at least approximately, exhibiting symmetry. We saw dimers, dimers of heterodimers, trimers (though the model isn't as good at these) etc.

RFdiffusion, Chroma, and others also generate symmetric structures. But, as far as we know, they only do so if some sort of symmetry constraint is imposed during generation. We aren't doing anything like that. This model just does it spontaneously!

With the ecosystem all in place, this is what it takes to 1) specify the entire model, and 2) the flow matching training objective. Just mixing and matching components:

Like MultiFlow this does struct+seq, with struct similar to FoldFlow-SFM, and with seq per Meta's Discrete Flow Matching. We trained this for a week on one RTX6000 Ada, expecting sketchy helices, and maybe the odd beta sheet pairing if we were lucky. Instead, we got these:

Side note: If you're looking to do a PhD in this space, there is just over a week left to apply for this position with us: bsky.app/profile/benj...

BenjMurrell@benjmurrell.bsky.social

We've been tinkering in the protein design space for a bit now and, since we stubbornly refuse to work in any language other than Julia, we've set up a @julialang.org ecosystem for flow matching, a collection of useful layers, a protein data processing pipeline, etc etc.

Karolinska has a rich work environment, and you'll get to collaborate closely with experimentalists. If you already have a strong biology background then you'll fit right in, and if you don't then this is an excellent way to broaden your education.

kidoktorand.varbi.comDoctoral (PhD) student position in probabilistic models of proteinsDo you want to contribute to top quality medical research? To be a doctoral student means to devote oneself to a research project under supervision of experienced researchers and following an individu

This position is part of (and funded by) the @scilifelab.se "Data Driven Life Sciences" initiative, and you will integrate into their new DDLS PhD "Research School". Please feel free to email me if you have questions, and please share this! Ad link here: kidoktorand.varbi.com/en/what:job/...

Coming from South Africa (via five years in the US), I can say that Stockholm is a wonderful place to live. The city itself is amazing, and, if you need some nature, there are forests and water everywhere. And you never need to speak any Swedish!