Expanding Flows for Fast and Flexible Generation Beyond the Fixed Canvas

[00:00:01]So today we're very happy to have Sophia Tang um here with us and she's a researcher at the University of Pennsylvania advised by Dr. Bernnam Chadi and this are currently a visiting researcher at the camera institute at Harvard University. Her research spans multiple areas of the A for science interative modeling with the guidance techniques for the scriptic fusion to theoretical shredding or bridge frameworks for generative modeling of branching and interactive interacting particle systems and uh yeah today she'll talk about expanding flows for fast and flexible generation beyond the fixed canvas. So with that thanks for coming and take it away.

[00:00:42]>> Um hi everyone, my name is Sophia. Thank you for the introduction. Um today I'll be talking about some new work um on expanding flows and flow maps for um fast and flexible generative modeling.

[00:00:54]And so a bit of an agenda of what we'll um talk about today. First I'll give a quick overview of flow matching and flow maps. Um, but I'm sure a lot of you guys are already familiar with those topics.

[00:01:05]And then I'll move on to kind of the motivation on why we want to um increase the expressiveness of these flows and flow maps to um expanding canvases and then I'll introduce the first idea in our paper um which is expanding generative flows and so expanding kind of like the the kind of standard OD to expanding um campuses and then expanding flow maps and then extend to discrete state spaces.

[00:01:37]So for a bit of background on flow matching and flow maps um as we all have seen many times probably um we can kind of define a continuous path between two um data points from two different distributions. So for example, x0 can be from a noisy distribution and x1 is from a clean data distribution. And then we basically integrate that flow um along this continuous path and we get a once we do this for many different samples from both distributions, we get a um marginal that follows the law of this interpolate. And um in generative modeling, we really want to be able to transport from samples from P 0 to samples from P1. And we do that by learning this marginal velocity which is the expectation of the times derivative of this interpolant um over all the different paths that pass this point x.

[00:02:33]And so the standard conditional flow matching we basically learn this by minimizing a conditional flow matching loss. And that's just the squared error between um some parameterized networks bhat t um over the interpolence minus the time derivative of this interpolate that we've defined.

[00:02:53]And with this conditional um velocity field, we're able to either integrate an OD um to get samples from P1 or integrate an SDE and get samples from P1.

[00:03:05]But the cost of kind of this um generative modeling technique is that we have to take many steps along this OD trajectory to be able to get accurate samples from P1. Um and that's kind of the motivation behind recent work on flow maps um which aim to kind of take larger jumps across this trajectory and to sample um from P1 with only one or two steps.

[00:03:31]And so um expanding on that a bit, a flow map is basically defined as um a jump between s and t which are the two time steps. And this can be any arbitrary pair of s and t between zero and one. And we basically take a step from x s um with an interval x uh t minus s and this mean velocity over this interval.

[00:03:55]And so um when we're learning this flow map on the diagonal when t is equal to s it's just the velocity field um or the instantaneous um od velocity that we've defined previously. But on the off diagonal it is this um mean velocity that we need to learn through consistency objective.

[00:04:16]And so there are three equivalent conditions for a valid flow map. Um the three are lrangeian. So basically this is saying that the map's end point follows this instantaneous velocity. Um and the next one is ulan which means that shifting the start time um still preserves the end point that the flow map gives at time t. And then the semigroup condition is when you basically take two jumps as we've seen in the previous slide. Um it's equal to just taking one big jump across the entire trajectory. And so this is a deterministic kind of map between time steps along this trajectory.

[00:04:55]And so we can learn this flow map by minimizing both the diagonal loss which is just the normal conditional flow matching loss and then also a consistency loss that either um is the squared residual of any of these consistency objectives. And so this kind of gives us a very nice way to either learn from some teacher model that already knows the instantaneous velocity or um self-distill. So learn the diagonal and then kind of anneal your way across the um offdagonal terms.

[00:05:28]And so more recently um this has been extended to discrete data um with the mean deninoiser which basically defines um a uh a single step that is between s and t kind of the mean velocity but also takes the jump all the way to time one and this can be defined as um you can see this kind of reversed into the mean velocity in this um in this slide as well. And so when we're taking the step on this mean deninoiser, we can take um kind of we can substitute in this mean den noiser and get this expression to take a single step along the flow map.

[00:06:05]And to minimize this, the nice thing about the mean deninoiser is that it's on the simplex space. So we can easily minimize um both a normal um classical dnoising loss which basically just takes the cross entropy between a dinoiser that just predicts the clean state from any intermediate state um and then also the consistency loss which now is taking a mean deninoiser and um this five or the sidebar is defined with the same consistency objectives but a bit different which I'll go over um in a later slide.

