Trajectory Replay: Continuous-Time Mixture Transport for Lifelong Generative Models

Alessandro Palma; Charles Jones; Tristan Cinquin; Fabian J. Theis; Frederik Wenkel

Figure 6: Two-mode mixture updates. ( a ) Samples and one-dimensional marginals from the real mixture and from trajectory replay after the update, for historical and incoming modes at

\lambda\in\{0.1,0.5,0.9\}

. ( b ) Samples and marginals from the real mixture and from naive fine-tuning after the update, shown for

\lambda\in\{0.1,0.3,0.5,0.7,0.9\}

.

Figure 7: Energy distance versus teacher sampling budget. Energy distance to the target mixture, computed on the active coordinate that separates the modes, as a function of cumulative teacher NFEs for trajectory replay and generative replay, in ( a )

d=128

and ( b )

d=256

. Lines show the mean and shaded bands the variability across three seeds. The

x

-axis is logarithmic and lower energy is better.

Figure 8: Five-task Gaussian ring stream. ( a ) The five task distributions

T_{0},\dots,T_{4}

placed on a ring, with the arrow indicating the order in which they are introduced. ( b ) Samples drawn after the final update for the balanced target and for naive fine-tuning, generative replay, and trajectory replay. Crosses mark the task centers, and colors follow the task legend in ( a ).

Figure 9: Replay-weight policies on the five-task stream. ( a ) Trade-off between retaining earlier tasks and fitting the newest one after the final update: each marker is one final model, positioned by its energy distance to the historical task distributions (

x

) and to the newest task (

y

), so points closer to the origin are better on both axes. Colors distinguish generative from trajectory replay, marker shapes distinguish constant-

\lambda

, equal-task, and adaptive policies, and the number next to each constant-

\lambda

marker is its

\lambda

value. The circled point near the origin marks the uniform target. ( b ) Replay weight

\lambda_{k}

selected at each sequential update for adaptive generative and trajectory replay, with the equal-task rule

\lambda_{k}=k/(k+1)

shown as a dashed reference. Error bars show variability across seeds.

Figure 10: Stream evolution on the five-task ring benchmark. ( a ) Energy distance to the balanced mixture and ( b ) KL divergence of the generated task-mass composition from the uniform target, both as a function of the number of tasks seen, for naive fine-tuning, generative replay, and trajectory replay. Bold lines show the mean and faint lines three individual seeds. Numbers at the right give the final values, both axes are logarithmic, and lower is better.

Figure 11: Task-0 budget and backbone selection on ImageNet-100. FID on the task-0 classes as a function of the number of optimizer steps on task 0, for two latent UNet backbones: the medium model with block widths

(128,256,256,256)

and attention at one resolution (

35.8

M parameters), and the large LDM-style model with widths

(128,256,512,512)

and attention at the two lowest resolutions (

106.4

M parameters). The circled point marks the configuration chosen for the continual-learning experiments; lower is better.

Data
	Fashion-MNIST	CIFAR-10	ImageNet-100
State	$1{\times}32{\times}32$	$3{\times}32{\times}32$	$4{\times}32{\times}32$ latent
Tasks $\times$ classes	$5\times 2$	$5\times 2$	$3\times(34/33/33)$
Horizontal flips	no	yes	yes
Model
UNet	medium	medium	large

Appendix

Table 2: Per-dataset configuration for the class-incremental image experiments. Sampler and replay steps are Euler step sizes, with the resulting number of function evaluations in parentheses.

Figure 12: Training wall-time comparison . Full-CL-pipeline comparison of trajectory and generative replay in terms of accounted wall time across datasets. The two methods are compared in the equal-task

\lambda

setting.

Figure 13: Class-incremental image streams. ( a ) Stream FID and ( b ) stream KLD as a function of tasks completed, for naive fine-tuning, diffusion distillation, and generative and trajectory replay under selected and equal-task

\lambda

. FID is shown for Fashion-MNIST, CIFAR-10, and ImageNet-100. KLD for Fashion-MNIST and CIFAR-10. Lines show the mean and shaded bands the variability across three seeds. Filled markers with solid lines denote selected

\lambda

and open markers with dashed lines equal-task

\lambda

