CaLMFlow: Volterra Flow Matching using Causal Language Models

## Abstract

Abstract We introduce CaLMFlow (Causal Language Models for Flow Matching), a novel framework that casts flow matching as a Volterra integral equation (VIE), leveraging the power of large language models (LLMs) for continuous data generation. CaLMFlow enables the direct application of LLMs to learn complex flows by formulating flow matching as a sequence modeling task, bridging discrete language modeling and continuous generative modeling. Our method implements tokenization across space and time, thereby solving a VIE over these domains. This approach enables efficient handling of high-dimensional data and outperforms ODE solver-dependent methods like conditional flow matching (CFM). We demonstrate CaLMFlow’s effectiveness on synthetic and real-world data, including single-cell perturbation response prediction, showcasing its ability to incorporate textual context and generalize to unseen conditions. Our results highlight LLM-driven flow matching as a promising paradigm in generative modeling, offering improved scalability, flexibility, and context-awareness.

## 1 Introduction

Recent advances in deep learning have revolutionized generative modeling for complex, high-dimensional data. In particular, methods based on ordinary differential equations (ODEs), such as continuous normalizing flows (CNFs) (Chen et al., ) and flow matching (Lipman et al., ), have emerged as efficient tools for modeling continuous data distributions. However, many ODE systems suffer from stiffness making them numerically unstable and computationally expensive to solve accurately (Kushnir & Rokhlin, ; Zappala et al., ). In contrast, integral equations (IEs) offer a more generalized framework for capturing dynamics, with IE solvers demonstrating greater numerical stability than their ODE counterparts (Kushnir & Rokhlin, ; Zappala et al., ). Recent work in operator learning (Xiong et al., ; Cao, ; Zappala et al., ) has also connected solving integral equations with transformers, the foundational architecture of large language models (LLMs), inspiring the use of LLMs to model dynamical systems through the lens of IEs.

Building on these insights, our work introduces Causal Language Models for Flow Matching (CaLMFlow), a novel approach that models flow matching using Volterra integral equations (Zappala et al., ), enabling learning flows in a more robust manner. By leveraging causal language models (CLMs) to solve Volterra integral equations, our method capitalizes on the ability of CLMs to comprehend natural language, allowing for the modeling of complex data distributions conditioned on natural language prompts. Our approach provides a robust framework for applications ranging from synthetic data generation to complex system modeling in the biological sciences.

Our key contributions are:

- •
Flow matching using causal language models: We introduce a novel framework formulating flow matching as Volterra integral equations and leveraging causal language models (CLMs) to approximate the solutions, enhancing stability and performance of flow matching.
- •
Controllable generation of flows using natural language:
We present a flexible and effective approach to controllable generation conditioned on textual prompts by leveraging the causal language model’s natural language understanding. In Section 5.2, we demonstrate its superiority by applying it to perturbation response prediction in single-cell data, outperforming traditional flow matching and popular single-cell methods.
- •
Continuous space tokens via variational decoding: We introduce variational decoding to sample and generate continuous data. This approach models a continuous conditional distribution for next-token sampling, extending language modeling techniques to continuous domains. Our ablation study in Section 6.1 demonstrate that variational decoding is crucial for accurately modeling continuous data within our framework.
- •
Spatiotemporal and trajectory tokenization: We present a spatiotemporal tokenization scheme that enables CaLMFlow to model VIEs across both spatial and temporal domains. Additionally, by modeling multiple flows concurrently, CaLMFlow captures correlations between data samples. We demonstrate in Sections 5.1 and 6.3 this approach significantly improves performance.

Figure: Figure 1: Overview of the CaLMFlow framework. CaLMFlow takes as input textual conditions and flows and generates the next time point for the flows. The textual condition is tokenized and embedded using the LLM tokenizer and embedding layer while the conditional flows are transformed into spatial-temporal tokens using a learned projection. If multiple conditional flows are input simultaneously, the tokens are ordered by flow, space, and then time. The LLM applies causal language modeling and generates the next time point for each flow.
Refer to caption: extracted/5897793/images/overview_figure/overview_with_integral_z.png

## 2 Related Work

Flow Matching and Continuous Normalizing Flows: Flow matching has significantly enhanced the efficiency and scope of continuous normalizing flows (CNFs) (Chen et al., ; Papamakarios et al., ). Conditional Flow Matching (CFM) (Lipman et al., ) allows for precise control over the generative process by optimizing conditional vector fields tailored for specific distribution paths, including those based on Gaussian distributions. Tong et al. () generalized the conditional paths and introduced mini-batch optimal transport and Schrödinger bridge CFM, improving the efficiency and performance of CFM models. In Hu et al. (), flow matching is applied to text generation in a non-autoregressive manner, showing improvements compared to other diffusion-based text generation models such as DiffSeq (Gong et al., ). Our work, however, is primarily concerned with adapting LLMs to generate continuous data conditioned on text.

Text-conditional Generation: Text-conditional image generation has made significant strides through the integration of diffusion models (Sohl-Dickstein et al., ; Ho et al., ; Song et al., ) and large language models (LLMs). State-of-the-art systems like Stable Diffusion (Rombach et al., ), DALLE-2 (Ramesh et al., ), and DINOv2 (Oquab et al., ) leverage LLM embeddings to generate high-quality images from textual descriptions. Recent research (Ding et al., ; Yu et al., ; Ge et al., ; Zhan et al., ) has focused on adapting LLMs for multimodal generation, often employing vector quantization (van den Oord et al., ; Razavi et al., ; Ge et al., ) to extend LLM vocabularies with latent tokens representing non-textual data. Our CaLMFlow method introduces a novel approach as the first flow matching-based text-conditional generative model that produces continuous tokens, potentially offering greater flexibility and expressiveness compared to discrete token-based methods.

Integral Equations in Neural Network Frameworks: The fusion of neural networks and differential equations was introduced by Chen et al. () with neural ordinary differential equations. This concept has since been extended to integral equations, particularly Volterra equations, in several studies (Fu & Hirsa, ; Zappala et al., ; ). Notably, Zappala et al. () exploited the relationship between attention mechanisms and integral kernels (Xiong et al., ; Cao, ) to model integral equations using transformers. Concurrently, approaches like physics-informed neural networks (PINNs) (Raissi et al., ; Lu et al., ; Goswami et al., ) have emerged, incorporating physical laws as prior knowledge to enhance model accuracy and generalization in operator learning tasks. Our work builds upon these advancements, extending the application of Volterra integral equations to the flow matching framework.

## 3 Volterra Flow Matching

### 3.1 Flow Matching as Volterra Integral Equations

Flow matching (Lipman et al., ) is typically formulated as learning the time-dependent vector field $v(x,t)$ generating the flow $\phi(x,t)$ related by the ordinary differential equation (ODE)

$$ $\frac{d\phi}{dt}=v(\phi,t),\qquad\phi(x,0)=x,$ (1) $$

transforming an initial distribution $p_{0}$ into a target distribution $p_{1}$ through the application of a numerical solver applied to the learned vector field.