[00:06:42]And so kind of the motivation behind our work now is that we have these um flows and flow maps defined on this fixed canvas. And this can be like pixel space or a fixed sequence length. And so these continuous flows, they have very nice properties that you can um kind of adapt for a guidance and and fast generation, but there's no way to increase the um length of the sequence or the canvas size during generations. And that's kind of the limitation that we wanted to overcome with this work. And so in a lot of real tasks, if you can think of just language generation, auto reggressive models are automatically kind of um flexible length because you can like you can arbitrarily generate more and more tokens along um after um the last token.

[00:07:28]And so um our question here is how can we extend these flu fewstep generation of continuous discrete and multimodal data through these continuous flows um with these adaptive dimen dimensionalities. So how can we allow these flows to expand over time?

[00:07:46]And so kind of the big idea of our um of our paper is that we want to factor this flow map or this continuous um map between time points along the flow um into two different operations. The first one being an expand operation and the second one being a transport map. And so the expand operator basically either inserts new coordinates or tokens into this sequence or this canvas that we have that we're generating. And then the transport map does the same thing as um any normal trans like any normal flow map that um basically transports along this continuous flow. But now um we have um an expanding flow. So um the expansion operator allows us to kind of increase the dimensionality during um inference and during training.

[00:08:36]And so now this also recovers the fixed dimensional flow map um as a special case when the expand operator is just the identity function and doesn't um add new coordinates at all.

[00:08:48]So now I'll introduce the first idea which is kind of an extension of the um standard flow matching techniques into um expanding generative flows and how we can kind of create a well- definfined theoretical framework for these expanding flow maps.

[00:09:06]And so first I'll introduce kind of the idea of an expanding interbulent. So here we can see that we start with some dimensionality d0 and along this flow we're um going from d p 0 to p1 and at the end this um this dimensionality d1 is um a larger uh or equal dimensionality from its original dimensionality and along this flow um we can kind of imagine that um this dimension is increasing and the state space is increasing.

[00:09:39]Um but the problem is with this um kind of construction is that it's not really a well- definfined flow because there's no um kind of function that's able to map between these increasing dimensionality without introducing some kind of stoasticity or jump function. Um and so we want to overcome this with our kind of framework which we'll um define in the later slides.

[00:10:02]And so one of the key ideas that we um introduce here is that we along this um continuous um trajectory, we want to augment the source distribution with some extra noise coordinates. And these are sampled from this conditional distribution based on how big of a jump that we want to take between these coordinates. And that defi defines not only um what this noise is but also um how like how many dimensions of noise we want to sample from. And so this lifts this source dimension from time s into a target dimension um time t.

[00:10:41]And the first object that we'll introduce is this expand operator which is able to lift this source state into a target dimensionality. So you can see here um it takes XS which is dimension DS and as well as this noise coordinate that's sampled from the conditional noise distribution which is dimension DT minus DS which is the difference between these two um state spaces. And this ultimately gives us this new um expanded augmented state that is in the larger state space dt.

[00:11:15]And so again this combines the existing coordinates and the new coordinates via the noise coordinates um via some placement scheme which can be defined in many ways. The first way is some concatenation of this noise coordinates.

[00:11:29]So you can think of this as like auto reggressive modeling where we're concatenating um new tokens along the end of a sequence. Um so there's a lot of um different things that you can concatenate for example. So a point cloud you can concatenate more um nodes in a point cloud.

[00:11:46]Um the next one is positional insert insertion. So you can um insert tokens or insert kind of nodes in a graph. Um not necessarily at the end of a sequence or in a graph. Um but this kind of uh inserts this kind of um obtains kind of some positional dependence of these noise coordinates.

[00:12:09]And the third one you can think of as um a node expanding into children nodes.

[00:12:14]And you can think of this as like if you have a coarse grain atom, we can expand into like smaller atoms or um kind of it's the atoms that it's bonded to. And that's like an expansion. And so there's many ways you can define this arbitrary um um expand operator.

[00:12:32]And yeah, this augmented noise that we've defined can also be can either be deterministic. So you can um for example autograph modeling it's deterministically applying new tokens at the end of a sequence but it can also be learned via a neural network. So if we want to parameterize some insertions along the sequence we can learn that um during training as well.