, and lower is better on both metrics.

Figure 14: Class composition of generated samples after the final task. Stacked bars give the proportion of generated samples assigned to each class by the evaluation classifier, for ( a ) Fashion-MNIST and ( b ) CIFAR-10. Rows list the uniform target and each method: trajectory and generative replay under selected and equal-task

\lambda

, diffusion distillation, and naive fine-tuning. Numbers at the right report stream KLD from the target composition (mean

\pm

SE over three seeds). Lower is better.

Figure 15: Selected replay weights across the image streams. Mean replay weight

\lambda_{k}

chosen at each sequential update for generative and trajectory replay on ( a ) Fashion-MNIST, ( b ) CIFAR-10, and ( c ) ImageNet-100, with the equal-task rule

\lambda_{k}=k/(k+1)

shown as a dashed reference. Error bars show variability across three seeds.

Figure 16: Replay backend cost and quality on the image streams. ( a ) Final stream FID relative to the matched ODE run as a function of replay-field evaluations per update, and ( b ) final stream FID against accounted training time, for Fashion-MNIST, CIFAR-10, and ImageNet-100. Colors distinguish generative from trajectory replay, marker shapes distinguish the MF-1, MF-4, and ODE backends, and filled markers with solid lines denote adaptive

\lambda

while open markers with dashed lines denote equal-task

\lambda

. The dotted line in ( a ) marks parity with the matched ODE run. Error bars show variability across three seeds,

x

-axes are logarithmic, and lower is better.

Figure 17: Teacher NFE sweep on Fashion-MNIST. ( a ) Final stream FID and ( b ) amortized wall time per update as a function of the configured teacher NFE for the ODE teacher. Dashed and dotted horizontal lines in ( b ) mark the corresponding cost of the MF-1 and MF-4 teachers. Error bars show variability across three seeds and the

x

-axis is logarithmic.

Method	Fashion-MNIST		CIFAR-10		ImageNet-100
	$\downarrow$ FID	$\downarrow$ KLD	$\downarrow$ FID	$\downarrow$ KLD	$\downarrow$ FID
LoRA ( $r{=}8$ , $\alpha{=}16$ )	$68.4$ $(\pm\,11.1)$	$4.508$ $(\pm\,1.956)$	$50.9$ $(\pm\,3.4)$	$1.006$ $(\pm\,0.114)$	$109.8$ $(\pm\,0.5)$
LoRA ( $r{=}16$ , $\alpha{=}32$ )	$67.3$ $(\pm\,10.9)$	$5.129$ $(\pm\,2.154)$	$50.2$ $(\pm\,3.2)$	$1.055$ $(\pm\,0.123)$	$107.6$ $(\pm\,0.9)$
LoRA ( $r{=}32$ , $\alpha{=}64$ )	$66.2$ $(\pm\,11.2)$	$5.107$ $(\pm\,2.146)$	$50.0$ $(\pm\,3.0)$	$1.069$ $(\pm\,0.115)$	$104.6$ $(\pm\,2.3)$
GFT (sigmoid $\beta$ schedule)	$65.1$ $(\pm\,12.0)$	$4.264$ $(\pm\,1.309)$	$50.7$ $(\pm\,4.0)$	$1.138$ $(\pm\,0.171)$	$119.5$ $(\pm\,2.7)$
GFT (linear $\beta$ schedule)	$65.8$ $(\pm\,12.8)$	$4.873$ $(\pm\,2.188)$	$51.0$ $(\pm\,4.3)$	$1.101$ $(\pm\,0.176)$	$119.1$ $(\pm\,3.4)$

Appendix

Table 3: Memory-free baselines for class-incremental image generation. Final-stream means with standard errors over three seeds after five tasks (Fashion-MNIST and CIFAR-10) or three tasks (ImageNet-100). LoRA uses

\alpha=2r

; GFT decays the teacher-regularization weight

\beta

from 1 to 0. Blue rows show trajectory-replay references (adaptive or equal-task

\lambda

; MF-

N

uses

N

replay steps). Lower is better.

Figure 18: Fashion-MNIST samples after the final update. Held-out real images and samples from generative and trajectory replay, shown for seeds ( a ) 42, ( b ) 69, and ( c ) 1714. Columns correspond to the tasks of that seed’s class stream, with the task index and its two classes given above each block. Rows within a block share the same task.