However, ODEs can suffer from stiffness, especially when modeling systems with rapid changes or long-range dependencies (Rokhlin, ; ; Kushnir & Rokhlin, ; Zappala et al., ) and, as a result, solving such systems are highly numerically unstable and computationally expensive. To address these challenges, we reformulate the flow matching problem using Volterra integral equations (VIEs) of the second kind:

$$ $z(t)=z(0)+\int_{0}^{t}G(z(s),t,s)ds$ (2) $$

where $G(z(s),t,s)$ is a Urysohn kernel function encoding the influence of past states on the current state. VIEs generalize ODEs, as shown by the transformation of Equation (1) into the equivalent integral form:

$$ $\phi(t)=\phi_{0}+\int_{0}^{t}v(\phi(s),s)ds.$ (3) $$

This VIE formulation offers several advantages: it inherently accounts for nonlocal components in the dynamics, is more flexible and robust for modeling complex systems with memory effects, and avoids issues like stiffness and underflowing that ODE solvers encounter (Kushnir & Rokhlin, ; Zappala et al., ). Consequently, Volterra flow matching provides a more general and stable approach to modeling continuous flows between distributions.

### 3.2 Solving Volterra Integral Equations with Causal Language Models

In CaLMFlow, we define the flow using a more general form of Equation 3:

$$ $z_{t}=f(z_{t},t)+\int_{0}^{t}G(z_{s},t,s)ds,$ (4) $$

where 1) $z_{t}$ (also referred to as $z(t)$) is the state at time $t$, 2) $f(z_{t},t)$ is the inhomogeneous term ($z_{0}$ in the case of Equation 3), and 3) The integral term $\int_{0}^{t}G(z_{s},t,s)\,ds$ captures the accumulated influence of past states.

We discretize and model the entire integral equation using a causal language model (CLM) following Xiong et al. (), Cao (), and Zappala et al. (). The time interval $[0,1]$ is first discretized into $N$ time steps $t_{0},t_{1},\ldots,t_{N}$, resulting in a sequence $\left(z_{t_{0}},z_{t_{1}},\dots,z_{t_{N}}\right)$. We can define a sequence of functions $\left(z^{0}(t),z^{1}(t),\ldots,z^{N}(t)\right)$ where $z^{i}(t)$ is defined on $[0,t_{i}]$ and $z^{i}(t_{j})=z_{t_{j}}$. The CLM then acts as an iterative integral equation solver on the discretized VIE by predicting $z^{i+1}$ from $z^{i}$:

$$ $\hat{z}^{i+1}=f(z^{i},t_{i+1})+\sum_{j=0}^{i}\Delta t_{i+1}G(z_{j},t_{i+1},t_{ j}),$ (5) $$

where $\Delta t_{k}=t_{k}-t_{k-1}$. Further theoretical discussion in the language of Banach spaces can be found in Appendix C, and a generalization to higher dimensional integrals where space is also included is straightforward.

Similar to the discussion in Lipman et al. (), the naïve Volterra flow matching objective

$$ $\mathcal{L}_{\text{VFM}}=\mathbb{E}_{p(z^{N})}\left\|z^{N}-\hat{z}^{N}\right\| ^{2},$ (6) $$

where $\hat{z}^{N}=(\hat{z}_{t_{0}},\hat{z}_{t_{1}},\ldots,\hat{z}_{t_{N}})$ is the sequence of states predicted from the ground truth marginal trajectories (Equation 5) and $p(z^{N})$ is the distribution of trajectories over the flow, is intractable due to the inaccessibility of ground truth marginal trajectories $z^{N}$. Instead, we use the optimal transport conditional paths (Lipman et al., ; Tong et al., ) $z^{N}_{z_{0},z_{N}}(t):=(1-t)z_{0}+tz_{N}$ and optimize the conditional Volterra flow matching objective

$$ $\mathcal{L}_{\text{CVFM}}=\mathbb{E}_{p_{0}(z_{0}),q(z_{N})}\left\|z^{N}_{z_{0 },z_{N}}-\hat{z}^{N}\right\|^{2},$ (7) $$

where $p_{0}$ is the initial source distribution (e.g., a Gaussian) and $q$ is the target data distribution. A theoretical discussion of the different objective functions can be found in Appendix D. We note that next-token prediction in CLMs is simulation-free during training, as the next token is predicted based on the ground truth history. Full trajectory simulation occurs only during inference. As such, CaLMFlow, like CFM, operates as a simulation-free approach during training.

### 3.3 Continuous Space Tokens via Variational Decoding

We introduce variational decoding with a Kullback-Leibler divergence regularizer in our experiments. Unlike standard CLMs, which rely on a fixed vocabulary and model the next-token distribution through softmax probabilities, our approach enables the CLM to model a continuous distribution of next tokens. Specifically, we use a probabilistic encoder $q_{\phi}(\mathbf{z}|\mathbf{x})$ to map the CLM output tokens $\mathbf{x}$ to a posterior latent distribution $\mathcal{N}(\mathbf{z};\bm{\mu},\bm{\sigma}^{2}\mathbf{I})$, and a probabilistic decoder $p_{\psi}(\mathbf{x}|\mathbf{z})$ to reconstruct the tokens $\mathbf{x}$, where latent variable $\mathbf{z}$ acts as a continuous representation in a latent space.

Both $q_{\phi}$ and $p_{\psi}$ are optimized by maximizing the evidence lower bound (ELBO) (Kingma & Welling, ):

$$ $-\mathcal{L}_{\text{VAE}}=\mathbb{E}_{q_{\phi}(\mathbf{z}|\mathbf{x})}[\log p_ {\psi}(\mathbf{x}|\mathbf{z})]-\beta\,\text{KL}(q_{\phi}(\mathbf{z}|\mathbf{x} )\|p(\mathbf{z})),$ (8) $$

where $p(\mathbf{z}):=\mathcal{N}(\mathbf{z};\mathbf{0},\mathbf{I})$ is a prior over the latent variable $\mathbf{z}$, $\beta$ is a scaling hyperparameter as introduced in (Higgins et al., ), and $\text{KL}(q\|p)$ is the Kullback-Leibler divergence. Using the ELBO, our Volterra flow matching objective becomes

$$ $\mathcal{L}_{\text{VCVFM}}=\mathcal{L}_{\text{CVFM}}+\beta\,\text{KL}(q_{\phi} (\mathbf{z}|\mathbf{x})\|p(\mathbf{z}))$ (9) $$

To control the generation of continuous tokens at inference, we use a temperature parameter $\tau$ that scales the variance of the encoded posterior, analogous to the use of temperature in LLMs (Renze & Guven, ; Peeperkorn et al., ). Specifically, the posterior latent distribution is modified as $q_{\phi}(\mathbf{z}|\mathbf{x})\sim\mathcal{N}(\mathbf{z};\bm{\mu},\tau\bm{
\sigma}^{2}\mathbf{I})$, where $\tau\geq 0$ is the temperature parameter. An ablation experiment for the temperature parameter is in Section 6.1.

## 4 Spatiotemporal and Multi-trajectory Tokenization

In this section, we introduce our spatiotemporal and multi-trajectory tokenization method depicted in Figure 1 to effectively represent data for CLM training.