[00:12:56]And so this expansion also kind of defines this um conditional distribution time t given a state at time x s and this is um defined as this push forward of the um augmented noise distribution.

[00:13:13]And so once we've inserted these noise um kind of coordinates now they're kind of on a different time scale as the already inserted coordinates. So they're not all on like the same um kind of fixed dimension time kind of schedule that we would have in a flow matching interpolence. And now since new coordinates are being inserted, we wanted to find some local time dependence between these different coordinates. So the model kind of has knowledge of which tokens were inserted later and which were first inserted um at the beginning of this um trajectory.

[00:13:45]And so we define this um as a local clock that starts at zero at the insertion time and then moves along this same kind any [clears throat] type of schedule but here I define it as a linear schedule um to the terminal time.

[00:13:59]And so um each kind of coordinate now is kind of it's on its own kind of time trajectory uh to this terminal distribution.

[00:14:08]And so we define also like a global time s um that is either um the at time at the insertion time it's si which is the um the time of this like position i is equal to zero and at all times after the insertion time it's this expression which we um show here.

[00:14:33]And so now that we've defined this expand operator, we can now define also a transport operator along this expanding interpolant. Um so we can start with uh this um xs with the um superscript um epsilon which is the expanded augmented state and then we integrate this flow b tow um over the time interval s and t.

[00:14:59]And you can also see this on the local clock as we know that now all the different coordinates are on different clocks. Um we can also kind of show here that this is not just the global time that we have on across all coordinates but also local time um for each individual coordinate.

[00:15:17]And um importantly when we have this expanding state space um for the coordinates that are not yet inserted along this clock um their velocity is zero because we wanted to stay the same.

[00:15:29]Um but after they've been inserted they kind of flow along this clock that um is defined by their insertion time.

[00:15:38]And with this BT we can train this with a normal um conditional flow matching objective um on these active coordinates. And then the inactive or not yet inserted coordinates are um given a velocity of zero.

[00:15:54]And so uh we can now define with this expansion operator and this transport operator a expanding generative flow which we use with the large fi um that kind of transports from DS to DT state spaces.

[00:16:12]And again the transport operator now only operates on this DT state space after we've augmented it with this noise.

[00:16:22]And we can also define this distribution of this augmented state based on this conditional distribution that we've defined for the noise as well as the original distribution. And it's again a push forward of the um expand operator.

[00:16:38]And so this satisfies two large um conditions which is just the marginal consistency. So if we push forward this um noise distribution conditional noise dist or this augmented distribution or this joint distribution based on the transport operator we get the terminal um distribution at time t and also a coupling consistency. So when we have xst um it agrees with this interpolent flow between the source and the target coordinate.

[00:17:10]And so um if we think of this as kind of augmenting with different noise coordinates, you can think that um it's not really a continuous flow anymore because we're kind of taking these discrete jumps in between. And so to overcome this kind of um theoretical inconsistency, we introduced this idea of e-flows being these peace-wise determinist markoff processes which have a continuous component which you can see in the first term. um which is a smooth denoising from time um across the time coordinate as well as this expansion kind of operator or kernel that takes this expansion operator um at a specific time and increases the dimensionality based on this noise um coordinate. And this is all determined by some intensity or this um scalar which we can define as an insertion intensity based on whether we insert coordinates or concatenate coordinates um at the end of sequence.

[00:18:08]And so the state space now becomes a kind of union of all these different dimensionalities that can take form across this interolent.

[00:18:17]Um and now with this definition we not only can define a continuous flow between two specific state spaces based on this augmented noise that we add to it but we can also define like this continuous flow along the entire interval from 0 to one.

[00:18:36]And now we can also define this change of variables along this expanding flow um from an interval from s to t. And this is just a standard kind of flow um change of variables formula.

[00:18:50]And so for all s less than t we can define this continuous flow based on some augmented noise that we insert. Um and this um importantly we can see that this velocity field is only defined over the active coordinates that have already been inserted um before time t.

[00:19:08]And so the non-inserted coordinates don't contribute any kind of um flow before they are inserted along this flow.

[00:19:17]And now that we've introduced kind of this um idea of expanding generative flows, we can now kind of extend it to flow maps, which is kind of the nice um idea that we want kind of the goal of this whole thing is to be able to get some kind of fast flexible generation with these continuous trajectories. And so we can define this expanding flow map very similarly to how we've defined the expanding flow which is with this expand operator from time s to t and this transport operator from time s to t. But now instead of taking a continuous um velocity field of t, we can now take a jump with this mean velocity um v st. So this is jumping from time s to time t along this augmented state space. So now this expand operator first augments or first expands this um dimensionality with this augmented noise and now this mean velocity is able to transport from time s to time t.