Figure 19: CIFAR-10 samples after the final update. Held-out real images and samples from generative and trajectory replay, shown for seeds ( a ) 42, ( b ) 69, and ( c ) 1714. Columns correspond to the tasks of that seed’s class stream, with the task index and its two classes given above each block. Rows within a block share the same task.

Figure 20: ImageNet-100 samples after the final update. Held-out real images and samples from generative and trajectory replay, grouped by the three tasks of the stream. Within each group, rows share the same source and columns show three classes drawn from that task.

Figure 21: Synthetic perturbation-shift generalization. ( a ) Layout of the synthetic data, where color encodes the context

c

and marker shape the perturbation

y

. The held-out combination

(c^{\star},y^{\star})

is marked separately. ( b ) Energy distance on the held-out combination as a function of the replay weight

\lambda

for generative and trajectory replay, with the dashed vertical line at

\lambda=0.5

and the naive fine-tuning value indicated at the top right. Lines show the mean and shaded bands the variability across seeds, and lower is better. ( c ) Samples generated by trajectory replay for the held-out combination (blue) against the observed held-out cells (grey), at

\lambda\in\{0.05,0.5,0.95\}

. The two clouds coincide at

\lambda=0.5

, whereas both extremes displace the generated population away from its target, toward the incoming context at low

\lambda

and toward the historical one at high

\lambda

.

Data
Setting	Replogle-Nadig	PBMC
Cells	$643{,}413$	$9{,}697{,}974$
Genes (highly variable)	$2{,}000$	$2{,}000$
Contexts	$4$ cell lines	$12$ donors
Perturbations	$2{,}023$ genes	$90$ cytokines
Control	non-targeting	PBS

Appendix

Table 4: Configuration for the continual transcriptomics experiments.

Dataset	PCA width	$\uparrow$ Pilot fit $\Delta$
(a) PCA width sweep
Replogle	$50$	$\mathbf{0.486}$
Replogle	$100$	$0.458$
Replogle	$256$	$0.423$
Replogle	$512$	$0.389$
PBMC	$50$	$0.353$

Appendix

Table 5: Transcriptomics architecture selection. ( a ) PCA width sweep with a fixed MLP

1024{\times}4

backbone, reporting fit Pearson

\Delta

from single-seed pilots. ( b ) Backbone width and depth under the final sequential protocol with

500

generated cells per non-control stratum, reporting fit, retention, and task-0 held-out Pearson

\Delta

. All rows use seed

42

, and blue rows mark the selected configuration.

Method	Pearson $\Delta$			DE Overlap@N
	$\uparrow$ Fit	$\uparrow$ Retention	$\uparrow$ Held-out	$\uparrow$ Fit	$\uparrow$ Retention	$\uparrow$ Held-out
Naive (vf ${=}0$ , pert ${=}0$ , ctx ${=}0$ )	$0.193$ $(\pm\,0.004)$	$0.629$ $(\pm\,0.001)$	$0.163$ $(\pm\,0.004)$	$0.172$ $(\pm\,0.001)$	$0.287$ $(\pm\,0.001)$	$0.161$ $(\pm\,0.001)$
Naive (vf ${=}0$ , pert ${=}1$ , ctx ${=}1$ )	$0.359$ $(\pm\,0.003)$	$0.595$ $(\pm\,0.001)$	$0.169$ $(\pm\,0.012)$	$0.213$ $(\pm\,0.000)$	$0.275$ $(\pm\,0.000)$	$0.165$ $(\pm\,0.002)$
Naive (vf ${=}1$ , pert ${=}0$ , ctx ${=}0$ )	$0.616$ $(\pm\,0.001)$	$0.216$ $(\pm\,0.006)$	$0.351$ $(\pm\,0.002)$	$0.260$ $(\pm\,0.001)$	$0.156$ $(\pm\,0.001)$	$0.198$ $(\pm\,0.001)$
Naive (vf ${=}1$ , pert ${=}0$ , ctx ${=}1$ )	$0.616$ $(\pm\,0.001)$	$0.216$ $(\pm\,0.006)$	$0.351$ $(\pm\,0.002)$	$0.260$ $(\pm\,0.001)$	$0.156$ $(\pm\,0.001)$	$0.198$ $(\pm\,0.001)$
Naive (vf ${=}1$ , pert ${=}1$ , ctx ${=}0$ )	$0.616$ $(\pm\,0.001)$	$0.214$ $(\pm\,0.007)$	$0.353$ $(\pm\,0.002)$	$0.260$ $(\pm\,0.001)$	$0.156$ $(\pm\,0.001)$	$0.198$ $(\pm\,0.000)$