### 4.1 Spatiotemporal Tokenization

We discretize the temporal domain $[0,1]$ using a fixed, evenly spaced grid $\{t_{i}\}_{i=1}^{N}$ of length $N$.

The discretized temporal dynamics are encoded as input matrix $\mathbf{X}\in\mathbb{R}^{T\times D_{in}}$ comprising T tokens, where each token $\mathbf{x}_{i}$ is of $D_{in}$ dimensions.

To introduce spatial structure into the representation, we further subdivide each temporal token $\mathbf{x}_{i}$ into $K$ tokens $\{\mathbf{x}_{ij}\}_{j=1}^{K}$. The splitting can be either learned or predefined, depending on the data. For instance, for single-cell data, we map the encoded gene expression features with a learned network $S_{\theta}:\mathbb{R}^{D_{in}}\to\mathbb{R}^{K\cdot D_{in}}$ and split the token to $K$ spatial tokens. In the case of images, we use an encoder to extract low-level features, resulting in a feature map that is split into tokens.

The final spatiotemporal token sequence $\mathbf{X}_{st}\in\mathbb{R}^{K\cdot D_{in}}$ is represented as

$$ $\displaystyle\mathbf{X}_{st}:=\left[\mathbf{x}_{11},\dots,\mathbf{x}_{1K}, \dots,\mathbf{x}_{N1},\dots,\mathbf{x}_{NK}\right]^{\top}.$ (10) $$

This spatiotemporal sequence of tokens is processed by the CLM, which approximates a Volterra integral equation over both space and time. The discussion in Section 3.2 extends to the spatiotemporal VIE:

$$ $z(x,t)=f(z,x,t)+\int_{0}^{t}\int_{\Omega}G(z,x,x^{\prime},t,s)dx^{\prime}ds.$ (11) $$

### 4.2 Multi-trajectory Tokenization

We tokenize more than one conditional trajectory as input to the CLM by sampling $M$ spatiotemporal sequences of tokens as above. The overall tokenized sequence $\mathbf{X}_{sdt}\in\mathbb{R}^{M\cdot K\cdot D_{in}}$ becomes

$$ $\displaystyle\mathbf{X}_{sdt}:=\left[\mathbf{x}_{111},\dots,\mathbf{x}_{1K1}, \dots,\mathbf{x}_{1KM},\dots,\mathbf{x}_{N11},\dots,\mathbf{x}_{NKM}\right]^{ \top}.$ (12) $$

We empirically find that providing information across a batch of trajectories as context benefits model performance. We observe that such benefits are unique to CLMs, whereas methods like CFM cannot natively model multiple trajectories. While multi-trajectory training and inference are related to integration over function spaces, exploring this connection is beyond the scope of our work and we leave it for future exploration.

## 5 Experiments

We evaluate CaLMFlow on synthetic datasets to showcase the advantages of integral equations for modeling high-dimensional dynamical systems, and apply it to single-cell data generation, demonstrating its ability to model complex distributions and leverage natural language understanding.

### 5.1 Synthetic Datasets

#### 5.1.1 High Dimensional Data

Stiffness is a well-known challenge in the numerical integration of ODEs (Kushnir & Rokhlin, ; Zappala et al., ), particularly in systems with high dimensionality. Conversely, the embedding dimensions of causal language models (CLMs) are inherently large, as demonstrated by pretrained models like GPT-2, which has an embedding dimension of 768. As such, we hypothesize that CaLMFlow is better at modeling high dimensional data than ODE-based methods. To evaluate the robustness of CaLMFlow in high-dimensional settings, we compare its performance against traditional ODE-based methods, which typically degrade as dimensionality increases.

Our results, summarized in Table 1, demonstrate that while CFM breaks down at higher dimensions, CaLMFlow maintains strong performance. This suggests that CaLMFlow is an effective alternative to ODE-based approaches for modeling high-dimensional problems, providing stability and accuracy where methods like CFM fail.

Table: Table 1: Performance comparison of CaLMFlow and CFM variants across different distribution pairs (Gaussian $\to$ 2 Gaussians, Gaussian $\to$ 8 Gaussians, Gaussian $\to$ 2 Moons) and dimensions (100, 1000). We report 2-Wasserstein distance (2-Wass) and Maximum Mean Discrepancy (MMD) ($\mu\pm\sigma$ over 5 runs), with best results highlighted in bold. As the dimensionality increases, CaLMFlow consistently outperforms CFM variants, achieving lower 2-Wasserstein distances and MMD values, particularly in high-dimensional settings where we expect traditional ODE-based methods, such as CFM, struggle.

#### 5.1.2 Multi-trajectory Context

While approaches like CFM focus on modeling the flow of individual points, CaLMFlow is able to sample multiple trajectories and model them simultaneously. The results in Table 2 show this approach improves the performance of CaLMFlow on generating synthetic data. Visualizations of generated results, as shown in Figure 2, further demonstrate that modeling several trajectories at the same time allows CaLMFlow to distribute data more evenly and accurately by leveraging trajectory context.

Figure: (a) Ground Truth
Refer to caption: extracted/5897793/images/multi_trajectories_square/heat_ground_mt.png

Table: Table 2: Comparisons of CFM, CaLMFlow, and CaLMFlow with multiple trajectories on different distribution pairs in 100 dimensions. We report 2-Wasserstein distance (2-Wass) and Maximum Mean Discrepancy (MMD), averaged over 5 runs. Best results are highlighted in bold. The table shows that CaLMFlow, especially when utilizing multiple trajectories, significantly outperforms CFM in terms of both fit and distribution accuracy. This demonstrates the advantage of leveraging multi-trajectory modeling in CaLMFlow.

### 5.2 Single-cell Generation

We apply CaLMFlow to immune tissue single-cell expression data (Dong et al., ) to demonstrate its effectiveness in both unconditional and conditional generation of complex, high-dimensional real-world data. We utilize the first 1,000 principal components of the gene expression data as features. The dataset comprises annotations for 7 cell types, 10 perturbations, and 2 chronicities, leading to 140 unique combinatorial labels. In the unconditional generation experiment, the model generates the overall target distribution from Gaussian noise as initial conditions, regardless of labels. In the conditional generation experiment, five combinations of the labels are held out as a test set, and the models are tasked with generating the held-out target distribution conditioned on the unseen combinatorial labels.

Our method is benchmarked against several models, including CFM (Tong et al., ) and its variants CFM-OT and CFM-SB, the denoising diffusion probabilistic model (DDPM) (Ho et al., ), single-cell generative models scVI (Lopez et al., ) and scGPT (Cui et al., ), and the Compositional Perturbation Autoencoder (Lotfollahi et al., ), depending on the task. To assess the quality of data generated by each model, we compute distributional metrics, such as maximum mean discrepancy (MMD), 2-Wasserstein distance, and Kullback-Leibler Divergence (KLD), between model-generated data and the ground truth test data. For KLD , we use Leiden clustering to generate a distribution of points across clusters (see Appendix A.1.1 for details).

#### 5.2.1 Unconditional Generation of Single-cell Data