[00:20:20]And so this is kind of the same definition uh as the standard flow map except now we have an expanding um state as the initial state of this Do do we still know that these flow maps will be defemorphism because you have the operation of the insertion which is pretty discontinuous operation.

[00:20:44]Would they is it still guaranteed that you have theorphism here at all?

[00:20:51]>> On we can think of this as like learning on the octave only on the active coordinates along this entire trajectory. So if you think of it as as like at the time like at the terminal time we have this um larger state space then we kind of project that back into the previous time and then we define the flow along that um and if we train on enough kind of data points with this expanding interpoline we're able to kind of recover um something similar to it if you if you morphism but not exactly since we do have these discontinuous jumps.

[00:21:25]Yeah. So that's kind of um what we try to uh recover with the the markoff kind of um definition of this this trajectory.

[00:21:41]And so this also satisfies a very similar tangent condition to a standard flow map where if we take the time integral or the limit from as s goes to t um it recovers this standard um instantaneous velocity. And so you can see here um the expanding flow map as time s goes to time t we recover this instantaneous velocity at time at this dimensionality dt.

[00:22:09]And um we can train this with a very similar objective as the standard um flow map which is including a diagonal objective that um first distills either distills from some teacher flow or from self-distillation.

[00:22:24]And yeah, we can define this as either a self dis distillation from data where instead of um having you can see here that x um superscript t1 is just the terminal state but now we basically constrain it to all the only the coordinates at time t. Um and then we have uh the instantaneous velocity if we're distilling from an ex existing um expanding generative flow and this also takes the same kind of identities as um flow maps but now we have this additional kind of expanding um state. So you can take you can define the lronian um as the same um being the end point following the same instantaneous velocity as defined as shifting the start time leaving the output unchanged. So the divergence goes to zero and then the semigroup if we take one jump along this expanding flow it recovers the same kind of um jump along from s and I think the only difference here is that um we can see here that vsss is um theoretically defined on a state space lower than dt which um this whole equation is kind the second term of that equation is on dt. So how we kind of overcome this is that we kind of define it as um a operation that is on the state space ds but now on the inactive coordinates it is just zero. So it's still like well defined that way. Looks like there's a question.

[00:24:12]>> May I ask? Do you do the same thing for semigroup? Like if your if your network predicts uh different dimensionalities from taking two small jumps versus one large jump, is it also that you set some to zero?

[00:24:25]>> Yes. Yeah. Exactly. So how we kind of overcome this like different dimensionality is just setting um some of these um velocities to zero during during training.

[00:24:36]>> Thank you.

[00:24:41]And um this leads us to kind of the same consistency training objectives as um now kind of expanded in um shown the complete form. Um but you can write down these kind of um objectives for each of these distillation or consistency um consistency identities.

[00:25:02]And now we can see here that um not only are we training this flow map, but we're also training this expand operator which we either which we can define in many ways. But I'll define kind of one of the most um one of the ways that we define it. And if the expand operator is not trainable, so if we have a deterministic function instead, we don't need to um kind of we don't need to train that part um and only train the uh the transport.

[00:25:30]And so you can see the fi kind of encompasses both the expand operator and this mean velocity.

[00:25:38]And another way we can think of um the expanding flow map is as a stochcastic flow map. So um you can see that adding this augmented noise if we don't have knowledge of what this noise is it inserts some kind of stochasticity into the process and that can be defined as a stochcastic um kind of map between initial state xs um any kind of augmented noise sampled from this distribution conditioned on the jump between s and t and then this gives us a kind of distribution over states um at time t. And so um if we're just given XS then this is kind of producing a distribution over states at time t rather than just a single deterministic deterministic state. Um but if we're given XS and an epsilon so like if epsilon is deterministic then it is just one state.

[00:26:33]And so uh kind of a nice unique uh or a nice kind of instantiation of the expanding flow map is in the discrete state space where we've seen with um recent work has really allowed kind of fast generation of text and sequences as well as graphs. So um to define this discrete expanding interpolant um instead we now have um a very kind of uh nice way where we have this instantiation of the expand operator as insertions along this sequence. And so we can define this insertion schedule alpha t. And this is just the probability a of a token being inserted before time t um or uh yeah so the probability of a token i being inserted before time t. And this is zero at alpha 0 or at t equals z and one at um t equals.