Appendix

Table 6: Continual perturbation-response modeling on Replogle-Nadig. Final-checkpoint Pearson

\Delta

and DE Overlap@N after adaptation from three task-0 cell lines to HepG2, reported as fit on the later task, retention on task 0, and held-out on perturbations excluded from training. Means with standard errors over seeds 42, 69, and 1714. Adaptive

\lambda

is selected on validation strata and equal-task uses

\lambda_{k}=k/(k+1)

. Gray marks the privileged real-replay reference and blue marks trajectory replay. Higher is better.

Method	Pearson $\Delta$			DE Overlap@N
	$\uparrow$ Fit	$\uparrow$ Retention	$\uparrow$ Held-out	$\uparrow$ Fit	$\uparrow$ Retention	$\uparrow$ Held-out
Naive (vf ${=}0$ , pert ${=}0$ , ctx ${=}0$ )	$0.233$ $(\pm\,0.003)$	$0.491$ $(\pm\,0.001)$	$0.235$ $(\pm\,0.004)$	$0.100$ $(\pm\,0.004)$	$0.144$ $(\pm\,0.001)$	$0.104$ $(\pm\,0.003)$
Naive (vf ${=}0$ , pert ${=}1$ , ctx ${=}1$ )	$0.232$ $(\pm\,0.005)$	$0.400$ $(\pm\,0.003)$	$0.207$ $(\pm\,0.002)$	$0.112$ $(\pm\,0.002)$	$0.133$ $(\pm\,0.001)$	$0.096$ $(\pm\,0.002)$
Naive (vf ${=}1$ , pert ${=}0$ , ctx ${=}0$ )	$0.290$ $(\pm\,0.000)$	$0.056$ $(\pm\,0.000)$	$0.151$ $(\pm\,0.001)$	$0.114$ $(\pm\,0.001)$	$0.074$ $(\pm\,0.001)$	$0.087$ $(\pm\,0.002)$
Naive (vf ${=}1$ , pert ${=}0$ , ctx ${=}1$ )	$0.290$ $(\pm\,0.000)$	$0.056$ $(\pm\,0.000)$	$0.151$ $(\pm\,0.001)$	$0.114$ $(\pm\,0.001)$	$0.074$ $(\pm\,0.001)$	$0.087$ $(\pm\,0.002)$
Naive (vf ${=}1$ , pert ${=}1$ , ctx ${=}0$ )	$0.289$ $(\pm\,0.001)$	$0.064$ $(\pm\,0.001)$	$0.154$ $(\pm\,0.002)$	$0.114$ $(\pm\,0.002)$	$0.078$ $(\pm\,0.001)$	$0.092$ $(\pm\,0.001)$

Appendix

Table 7: Continual perturbation-response modeling on Parse-PBMC. Final-checkpoint Pearson

\Delta

and DE Overlap@N after four sequential donor updates, reported as fit on the later tasks, retention on task 0, and held-out on perturbations excluded from training. Means with standard errors over seeds 42, 69, and 1714. Adaptive

\lambda

is selected on validation strata and equal-task uses

\lambda_{k}=k/(k+1)

. Gray marks the privileged real-replay reference and blue marks trajectory replay. Higher is better.

Figure 22: Fit–retention trade-off under alternative metrics. Retention against fit for every configuration on Replogle-Nadig (left) and Parse-PBMC (right), measured by L1 discrimination (top) and residual Pearson

\Delta

(bottom). Marker shape denotes the method family and marker color the heldout score for the same metric, with a separate color scale per panel. Higher is better on all axes.

Figure 23: Parse-PBMC donor stream. ( a ) Held-out Pearson