The results in Table 3, demonstrate that CaLMFlow consistently outperforms CFM and all other methods across all metrics. Furthermore, as illustrated in Figure 6, CaLMFlow generates cells with distributions more closely aligned to the ground truth data compared to other methods. These findings underscore CaLMFlow’s superior performance in capturing and reproducing the complex high-dimensional distributions inherent in single-cell expression data.

Table: Table 3: Unconditional single-cell generation results comparing generated data to ground truth data. To evaluate CaLMFlow’s ability to accurately generate its training single-cell dataset (Dong et al., ), we computed distributional metrics MMD, 2-Wasserstein, KLD, and adMMD (MMD with a $k$-NN-based adaptive kernel) across 5 seeds. Our default CaLMFlow outperforms all methods across all metrics including CFM and its variants CFM-OT and CFM-SB, demonstrating CaLMFlow’s ability to model the data distribution. Further improvement is seen with CaLMFlow (5 traj.), showing the benefit of multi-trajectory tokenization. scGPT’s Leiden KLD score is omitted due to the model’s poor performance on this metric being less informative for comparison purposes. See Figure 6 for a visual comparison of CaLMFlow and CFM. Experimental details are in Appendices A.1.1 and A.1.2

#### 5.2.2 Single-cell Perturbation Response Prediction

We leverage CLMs’ inherent capabilities to encode and comprehend natural language by representing perturbation conditions as simple text prompts (see A.1.3 for details). These prompts are prepended to the embedded flow-matching conditional trajectories and processed through the CLM’s tokenizer and embedding layers. For details on conditional encodings for other models, see A.1.3.

Our architecture is based on a customized Pythia (Biderman et al., ) model as the CLM. To investigate the benefit of natural language capabilities of CLMs, we compare random initialization to natural language pretraining.

The results shown in Table 4 highlight CaLMFlow’s ability to generate data distributions that closely align with the ground truth, outperforming competing models. The correlation statistics shown in Table 5 underscore CaLMFlow’s effectiveness in producing realistic cell expression profiles that correspond to the specified combinatorial labels. Notably, both tables show that leveraging pretrained CLM weights enhances CaLMFlow’s performance, showcasing the power of utilizing natural language understanding abilities of CLMs in the CaLMFlow framework.

Furthermore, as shown in Figure 3, both variants of CaLMFlow generate data that closely overlaps with the ground truth distribution, demonstrating CaLMFlow’s superior ability to model data under unseen conditions. Figure 7 illustrates the distributions produced by each model, showing that CaLMFlow’s generated data most accurately reflects the ground truth. In contrast, other models are either unable to differentiate between combinatorial labels or generate unrealistic distributions. These visualizations reinforce the high quality of data generated by CaLMFlow, emphasizing its capability to model complex distributions and effectively utilize natural language prompts.

Table: Table 4: Single-cell perturbation response prediction comparison in terms of fit ($\mu\pm\sigma$ over five repeated runs). CFM, scVI, scGPT and CPA are compared. Best results in bold. For CaLMFlow, R.I. stands for randomly initialized CLM and N.L. stands for natural language pretrained CLM. "—" in Leiden KLD represents infinite value due to not being able to generate all classes. ^∗: for CPA, no standard deviation is reported as it is deterministic; CPA generated data that resulted in numerical instability when computing MMD, leading to the absence of valid MMD values

Table: Table 5: Single-cell perturbation response prediction comparison in terms of correlation. $R^{2}$, Pearson $R$ and Spearman $R$ on the full cell and the top 100 most differentially expressed genes are reported ($\mu\pm\sigma$ over five repeated runs and five unique combinatorial labels). For CaLMFlow, R.I. stands for randomly initialized CLM and N.L. stands for natural language pretrained CLM.

Figure: (a) CaLMFlow (R.I.)
Refer to caption: extracted/5897793/images/conditional_single_cell/combined_umaps/CaLMFlow_RI_.png

## 6 Ablation Experiments

### 6.1 Temperature

To investigate the impact of the temperature parameter on CaLMFlow’s performance at inference, we varied the temperature for the 8-Gaussians to 2-Moons dataset. Figure 4 shows the best MMD and 2-Wasserstein values at $\tau=0.2$, where the generated data closely matches the ground truth. Deviations from this value lead to less accurate transformations. Interestingly, it has been empirically found that the optimal temperature in LLMs is often below 1.0 to mitigate inference noise, an observation that aligns with our findings. The experiment also highlights the importance of the VAE component, as its removal significantly degrades performance.

Figure: Figure 4: Ablation results on temperature. Left: CaLMFlow generated data from 8gaussians to 2moons, using different temperature values. Right: 2-Wasserstein and MMD performances as a function of temperature. The plots show that a low, non-zero temperature value ($\tau$=0.2) produces the best performance and that the VAE is necessary.
Refer to caption: extracted/5897793/images/temperature_ablation/temperature_scatter.png

### 6.2 Number of time points

To evaluate the impact of the number of time points on CaLMFlow’s performance, we varied the number of time points during training and inference for the 8 Gaussians to 2 Moons dataset. Figure 8 shows that as the number of time points increases, both the MMD and 2-Wasserstein consistently decrease, indicating improved model accuracy. This demonstrates that increasing the number of time points improves CaLMFlow’s ability to capture the transformation dynamics, leading to better performance.

### 6.3 Number of Spatiotemporal tokens and Trajectories

To test the impact of spatiotemporal tokenization, we use a similar setup as in Tong et al. () on the MNIST dataset (details in Subsection 4.1). All key hyperparameters, optimizers, and training configurations were kept identical to ensure consistency. Our results, as shown in Table 6, demonstrate that CaLMFlow outperforms other methods and increasing the number of spatial tokens improves inception scores.

Table: Table 6: Comparison of inception scores on the MNIST dataset. CaLMFlow outperforms other methods. The results show an improvement in inception scores as the number of space tokens increases, with CaLMFlow (8 Space tokens) achieving the highest score with statistical significance. A comparison of generated images between CFM, CFM-OT, CFM-SB, and CaLMFlow can be found in Figure 9

Similarly, to test the impact of multi-trajectory tokenization, we varied the number of trajectories used to transform data from 2-Moons to 8-Gaussians toy dataset. Table 7 shows a marked improvement in MMD and 2-Wasserstein with the addition of more trajectories.

**Table 7: Ablation results on multi-trajectory implementations of CaLMFlow transforming 2 moons into 8 Gaussians in 2 dimensions. We evaluate the MMD and 2-Wasserstein between the generated target distribution and ground truth distribution over 5 seeds. The results show increasing the number of trajectories improves CaLMFlow’s ability to accurately generate the target distribution.**
| Method $\downarrow$ Metric $\rightarrow$ | 2-Wass ($\downarrow$) | MMD ($\downarrow$) |
| --- | --- | --- |
| CaLMFlow (1 traj.) | 1.3756 $\pm$ 0.0668 | 0.0057 $\pm$ 0.0002 |
| CaLMFlow (4 traj.) | 1.0477 $\pm$ 0.0771 | 0.0041 $\pm$ 0.0003 |
| CaLMFlow (10 traj.) | 0.8753 $\pm$ 0.1243 | 0.0020 $\pm$ 0.0003 |