[00:27:27]And so now we have defined this um state which is empty when it's not yet inserted. And then it is also a state along this um continuous flow um after it's been inserted. So you can see here this is just a linear interolent but now it's starting at an insertion time that might not be just t equals z.

[00:27:51]And so the sequence at time t is just a subsequence of the currently inserted tokens d noiseis to its local time ti.

[00:28:01]And this is some dimensionality dt which is now the length of the sequence um times v which is the vocabulary size.

[00:28:11]And so the way we can parameterize this insertion operator is with um here we define two time insertion um expectation which is defined with the rate of insertion alpha t. And this can give us um if you're integrating over time t to t time t equals or time s to time t you get this expression alpha t minus alpha s over 1 - alpha s and this is multiplying by the total gap count at time s. So you can see here um again here you can see the gap the one time kind of insertion count is just the expectation of the number of missing indices between um position i and position position i plus one and this is what gives us um this gap count and this also defines kind of um this insertion expectation but when we're training we can also define it as I'm training this diagonal insert insertion which is just the total insertions that we have left at time t and then we can train the off diagonal terms with the number of actual missing tokens that we have when we're um defining this interpolence. So during training we obviously have this clean state and then we either delete kind of tokens along the sequence and then noise them and so we know exactly how many tokens should be inserted over this time interval ST and so that's how we're um supervising this training and you can see this fi term is just the um pluson divergence between these two terms and so if you take the minimizer of that it's just um giving us the first first uh the first argument of this um of this slide.

[00:30:06]And so now we have this discrete expand operator um with this insertion uh kind of parameterization. We can define this insertion operator now as just inserting these noisy tokens um somewhere along this sequence. And so um you can see here I just show that you can concatenate these noisy tokens that have local time tials 0 um at their insertion time um along the sequence that is um kind of being denoised along their local times.

[00:30:36]And so this insertion law um during inference is uh defined we define it as a binomial distribution um give that is capped at the maximum length that we want to generate. Um so this can be ar arbitrarily large depending on what we want to actually actually generate. Um and it's also kind of predicted at the mean of this um insertion um parameterized insertion jump between s and time as time s and time t. And so this gives us a very um nice way to do a few step generation. So in many step generation you can think of um taking these small steps. So we're only inserting a few coordinates at time and this converges to just sampling li from this pluson distribution at um this mean um i hat that we've parameterized. And so this also is equal to just an infinite infinite decimal jump um for the pdmp across this law that we've defined as the insertion intensity. But when we're taking few steps um this probability of generating like all the tokens as at once um kind of goes to one. So we have this insertion in intensity going to one.

[00:31:49]And so in this case uh we don't want to just sample from the quanton distribution because it really has um high variance and is unbounded. So instead we sample from this binomial distribution which is bounded and has variance that vanishes as we take larger and larger steps along this interpolent.

[00:32:06]Um so this is a nice way that we can parameterize this insertion process while making sure that the sequence doesn't grow arbitrarily large or blow up um when we're sampling these lengths.

[00:32:24]So now we have the expanding operator and now we can just define um this discrete expanding flow map with um a similar parameterization as the recent work um in discrete flow maps which basically defines it as a mean deninoiser but now it's on this augmented state space um with this augmented noise.

[00:32:47]And so first the expansion operator inserts um these new tokens equal to dt minus ds and then we d noiseise um 2 * t with this mean deninoiser that lies on the simplex.

[00:33:00]And so as we've shown in the previous slides, the mean deninoiser is on is lies on the simplex which allows us to um minimize it with some very nice cross entropy losses which um are very stable for uh the simplex space especially with higher and higher dimensionality.

[00:33:20]And so now we can define um these cross entropy objectives. So for this um objective we basically take the kind of similar form as existing work where we have this diagonal loss which is just the dnoising or cross entropy between this dino noiser at a single time t um and this can be defined either with a teacher dino noiser or with um the diagonal entries of um its own deninoiser so with self distillation um and we're training this um mean deninoiser uh sigh um against the one-time deninoiser. And then the second term which is highlighted in blue here is just enforcing um some consistency uh equation which we choose to be the semigroup consistency equation which um takes this form of a stop gradient between um the jump from s to on this augmented state space. So you can see with the um superscript epsilon that it's on this augmented state space and then the jump from u to t and this um omega term is just defined as this kind of expression of the times which is also derived in a previous previous work on discrete flow maps.