\Delta

after each donor update for naive fine-tuning, real replay, and generative and trajectory replay with an ODE backend. Filled markers with solid lines denote adaptive

\lambda

and open markers with dashed lines equal-task

\lambda

, and shaded bands show variability across seeds. ( b ) Mean replay weight

\lambda_{k}

selected at each update by the adaptive policy for the same replay methods, with the equal-task rule

\lambda_{k}=k/(k+1)

as a dashed reference. Error bars show variability over seeds

42

,

69

, and

1714

. Higher is better in ( a ).

Figure 24: Replay budget on the transcriptomics streams. Final held-out Pearson

\Delta

relative to the matched ODE run as a function of replay-field evaluations per update, for Replogle-Nadig (left) and Parse-PBMC (right), with the MF-1, MF-4, and ODE backends labeled by their evaluation count. Colors distinguish generative from trajectory replay, filled markers with solid lines denote adaptive

\lambda

and open markers with dashed lines equal-task

\lambda

, and the dotted line marks parity with the matched ODE run. Error bars show variability over seeds

42

,

69

, and

1714

. Higher is better.

Figure 25: Parse-PBMC performance by cell type. Pearson

\Delta

for each annotated cell type under naive full fine-tuning, real replay, and generative and trajectory replay with an ODE backend, all using equal-task

\lambda

. Each cell averages over the donor-cytokine strata of that cell type and over seeds 42, 69, and 1714. Higher is better.

Figure 26: Decoder reconstruction quality. Joint UMAP of observed cells and their reconstructions from the observed PCA latents through the frozen affine decoder, for Replogle-Nadig HepG2 test cells with PCA-50 (left) and Parse-PBMC cells from all donors with PCA-100 (right). Panel titles report the mean per-cell reconstruction correlation. Each panel embeds observed and reconstructed cells together, so no flow-matching generation is involved.

Data
Setting	Value
Cell line / experiments	U2OS, $3$ (U2OS-01/02/03)
Input	$6\times 96\times 96$ , uint8 $\rightarrow[-1,1]$
Train augmentation	uniform dequantization $1/128$ , random H/V flips
Perturbations	$1{,}071$ total, $100$ evaluated
Validation split	$10\%$ of each experiment’s train rows

Appendix

Table 8: Configuration for the continual microscopy experiments on RxRx1.

Method	Per-experiment FID $\downarrow$			Pooled stream $\downarrow$
	U2OS-01	U2OS-02	U2OS-03	FID	DINOv2 FDD
Joint training (all data; upper bound)	$29.7$ $(\pm\,0.2)$	$19.9$ $(\pm\,0.2)$	$26.6$ $(\pm\,0.2)$	$12.3$ $(\pm\,0.2)$	$66.4$ $(\pm\,1.4)$
Real replay (stored data)	$29.7$ $(\pm\,0.4)$	$19.6$ $(\pm\,0.2)$	$26.3$ $(\pm\,0.3)$	$12.2$ $(\pm\,0.3)$	$66.1$ $(\pm\,1.0)$
Naive fine-tuning	$\mathbf{36.0}$ $(\pm\,0.6)$	$83.0$ $(\pm\,1.6)$	$\mathbf{26.0}$ $(\pm\,0.1)$	$24.5$ $(\pm\,0.3)$	$110.1$ $(\pm\,1.3)$
Generative distillation (diffusion) ( Masip et al., 2025 )	$80.9$ $(\pm\,6.9)$	$82.1$ $(\pm\,2.6)$	$42.0$ $(\pm\,0.3)$	$46.4$ $(\pm\,3.6)$	$218.2$ $(\pm\,13.4)$
Generative replay (ODE, adaptive $\lambda$ )	$71.4$ $(\pm\,0.3)$	$27.1$ $(\pm\,0.2)$	$27.0$ $(\pm\,0.2)$	$22.0$ $(\pm\,0.1)$	$119.0$ $(\pm\,1.6)$

Appendix

Table 9: Continual phenomic generation on RxRx1. Final-checkpoint performance after the U2OS-01

\rightarrow

U2OS-02

\rightarrow