## 7 Conclusion and Future Work

We introduce CaLMFlow, a novel framework for flow matching that leverages causal language models by casting flow matching as a Volterra integral equation. CaLMFlow outperforms traditional flow matching models like CFM, especially on high-dimensional datasets, such as single-cell generation and perturbation response prediction. It generates more accurate and realistic data in both synthetic and real-world tasks. Future work will formalize CaLMFlow’s multi-trajectory approach using integral equations over function spaces and explore its potential as an iterative solver to refine entire trajectory outputs, enhancing its ability to model systems with complex global dynamics.

##### Reproducibility Statement

We have supported reproducibility by detailing our experimental implementations, metric computations, and data sources in the appendices. The main text and supplementary materials thoroughly explain our models and methods, along with how the metrics were computed.

## Appendix A Detailed Experiment Setup

All models are implemented in PyTorch (Paszke et al., ), trained using the Adam optimizer (Kingma & Ba, ). Additionally, FlashAttention-2 (Dao, ) is employed to accelerate training.

For the synthetic experiments, we utilize GPT-2 (Radford et al., ) as the CLM, while in the single-cell experiments, we adopt Pythia (Biderman et al., ). Both models are hosted on the open-source LLM platform Hugging Face (Wolf et al., ).

### A.1 Single Cell Generation

#### A.1.1 Details on Distribution and Correlation Metrics

##### MMD and 2-Wasserstein Distance

Distributional metrics, even with kernel methods, are shown to suffer greatly from curse of dimensionality as shown in previous work(Reddi et al., ; Arias-Castro et al., ). Therefore, rather than computing MMD and 2-Wasserstein in the original 1000-dimensional feature space, we first embed the ground truth and generated distribution together with UMAP and compute the metrics in the joint UMAP space. The dimension of the UMAP is chosen to be 10.

For MMD, we adopt a radial basis function (RBF) kernel $k:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$ with

$$ $k(x,y)=\exp\left(-\lVert x-y\rVert^{2}\right)$ $$

The estimator of MMD between two the generated distribution $P=(x_{1},...,x_{N})$ and the ground truth distribution $Q=(y_{1},...,y_{N})$ is computed as

$$ $\text{MMD}(P,Q)=\frac{1}{N(N-1)}\sum_{i\neq j}k(x_{i},x_{j})+\frac{1}{N(N-1)} \sum_{i\neq j}k(y_{i},y_{j})-\frac{2}{N^{2}}\sum_{i=1}^{N}\sum_{j=1}^{N}k(x_{i },y_{j})$ $$

The 2-Wassertein distance is computed as

$$ $W_{2}(P,Q)=\inf_{\pi}\sqrt{\left(\frac{1}{n}\sum_{i=1}^{n}\lVert x_{i}-y_{\pi( i)}\rVert^{2}\right)},$ $$

where $\pi$ are permutations on $n$ elements. This is implemented with the Python Optimal Transport toolbox (Flamary et al., ).

We compute adaptive MMD scores using an adaptive Radial Basis Function (RBF) kernel, where the bandwidth is dynamically adjusted according to the local density of points.

$$ $K(X,Y)=\exp\left(-\frac{\|X-Y\|^{2}}{2\sigma_{X}\sigma_{Y}}\right)$ $$

The bandwidth for each point is determined based on the distances to its $k$-nearest neighbors (k-NN), enabling the kernel to adapt to varying local densities. This adaptive mechanism enhances the kernel’s ability to capture distributional differences in non-homogeneous data, improving the sensitivity of the MMD to local structure.

##### Leiden KL Divergence

We employ the Leiden clustering algorithm (Traag et al., ) to identify community structures within the ground truth data. We run Leiden clustering with the resolution parameter set to 0.3, resulting in 8 clusters as shown in Figure 5(a). We then train a simple two-layer MLP with ReLU activation to predict the Leiden labels from the ground truth data, achieving a prediction accuracy of 95$\%$ on a held-out test set. This MLP classifier is subsequently used to assign predicted Leiden cluster labels to data generated by different models. To assess model performance, we compute the KL divergence between the distributions of Leiden clusters in the ground truth and generated data.

Figure: (a) Ground Truth
Refer to caption: extracted/5897793/images/conditional_single_cell/appdix_leiden/gt.png

##### Inception Score

For the conditional generation task, to assess the quality of the generated cells, we utilize the Leiden cluster labels and the trained MLP classifier to compute an inception score, inspired by its use in the computer vision domain (Salimans et al., ). Let $p(y|\mathbf{x})$ represent the conditional distribution of the MLP classifier assigning a label $y$ to a sample $\mathbf{x}$, the Inception Score is calculated as