[00:34:40]And so together with this insertion loss and this um flow map loss, we're able to train a kind of expanding flow map that is able to both expand along this interpolent in few steps as well as d noiseise in few steps. And this together can kind of define allow us to do fast inference um with um expanding sequences.

[00:35:05]And so now um before I get into experiments, I wanted to pause for questions if there are any. Um >> you made the assumption that the dimensions are increasing. Is that necessary? So can you have deletion and insertion and so on?

[00:35:21]>> I I for our um construction we do um have it be necessary for the dimensionality to increase but I think an interesting future work might be to also allow um decreasing dimensionality.

[00:35:34]But we can also think of I guess the decreasing dimensionality as just um kind of as an all like as an intermediate step between this increasing dimensionality. So um yeah I think there are many ways that we can construct it with also deletions along this inter line.

[00:35:55]Okay. So um the first experiment that we wanted to kind of show was with um continuous state space which uh a nice kind of um application of these expanding flows is if we have a coarse um kind of construction of a molecule we can expand it into kind of more fine grained um coordinates. So you can think of like larger heavy atoms as being these coarse representations of the coordinates and then expanding into hydrogen atoms and a more fine grain kind of um molecular structure. And so we frame this as a um we kind of our experiment is for conform generation and in this case the number of atoms is fixed. So we have a deterministic expanding operator and so we train this and we showed that it had does actually really well in few step generation of these compformers um on both QM9 and Gome drugs and you can see it's very comparable against very um large numbers of steps um on diffusion models. So we compare it against kind of the state-of-the-art of diffusion based conformer generation and show that both the e-flow which is our non-distilled model um can get away with only a few as 20 steps and generate very um high quality confirmers and if we distill it further with efm so with the distillation loss we can get even um accurate confirmer generation with um sposs four to two in even one app.

[00:37:30]>> So in this cases and the experiments that you run the final dimension is the full dimension of the thing that you're trying to generate and you start from very small ones.

[00:37:40]>> Yeah. So we start here from the dimensionality of the coarse grain or the heavy atoms and then we expand to the dimensionality containing the hydrogens.

[00:37:48]>> Oh, so it has it has domain specific knowledge.

[00:37:53]So yeah, it's given the like it's given the molecular graph that we want to generate and it's just generating the coordinates. So it's a continuous flow.

[00:38:00]>> I see there is this paper I think called transdimensional or something where I think they generate the full molecule by starting basically from nothing and then they add the coordinate. It's similar in the idea that you need to add dimension but they do it by having a diffusion model. Um >> yeah, >> so I don't know maybe you compared it in it's not exactly the same experiment but maybe you made a comparison. Don't know if you're familiar with that paper.

[00:38:27]>> Yeah, I think some of these are from very recently. So I think we tried to compare against most state-of-the-art models that we've we've seen. So a lot of these diffusion models may be able to expand, but they aren't trained with the same distillation losses as we've defined in our paper. So I guess in that sense it's very different um how we're kind of training these models to expand and um integrate this flow.

[00:38:49]>> Yeah. Yeah, it is different because what they do they just combine a a continuous diffusion model in a discrete >> diffusion model where the discrete it just decides do I add another coordinate or not. So it's certainly different.

[00:39:04]>> Yeah, very interesting. But yeah, we did compare with um diffusion models on this as well as some other there wasn't room on the slide to put it, but we did put like compare against some other more um older models that also did confirm regeneration.

[00:39:19]>> One question here because you mentioned at the beginning that you have like three different ways of inserting like new tokens.

[00:39:26]How like which one are you like are you using like how you're deciding whether like where to >> Yeah. So for this one we're actually using a deterministic insertion schedule. So um since we have a confirmer that we already know how many atoms are with each heavy atom. We just um insert them at different times along this interpolant and then d noiseise them from there. Um but yeah I think with molecular generation we can which is is actually in the next slide um we do graph generation where we don't actually know the number of atoms that we're generating >> and yes >> I guess when you insert the hydrometers and atoms that depends on the specific structure that you have at that time. So is the number of inserted uh like uh positions also dependent on the current structure? Um in this case we do uh we don't parameterize when it's inserted we just parameterize um the kind of the flow depending on these um differentially inserted tokens >> I see >> or inserted um atoms but yeah but I think that's another way that we can parameterize this but we see that it works well with the deterministic case though um in the next slide I'll kind of go into this parameterize insertion which um we show on molecular graph generation so in this case the task is more um to generate novel molecules and we don't actually know how many atoms that we have in this molecule. So, we validate this um against the QM9 data set and show that um against this multi-step um this kind of state-of-the-art multi-step um discrete flow model for graphs um we're able to kind of beat it across all the different steps as well as lower the FCD um of this generated distribution and against um CFM which is categor categorical flow maps is another paper that came out recent rec recently and um kind of showed uh generation on these um discrete graphs, we show that our method is able to um feed it especially at um one step which is uh kind of really cool. Um we can see that the SCT SCD is much lower. Um and that can be attributed to better >> validity here is molecular validity or atom validity.