U2OS-03 experiment stream. Per-experiment FID compares composition-matched generated and test images within each biological experiment and stream metrics pool all three experiments, with FDD denoting Fréchet distance in DINOv2 ( Oquab et al., 2024 ) feature space. Means with standard errors over seeds 42, 69, and 1714. Adaptive

\lambda

is selected using validation FID and equal-task uses

\lambda_{k}=k/(k+1)

. Gray marks the privileged references that retain real historical data and blue marks trajectory replay. Bold marks the best memory-free mean in each column. Lower is better.

Figure 27: RxRx1 experiment stream. ( a ) FID on U2OS-01 and ( b ) pooled stream FID as a function of tasks completed, for naive fine-tuning, diffusion distillation, and generative and trajectory replay under adaptive and equal-task

\lambda

. The pooled reference set in ( b ) grows with each task, as indicated below the tick labels, whereas ( a ) keeps U2OS-01 fixed. Lines show the mean and shaded bands the variability over seeds 42, 69, and 1714. Lower is better.

Figure 28: Replay-weight and backend ablations on RxRx1. ( a ) Final pooled stream FID under adaptive and equal-task

\lambda

for generative and trajectory replay (left), with filled markers denoting adaptive

\lambda

and open markers equal-task, and the mean replay weight

\lambda_{k}

selected at each update (right), with the equal-task rule

\lambda_{k}=k/(k+1)

as a dashed reference. ( b ) Final pooled stream FID in Inception features (left) and final pooled stream Fréchet DINOv2 distance (FDD) in DINOv2 features (right) for trajectory replay with the MF-1, MF-4, and ODE backends, labeled by their replay-field evaluations per update. Error bars show variability over seeds 42, 69, and 1714. Lower is better for both metrics.

Figure 29: RxRx1 samples after the final update. Real test images and trajectory-replay samples for the perturbations TRNT1, NAA50, and EPPK1, shown for each biological experiment of the stream. Each block contains nine images from the corresponding perturbation and experiment, rendered through the fixed six-channel-to-RGB projection.

Method	Fashion-MNIST		CIFAR-10		ImageNet-100
	$\downarrow$ FID	$\downarrow$ KLD	$\downarrow$ FID	$\downarrow$ KLD	$\downarrow$ FID
Joint training (all data; upper bound)	$12.5$ $(\pm\,0.5)$	$0.040$ $(\pm\,0.010)$	$25.9$ $(\pm\,0.1)$	$0.100$ $(\pm\,0.020)$	$50.7$ $(\pm\,0.7)$
Diffusion naive fine-tuning	$67.7$ $(\pm\,6.6)$	$4.140$ $(\pm\,1.400)$	$53.6$ $(\pm\,4.4)$	$1.100$ $(\pm\,0.140)$	$105.8$ $(\pm\,3.1)$
Generative replay (diffusion) ( Masip et al., 2025 )	$184.9$ $(\pm\,12.6)$	$3.680$ $(\pm\,1.320)$	$84.3$ $(\pm\,9.1)$	$0.590$ $(\pm\,0.110)$	$106.3$ $(\pm\,2.5)$
Generative distillation (diffusion) ( Masip et al., 2025 )	$25.5$ $(\pm\,1.1)$	$0.290$ $(\pm\,0.040)$	$38.5$ $(\pm\,4.4)$	$0.230$ $(\pm\,0.080)$	$85.5$ $(\pm\,3.4)$
Flow matching naive fine-tuning	$64.6$ $(\pm\,11.4)$	$5.427$ $(\pm\,2.446)$	$51.2$ $(\pm\,4.0)$	$1.189$ $(\pm\,0.172)$	$103.4$ $(\pm\,2.6)$

Trajectory Replay: Continuous-Time Mixture Transport for Lifelong Generative Models

Abstract

Figures & tables

Appendix

Generative Modeling

World Action Models Enable Continual Imitation Learning with Recurrent Generative Replays

ARROW: Augmented Replay for RObust World models

Reliable Replay through Spatial Coherence in Online Continual Learning

Abstract

Generative Modeling

Explore similar work

World Action Models Enable Continual Imitation Learning with Recurrent Generative Replays

ARROW: Augmented Replay for RObust World models

Reliable Replay through Spatial Coherence in Online Continual Learning