$$ $\text{IS}=\exp{\left(\mathbb{E}_{x}\text{KL}(p(y|\mathbf{x})||p(y)\right)},$ $$

where $p(y)$ is the marginal distribution of the label $y$.

##### Correlation Metrics

For the conditional generation task, we additionally compute correlation metrics such as $R^{2}$, Pearson $R$, and Spearman $R$ between the average generated cells and the average ground truth cells for each condition, to assess how well the models generate cells based on the specified conditions. These metrics are calculated per combinatorial label and then averaged. For each condition (cell type, perturbation, and chronicity) we identify the top 100 most differentially expressed genes by comparing them to control cells of the same type and chronicity but without stimulation.

#### A.1.2 Single Cell Unconditional Generation Model Implementations

CaLMFlow’s architecture is a small Pythia model with 64 hidden dimensions, a 256-dimensional feedforward layer, 4 attention heads, and 2 blocks. The encoder for the cell expression vector, latent mean, latent sigma and the decoder were all 2-layer MLPs, with input and output dimensions adjusted for the number of spatial tokens. We used the standard setup for CFM following the torchCFM library—a 4-layer MLP with intermediate width 1024 for CFM, CFM-OT, and CFM-SB. At inference, we use 100 function evaluations using the adaptive "dopri5" solver in the NeuralODE package with the default parameters. For DDPM, we use a 2-layer MLP with 2048 intermediate dimension, and 100 denoising steps. We found similar performance with deeper and wider layers, so we used a smaller model to account for the number of parameters. For scVI, we trained on the full gene expression since its zero-inflated negative binomial decoder is not meant to handle non-sparse data. At inference, we randomly sample from the prior distribution to generation samples with the decoder. For scGPT, we trained the model with a 5000 gene context using the masked decoding decoding scheme from the official repository with masking ratios 0.25, 0.5, and 0.75. At inference, we mimicked the training procedure as closely as possible by using 2500 context genes and masking 2500 genes until an entire cell is generated. The first 2500 genes are taken from a randomly sampled ground truth cell expression vector.

#### A.1.3 Single Cell Conditional Encoding

##### CaLMFlow

We use the following templates to form the natural language prompts:

For perturbed cells, the prompts are:

"Generate a {cell type} cell stimulated with {perturbation} and exposure {chronicity}:"

For control cells, the prompts are:

"Generate a {cell type} cell with no stimulation and exposure {chronicity}:"

For example, one sample prompt for a perturbed cell is:

"Generate a CD4 T cell stimulated with IL-6 and exposure acute:"

and one sample prompt for a control cell is:

"Generate a B cell with no stimulation and exposure chronic:"

The prompts are tokenized using

AutoTokenizer.from_pretrained("EleutherAI/pythia-160m")

and embeded using the embedding layers of a customized PyThia-160M model. The embedding vectors are prepended to the corresponding conditional flow matching trajectories and processed by the CLM. Since we don’t train the CLM to generate text responses, no losses are applied on the text prompts during training.

During inference, we sample combinatorial labels from the test dataset, form prompts as above, sample initial Gaussian noises, tokenize and prompt the CLM to generate the trajectories and consequently the target distribution.

##### CFM

Conditional Flow Matching doesn’t provide an option for conditional generation out of the box, so adapted the torchCFM python library. In order to allow CFM to generate contionally, we create 3 one-hot vectors corresponding to each of the 3 conditions: cell type, perturbation, and chronicity. With 7 cell types, 10 perturbations, and 2 exposures, we obtain a 19-dimensional vector once concatenated together. This conditioning vector is then appended to the data vector, and CFM learns the vector field as usual. The NeuralODE python package was adapted to handle conditioning vectors so we could apply the same adaptive solver used in Tong et al. ().

##### scVI

To condition scVI, we applied the same one-hot encoding scheme as CFM to the latent vectors prior to decoding the posterior parameters.

##### scGPT

We trained scGPT (Cui et al., ) with a 5000 gene context using the masked decoding decoding scheme from the official repository with masking ratios 0.25, 0.5, and 0.75. The conditioning vectors were one-hot encoded, and we used learned embedding matrices for each type of condition and added the learned embedding to the cell expression. At inference, we applied the same procedure, always using 2500 context genes and masking 2500 genes until an entire cell is generated. The first 2500 genes are generated using a sampled ground truth cell expression vector from the test dataset, giving an advantage to scGPT other approaches did not have. Despite this, scGPT still showed poor performance relative to other methods.

##### CPA

CPA is a framework of perturbation response prediction. We follow their setup as in their official GitHub repository by setting our perturbation as perturbation_key and cell type and chronicity as categorical covariates.

### A.2 Multi-trajectory Experiments, Ablations on Temperature and Number of Time Points

##### CaLMFlow

The model architecture for CaLMFlow consists of a simple MLP to embed the input into the encoder, represented by GPT-2. The dimensionality of the embeddings and hidden states in GPT-2 is set to 32. The GPT-2 Transformer encoder uses 4 hidden layers with 1 attention head per attention layer. The decoder is implemented as a VAE, with a latent dimension of 16 and a hidden dimension of 32. For CaLMFlow, the scaling hyperparameter $\beta$, the number of time points, and the temperature parameter $\tau$ are set to empirically optimal values for each task.

##### CFM

For CFM, we use the standard architecture described in Tong et al. () for the 2D experiments. Trajectories are evaluated using the dopri5 option of the adaptive solver.

## Appendix B Additional results

Figure: (a) CFM
Refer to caption: extracted/5897793/images/unconditional_gen/cfm_uncond_umap.png

Figure: (a) Ground Truth
Refer to caption: extracted/5897793/images/conditional_single_cell/individual_umaps/ground_truth_umap.png

Figure: Figure 8: Ablation results on number of time points. Left: CaLMFlow generated data from 8gaussians to 2moons, using different number of time points. Right: 2-Wasserstein and MMD performances as a function of number of time points. The plots show that increasing the number of time points improves the model performance.
Refer to caption: extracted/5897793/images/time_point_ablation/time_points_scatter.png

Figure: Figure 9: Uncurated conditional image generation comparison between CaLMFlow, CFM models, and DDPM on MNIST. CaLMFlow generates higher image quality images compared to both CFM variants and DDPM.
Refer to caption: extracted/5897793/images/conditional_mnist/normal_generated_FM_images_step_14000.png

## Appendix C Banach Spaces, Volterra Integral Equations and CLMs

We adopt the standard next-token prediction paradigm used in CLMs. Given tokens $x_{0}$ to $x_{k}$, the model predicts $x_{k+1}$, where the sequence $(x_{0},\dots,x_{k})$ corresponds to portions of the conditional flow matching trajectory. In this section, we give the implementation details of the theoretical discussion given in Section 3.2.

Our training procedure enables the causal language model (CLM) to learn system dynamics by modeling sequences of varying lengths. During training, the model predicts the next state in the trajectory given previous states, similar to next-token prediction in language models, but with tokens representing continuous trajectory states.

Solving a nonlinear integral equation generally requires some iterative procedure, where an initial guess is refined iteratively. Since the evaluation of $z$ at time $t$ requires evaluation of $z$ at all time points between 0 0 and $t$ due to the integral $\int_{0}^{t}G(z_{s},t,s)ds$. Observe that once an approximation (guess) $z^{j}(t)$ has been obtained, all evaluations happen in parallel for each iteration, since we can integrate using $z^{j}(t)$, see Zappala et al. () for details. We present the model with sequences starting from the initial state $z_{0}$ and extending to various lengths (e.g., predicting from $z_{0}$ to $z_{1}$, $z_{0}$ to $z_{2}$, up to $z_{0}$ to $z_{N}$). This trains the model on multiple sub-trajectories, which can be used in inference in an iterative manner to output the complete trajectory.

This training methodology can be viewed through the lens of functional analysis, where we consider the CLM as learning operators on a direct sum of Banach spaces. Each Banach space corresponds to sequences (or discretized trajectories) of a particular length, and the direct sum represents the combination of these spaces. The norm on this direct sum space is defined as the sum of the norms of the individual spaces. By minimizing the loss, which aggregates the errors across all sequence lengths, the model learns to operate effectively on each of these spaces.

Let $\mathcal{B}_{k}$ denote some Banach space of functions with domain $[0,t_{k}]$, equipped with an appropriate norm $\|\cdot\|_{\mathcal{B}_{k}}$, e.g., the $L^{2}$ space. The direct sum $\mathcal{B}=\bigoplus_{k=1}^{N}\mathcal{B}_{k}$ is the direct sum of the $\mathcal{B}_{k}$ spaces. The total loss $\mathcal{L}$ in this space is the sum of the component norms:

$$ $\mathcal{L}=\sum_{k=1}^{N}\|z^{i}(t)-\hat{z}^{i}(t)\|_{\mathcal{B}_{k}}$ (13) $$

where $z^{i}$ indicates the ground-truth between $[0,t_{i}]$, and $\hat{z}^{i}$ is the model’s prediction on $[0,t_{i}]$. This loss function aggregates the errors over all sequence lengths, reflecting the norm on the direct sum space.

During inference, the model sequentially generates the trajectory by predicting the next state given the current sequence, starting from the initial state $z_{0}$. At each step, the model uses the operators it has learned during training to map the current sequence to the next state. This process can be interpreted as the model applying learned operators on the respective Banach spaces to generate the trajectory. We note that this process in inference is sequential, but during training is performed in parallel.

The LLM computes:

$$ $z_{i+1}=f_{\theta}(z_{i},t_{i})+\Delta t\sum_{j=0}^{i}G_{\theta}(z_{j},t_{i},t _{j}),$ $$

and as $\Delta t\to 0$ we obtain Equation 4, where we note that during inference there is no need of distinguishing between the $z$ used in the integral operator and the prediciton, as the prediction is used iteratively to produce the next $z^{i+1}$, contrary to training where the ground truth is used to perform the process in parallel.

By framing the problem as a sequence modeling task, the LLM effectively approximates the solution to the integral equation. The sequence of states $\{z_{t}\}$ can be viewed as tokens in a sequence, where each state depends on all previous states due to the integral over past times $s$ in Equation 4. The LLM, with its inherent capability to model long-range dependencies through attention mechanisms, captures this dependence without the need to explicitly compute $f_{\theta}$ and $G_{\theta}$.

In practice, we model Equation 4 using the LLM as follows:

- •
The LLM models the conditional distribution of each $z_{t_{i+1}}$ given the past states $\{z_{t_{0}},\dots,z_{t_{i}}\}$:
$p(z_{t_{i+1}}|z_{t_{0}},z_{t_{1}},\dots,z_{t_{i}})=\text{LLM}(z_{t_{0}},z_{t_{
1}},\dots,z_{t_{i}}),$
(14)
which implements an integral depending on $z_{t_{0}},\ldots,z_{t_{i}}$, i.e. $z(s)$ with $s\in[0,t_{i}]$.

We do not define an iterative procedure here for the training, but leverage the use of LLMs and learn to predict $n$ functions $z^{i}$, each of which model the solution $z(t)$ (i.e. the ground truth) between $t=0$ and $t=t_{i}$. This is in line with the use of LLMs and allows us to formulate the solver procedure on a direct sum Banach space, $X=\bigoplus_{i=1}^{n}L^{2}([0,t_{i}])$, where the target of each $L^{2}([0,t_{i}])$ is the ground truth of the trajectory. In inference, we apply an iterative procedure where the model predicts each component $z^{i}$, and uses this to approximate the integral to compute $z^{i+1}$:

$$ $z^{i+1}_{t}=f(z_{t},t)+\int_{0}^{t_{i}}G(z^{i}_{s},t,s)ds,$ (15) $$

for $t\in[t_{i},t_{i+1}]$, where the integral operator is well defined because $z^{i}$ is defined over $[0,t_{i}]$ – i.e. it is an element of $L^{2}([0,t_{i}])$ as stated above.

In order to perform an approximate integral over the discretized trajectory, we also need a discretized version of Equation (15), which becomes:

$$ $z_{i+1}=f_{\theta}(z_{i},t_{i})+\sum_{j=0}^{i}\Delta t_{i+1}G_{\theta}(z_{j},t _{i+1},t_{j}),$ (16) $$

where $\Delta t=t_{i+1}-t_{i}$, and $z_{i}=z(t_{i})$. Once again, observe that we need each $z_{j}$ with $j=0,\ldots,i$, and we use the ground-truth data for this.

## Appendix D Volterra Flow Matching Objective

The goal of this section is to connect the conditional Volterra flow matching object to the CFM objective. We recall some notation from Lipman et al. (). Let $t\in[0,1]$, $u_{t}(x)$ the marginal time dependent vector field associated with the flow $\phi$, $u_{t}(x|x_{1})$ the conditional time dependent vector field, $v_{t}(x)$ the model learning the vector field $u_{t}$, $q$ the data distribution, and $p_{t}$ the probability density path. Then we can define the flow matching objective and its conditional variant, $\mathcal{L}_{\text{FM}}(\theta)=\mathbb{E}_{t,p_{t}(x)}\left\|v_{t}(x)-u_{t}(x
)\right\|^{2}$ and $\mathcal{L}_{\text{CFM}}(\theta)=\mathbb{E}_{t,q(x_{1}),p_{t}(x|x_{1})}\left\|
v_{t}(x)-u_{t}(x|x_{1})\right\|^{2}$. The key observation is that the gradients of these two objective functions are equivalent:

###### Theorem ( (Lipman et al., 2022 ) ) .

Assuming that $p_{t}(x)>0$ for all $x\in\mathbb{R}^{d}$ and $t\in[0,1]$, then, up to a constant independent of $\theta$, $\mathcal{L}_{\text{CFM}}$ and $\mathcal{L}_{\text{FM}}$ are equal. Hence, $\nabla_{\theta}\mathcal{L}_{\text{FM}}(\theta)=\nabla_{\theta}\mathcal{L}_{
\text{CFM}}(\theta)$.

The crux of the proof of Theorem Theorem rests in the chain of equalities (see Theorem Appendix A for details):

$$ $\displaystyle\mathbb{E}_{p_{t}(x)}\langle v_{t}(x),u_{t}(x)\rangle$ $\displaystyle=\int\left\langle v_{t}(x),\int u_{t}(x|x_{1})p_{t}(x|x_{1})q(x_{ 1})dx_{1}\right\rangle p_{t}(x)dx$ $\displaystyle=\int\left\langle v_{t}(x),\int u_{t}(x|x_{1})p_{t}(x|x_{1})q(x_{ 1})dx_{1}\right\rangle dx$ $\displaystyle=\int\left\langle v_{t}(x),u_{t}(x_{1})\right\rangle p_{t}(x)q(x_ {1})dx_{1}dx$ $\displaystyle=\mathbb{E}_{q(x_{1}),p_{t}(x|x_{1})}\langle v_{t}(x),u_{t}(x|x_{ 1})\rangle.$ $$

We can expand $v_{t}$ in the penultimate line as:

$$ $\int\left\langle v_{t}(x),u_{t}(x_{1})\right\rangle p_{t}(x)q(x_{1})dx_{1}dx= \int\left\langle\int v_{t}(x|x_{1}^{\prime})dx_{1}^{\prime},u_{t}(x_{1})\right \rangle p_{t}(x)q(x_{1})dx_{1}dx.$ $$

Note that CaLMFlow takes conditioned inputs during training, i.e. it is history-aware. One approach to training CaLMFlow like CFM is to randomly sample pairs of conditional trajectories—one for the target trajectory and one for the input trajectory to generate the flow—and minimize the mean squared error, either of the next time step or the output of a numerical solver. However, in practice we have found our model performance to benefit from using the the CVFM objective in Equation 7 combined with the KL divergence regularizer from Equation 9. We leave the exploration of alternative optimization procedures to future work.

###### Theorem ( (Lipman et al., 2022 ) ) .

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 3
presentation: 1
questions: This is more of a question to the other reviewers as much as it is to the authors. I am not too familiar with usual practices when these flow-based approaches are introduced and compared. Are these experiments standard ones that people have used to compare different iterative refinement based samplers? The first set of experiments on mixtures of Gaussians look extremely low-dimensional, and the second set of experiments on single cell sequencing data look relatively unconventional (I like this problem but it is difficult for me to understand that this is one of the standard problems people use to compare these samplers.) I may be completely wrong here (and if so, please do point out,) but I would like the authors to clarify what are more or less standard benchmarks in this area and how the proposed approach compares to the existing ones on those.
soundness: 3
strengths: It is only natural to extend any of these iterative refinement based approaches to learning to sample from a complex distribution to make each refinement step less local and more global. Although it is natural, it has been challenging to do so until recently, as it was not clear whether we can build a powerful neural net that can take as input a long sequence of a trajectory and use it properly. With the recent advances in language models, this doubt is no more, and the authors in this paper demonstrate that indeed we can use such an autoregressive model to build a better sampler based on iterative refinement. Despite the unnecessarily convoluted way of presenting it via so-called Volterra flow, the idea is extremely straightforward, and the authors’ limited experiments do support that there is a benefit to be had from such non-local iterative refinement. It is furthermore interesting to see that they could easily extend it by prefixing their iterative refinement sampler with a natural language description to make it language conditional.
summary: In this paper, the authors propose to use an autoregressive (language) model to imitate a (discretized) flow from the prior to the target distribution. This has a clear advantage, as such an autoregressive model with an unbounded context can take into account non-local structures across multiple timesteps, unlike previous approaches that model a purely local transition (a la “difference”). Once training is over (in the next-token prediction way), one can simply run the autoregressive model forward starting from a sample drawn from the prior distribution to draw a sample from the target distribution.
weaknesses: Unfortunately the current manuscript is extremely difficult to read. One reason i can point out is due to the lack of clear exposition on how this underlying autoregressive model looks like; what does it take as input, and how the prefix is processed to result in the prediction of the next step’s refined observation? Furthermore, the authors’ use of the term “tokenization” confused me quite a lot, as “tokenization” is often used to refer to as the process by which we quantize a continuous observation into a finite set of discrete tokens, while the authors here are referring to discretizing continuous time and space. Such issues with presentation make it pretty much impossible for me to recommend the manuscript in the current form to be accepted at the venue. I would suggest the authors make a major revision to present the proposed approach (which I think may have an interesting impact on the community) more clearly and thoroughly.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 2
questions: What are the author's thoughts on generalizing the method to other kind of spatiotemporal data, e.g., high-dim time series, video data, etc?
soundness: 3
strengths: The technical idea is original and natural, blending LLMs into the framework of VIE flow matching. 

Extensive experiments were carried out on a range of tasks, covering synthetic and real-world data. 
The analysis and ablation studies also show the importance of each component/technique in the method, providing insights for future improvements. 

The paper is generally clear.
summary: The paper introduces CaLMFlow, a framework that leverages LLM for continuous data generation. The paper formulates flow matching as a sequence modeling task, and applies LLMs to learn the complex flows. By tokenizing the space and time, the approach enables efficient handling of high-dim data. On a range of tasks (e.g., single-cell perturbation response prediction), the method shows strong performance.
weaknesses: The method seems general but the only kind of real data in the experiments are single-cell data. Without further evidence, it is hard to judge whether the significance will be high or not, broad or not.

The writing has some typesetting issues. E.g., 
- line-213, $T$ tokens 
- line-219, $D_{text{in}}$

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: Problem Definition: What is the task setup for the synthetic experiments? How many training and test data samples are used, and what does each training sample look like (e.g., is it a two-dimensional Gaussian distribution)?

Parameter Comparison: There is no parameter comparison between different methods, which makes the evaluation less fair for data-fitting tasks.
soundness: 2
strengths: The writing is clear and accessible.

The idea is easy to understand.

Experiments are conducted on both synthetic and real data, enhancing the validity of the results.
summary: This paper proposes CALMFLOW, an autoregressive modeling method for continuous data. It tokenizes continuous data into small time segments and uses numerical solutions of ODEs to model the data. Experiments are conducted on synthetic data and real cell-perturbation data to evaluate the approach.
weaknesses: Trivial Task: The paper focuses on continuous data modeling, which may not present a sufficiently challenging or novel task. The modeling of Gaussian distributions and the MNIST dataset appears trivial, and many existing multimodal methods can already handle cell data modeling effectively.

Mismatched Motivation and Experiment: Although the paper initially highlights the difficulty of solving ODEs, the experiments focus only on continuous data modeling and do not fully support or address the stated challenges with ODEs.

Overclaim of Novelty: The claimed tokenization method seems indistinguishable from standard window-based segmentation techniques, which questions the novelty of the approach.

Meta-review: The paper introduces CaLMFlow (Causal Language Models for Flow Matching), a novel framework that reformulates flow matching as a Volterra integral equation (VIE) (instead of a differential equation), leveraging large language models (LLMs) for continuous data generation. The authors claim that CaLMFlow bridges discrete language modeling and continuous generative modeling by implementing tokenization across space and time, enabling efficient handling of high-dimensional data. The method is demonstrated on synthetic and a real-world data application (single cell sequencing dynamic), showcasing its ability to incorporate textual context and generalize to unseen conditions.

**Strength: Timely, Innovative Connection of two Active Areas of Research**. There is a long line of work aiming to repurpose neural network architectures typically used in language modeling (RNNs, LSTMs, Transformers) for continuous data domains. This paper fits withing that broader literatures, approaching this problem by marrying LLMs and flow matching, two very active areas of research, but so far mostly independent. Therein lies, in my view, the main strength of the paper: an innovative approach that presents a unique application of LLMs to flow matching, bringing the promise of flow matching to other domains. The use of integral (instead of differential) equations to model flow matching is interesting, but not an original contribution of this work (see Zappala et al.). Furthermore, I think its use requires more motivation. The paper makes a broad argument that IEs should be preferred over ODEs, but there's inherent computational-robustness tradeoffs that are not discussed here in depth. Additionally, the paper implicitly assumes that the conditions under which IEs might be preferable (singularities, non-locality, weak solutions) are ipso facto true in the applications of interest, but this does not seem to be case in the synthetic Gaussian experiments (despite their high-dimensionality), and unclear in the single cell perturbation ones. A stronger case for VIEs would have been made with an application where ODEs provably suffer from such issues, or to demonstrate that the decreased performance of ODE-based CFM in table 1 is due to stiffness or instability and not implementation issues. 

**Weakness: Methodological choices lack motivation (see above)**. 

**Weakness: Clarity and Presentation**. The reviewers noted that the manuscript is difficult to read due to unclear exposition and confusing terminology, particularly around the use of “tokenization.”. The revised version made important improvements in clarity, but introduced substantial changes. While such improvements are welcome, the fact that such an extensive rewriting was necessary suggests that the paper was submitted in a form not ready for prime time. Furthermore, while the revised version is definitely easier to read, many of the points raised by reviewer ktMR about missing details still stand. 

**Weakness: experimental evaluation**. As observed by all three reviewers, the experimental framework consists of very simplistic synthetic datasets and non-standard (for flow matching) single-cell generation task. This makes it hard to compare the method to alternatives. Additionally, there does not seem to be experimental results comparing the computational cost of this method against the baselines. This is particularly problematic given the known robustness-computation tradeoff of solving integral vs differential equations. 

Overall, I think this is a borderline paper with high unrealized potential that would benefit from a broader evaluation, rewriting, and an additional round of peer review.

Decision: Reject