[00:41:46]>> Um molecular valid with RD and so we show here that this framework is kind of robust across um multi-step and onestep generation of these variables graphs. And theoretically this can be um extended to larger graphs. Um we know that QM9 is kind of a smaller uh data set. So we assume that this can expand um to multiple different data sets that have large graphs.

[00:42:19]And then finally we show this on variable language modeling. So we can show um a kind of natural uh task at hand is to generate um sequences of text that where we can insert tokens between two different um text uh two different um tokens along the sequence along this generative trajectory. And so we show on LM1B that um this multi-step FLM or this multi-step um e-low framework is able to kind of be competitive with FLM on um discrete generation without um kind of even though it is kind of solving a harder variable length problem, it's still able to kind of preserve this generation quality which you can see both in the um lower gen PBL and kind of comparable entropy as well as these text samples that I have here.

[00:43:13]And then for fewstep distillation, we show that um this EFM frameworks is still competitive and coherent at two-step generation. Um and you can see kind of an example uh in this in the in the shown um but we also have more examples in the paper which will be out.

[00:43:34]And so we show basically that this um expand transport decomposition is very um generalizable across different data types um both in continuous state spaces as well as discrete state spaces. And we think that it's um a nice framework that uh really has a lot of applications um if you think hard enough a lot of different um applications can be applied with this framework and as well as like fixed lens generation is just kind of an instance of an expanding flow. um where we don't actually expand any tokens. But um yeah, so we hope that like these experiments kind of demonstrate that this can be a very generalizable framework.

[00:44:16]And so some conclusions um we basically in this paper introduced expanding generative flows as an overall framework of um expanding this dimensionality of a continuous time flow into uh this expanding interpolant and then extend that to expanding flow maps which allows us to do fast generation.

[00:44:36]And then we define this expand operator and this transport map that kind of increases the state space and then transports along it.

[00:44:44]And then we extend the both the continuous and discrete weight space um consistency identities for uh flow maps into expanding flows um that allow us to kind of increase the expressiveness of these flow maps.

[00:45:00]And this allows us to define kind of a unified recipe that can be applicable for um for example 3D molecular confirmers um discrete molecular graphs as well as language modeling.

[00:45:13]Yeah, thank you for listening. Um, I'll take any other questions if there are any, but um, yeah, feel free to also send over your questions via email. And then we also have the archive version of this paper, um, out sometime this week.

[00:45:27]So maybe let's open the floor for for questions. And I think there's a question by Edwick first and then Benjamin.

[00:45:42]Uh I actually think Benjamin was first so I'll let him go first.

[00:45:46]>> Sure.

[00:45:48]>> All right. Um I'm just trying to think uh architecturally say you're taking one step to generate from a single token into a full molecule. How like what is the network seeing at its input? Is it like one token that it's iterating over or you know what what is the dimension that it's it's operating on at that point? Uh so for onestep generation we actually do the expand operation first.

[00:46:13]So we expand it to some um dimensionality that we predict via the insertion network and then we den noiseise the full thing via just the onestep map. But now um since the model like we've actually discussed this more in the paper but since the model hasn't really seen a empty um graph at time or it's only seen empty graphs at time t equals z we um take the jump from this intermediate time which we define um more in the paper and take this onestep jump and that is performs kind of um pretty well onestep generation but we do uh discuss this more um in the paper as well was kind of discussed as a limitation of just the one step parameterization. But this for multiple steps, it's more well defined how we're taking these um these expanding and insertion or expanding and denoising operations.

[00:47:07]>> Thank you.

[00:47:09]Uh yes. So I have a question uh relating back to my earlier question about uh how you define the semigroup loss uh when uh the smaller steps and the large step dimensionality do not match. So it seems to me that if uh let's say the smaller step ones inserted more uh so the two small jumps maybe inserted from went from one coordinate to three coordinates and the large jump failed to match this. It went from one to two. So then you pad with a zero and then it seems like the loss from distribution matching is higher when the absolute distance from zero uh is larger um or like the loss from missing a dimension. So um in the paper do you show that this um model when trained sufficiently well recovers the data distribution or to me it seems like this would introduce some bias where um a dimension mesh is punished more if you're far away from zero so it makes enough translation invariant for example.

[00:48:22]Yeah, I think that's a very good question. I think for our like when we're training the distilled model, we take kind of this analing technique where we go from this these diagonal terms with small steps first um and then kind of first train also the um fixed state flow map and then also train the expanding operator not during like it kind of is a gets a warm-up between before it starts training this expand operator. But we do discuss kind of how we um train more in the paper.

[00:48:56]>> Thank you. I look forward to reading it.

[00:49:00]>> So maybe Trevor, you can go next.

[00:49:17]Hey. Uh, I have I have a question.

[00:49:20]>> Oh, sure, Ben.

[00:49:21]>> Yeah.

[00:49:22]>> So, this may be a bit of a naive question because this isn't really in my field, but I'm curious. Um, if you just said, okay, instead of expanding the dimension, the problem is you want to expand the dimension of as you go along. Well, what if you said I'm just going to set some maximum dimension size ahead of time and just have some sort of a null state as you go along that says when in generate if you know this particular dimension is just null and we're not going to include it in the final thing and I guess like and just use all your traditional methods that way. So what what would be like a use case or benefit of doing it kind of where you say I want to expand the dimension rather than just say I'm going to set an arbitrarily large dimension to start and then just fill in null states as I go.

[00:50:10]>> Yeah, that's a good question. I think for the expanding operator um we technically when we're training we do still have like some fixed maximum length that we have to because we want to use data >> the binomial right?

[00:50:23]>> Uh not necessarily. I mean that um if we have like data, we have to either pat it to some maximum length or more training, but the expand operator really allows us to diversify how we're expanding. So we can either parameterize it with um some like distribution that we can actually learn. So if you want to insert a token between two different tokens, um this has actually already been done with discrete diffusion models where we have these expanding sequence lengths. And so we have these um models that are trained to kind of take this context and then determine at each gap how many tokens do we want to insert and what are those tokens. So um if we kind of use this expand operator as a general framework um we can define either that insertion as just inserting like noisy states inside between two different tokens. Um and then at the end we have this padding that we can just discard as null state as you said. Um or we can even like learn some distribution that we're sampling from, which is why I define it as like a conditional um distribution of augmented noise. Um we can actually learn some distribution that it's sampling to insert tokens between um or like initialize tokens between two different tokens that we already have.

[00:51:37]So it kind of just expands the space or the design space of what kinds of um interpolence that we can define.

[00:51:46]But yeah, when we're training, we do have kind of like padding and fixed length that train with. Yeah.

[00:51:52]>> Yeah. Maybe let's go with Trevor again if you're if you have a question.

[00:52:00]>> Otherwise, there's a question in the chat on does this scale to proteins.

[00:52:05]>> Yeah, we actually haven't tried proteins, but um I think it like certainly you can try um training a protein model with this. I think with proteins um depending on if you're talking about protein sequences if it's protein sequences I think a lot of these techniques are very applicable um just as they are but if you are talking about um larger like structural um proteins I think there might be some training techniques that you may need to experiment with um to get it working well on those because those are larger and obviously very um has to kind of satisfy very many structural constraints to be a valid protein. So um yeah but with sequence I think you can definitely just train an expanding flow map on that.

[00:52:52]>> Uh I'll ask a question too. So you mentioned that >> uh ter you have a question.

[00:53:05][laughter] So okay I'll ask the question. Uh so you you mentioned that you can either insert noise or or insert or use a network to like determine what you're inserting which one did you use for these experiments?

[00:53:18]>> Yeah, so this one is actually just inserting noise but we like more recently tried kind of learning this um or learning some kind of mean um for this inserted token given like this previous or context of the before token just to see if it does better. Um, and we have inconclusive results as of now, but I think it is a possibility of just like priameterizing some additional network that's along with the insertion head that kind of predicts the mean of this noise that we sample from.

[00:53:45]>> Sure.

[00:53:46]>> And maybe related to this like do you insert like you when you insert you know how many dimensions you're going to insert at every point in time.

[00:53:55]>> Yeah. So the insertion thread would predict how many tokens we're inserting between two different tokens. And then this additional head that we're defining this new mean.

[00:54:04]>> So you're always broadcast. Yeah.

[00:54:06]>> You're always using an insertion head to determine how many that you always use.

[00:54:09]Yeah.

[00:54:09]>> Yeah.