Categorical Schrödinger Bridge Matching

Thank you for your comments and positive evaluation. We will correct the typos you mentioned. If you're interested in the topic, we recommend checking the other reviews and these works on the Schrödinger Bridge Problem [1, 2, 3].

[1] Kim, Jun Hyeong, et al. "Discrete Diffusion Schrödinger Bridge Matching for Graph Transformation." arXiv preprint arXiv:2410.01500 (2024).

[2] Shi, Yuyang, et al. "Diffusion Schrödinger bridge matching." Advances in Neural Information Processing Systems 36 (2023): 62183-62223.

[3] Gushchin, Nikita, et al. "Adversarial Schrödinger Bridge Matching." The Thirty-eighth Annual Conference on Neural Information Processing Systems.

Dear reviewer 1AGf thank you for your questions and commentaries.

[Q. 1] Can the author clarify the difference between proposed method to DDSBM, which does present theoretical results for continuous-time IMF in discrete spaces? I'm fairly familiar with DDSBM and, since both are based on IMF, the methods seem to collapse in practice when time is discretized.

The answer to this question can be found in our response to reviewer MBG2 [W. 1] and [W. 2].

[W. 1] I'm not convinced by Table 2. The author should report ASBM and DSBM in a similar GAN-based continuous latent spaces for a fair comparison.

Regarding the setup on the CelebA dataset, we agree with your concerns. We did attempt to train DSBM in the latent space. For a fair comparison, we ran DSBM on the same latent space used for CSBM, following the approach in [1, Appendix G]. However, the results were not satisfactory, as the model tended to collapse to the identity mapping with $\epsilon = 1$ and $\epsilon = 3$ (LINK: see figures with prefix 'latent'). Due to these limitations, we did not proceed with training ASBM and chose not to compare both methods with CSBM in such settings. One may ask why CSBM performs better in this setting. We hypothesize that this is due to the choice of the reference process, with $q^{\text{unif}}$ being more suitable for the latent space of VQ-GAN.

[1] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[W. 2] It seems faulty to claim the dimensionality of CelebA Faces to be $1024^{256}$. Given the factorized parametrization, the dimension that the proposed method handles should be $1024*256$.

This is not a mistake but rather a point that could be clarified more explicitly. When we refer to $1024^{256}$, we are indicating the complexity of the data, not the complexity of the model parametrization. The quantity that you mentioned, $1024 \times 256$, corresponds instead to the complexity of the generated samples under $q_\theta(x_1 | x_{t_{n-1}})$.

It is also worth noting that as the number of sampling steps increases, the complexity of the resulting composition of distributions also grows. As mentioned in our response to reviewer McJt [Q. 1], this increase in sampling steps can lead to higher-quality samples and help mitigate issues related to factorization.

[R. 1] Fig 3's caption should clarify that images are generated in VQ-GAN latent spaces. I think it's boarderline misleading to not mention specifically in the caption that it's a GAN-based latent-space image experiments. GAN's latent spaces are not only of lower dimension but also much structural.

We will specify this aspect of training in the caption, as you have suggested.

Concluding remarks. We hope that, with the above clarifications, you will kindly reevaluate our work and find it deserving of a higher rating.

Dear reviewer atNH thank you for your questions and commentaries.

[R. 1] Hoogeboom et al., and Gat et al, are mentioned but could use further discussion and perhaps comparison

Regarding the references [1, 2] you mentioned, we believe they are not suitable for comparison in our setting. For [1], the practical differences from D3PM [3] (which is our backbone method) are minimal, as [3] extends the discrete diffusion framework introduced in [1]. While the second method [2] could, in principle, be used as a backbone for our method, it is well known that lower values of $\epsilon$ (or $\alpha$ in our work) make it significantly harder for the model to approximate the target distribution (see [4], Figure 7). Current discussion will be added to the revised version.

[1] Hoogeboom, Emiel, et al. "Argmax flows and multinomial diffusion: Learning categorical distributions." Advances in neural information processing systems 34 (2021): 12454-12465.

[2] Gat, Itai, et al. "Discrete flow matching." Advances in Neural Information Processing Systems 37 (2024): 133345-133385.

[3] Austin, Jacob, et al. "Structured denoising diffusion models in discrete state-spaces." Advances in neural information processing systems 34 (2021): 17981-17993.

[4] Shi, Yuyang, et al. "Diffusion Schrödinger bridge matching." Advances in Neural Information Processing Systems 36 (2023): 62183-62223.

[W. 1] Could benefit from additional experiments on other categorical datasets (e.g., discrete text tokens or molecules) and perhaps more comparison with baselines.

We believe the current set of experiments is sufficient to support our claims. In particular, the practical implementation of the VQ-GAN setup does not differ significantly from potential setups with text data. The model and reference process will remain the same. While graphs are indeed more challenging due to their structural complexity, this domain has already been addressed by DDSBM. As explained in our response to reviewer MBG2 [W. 2], we do not include a direct comparison with DDSBM and instead focus on exploring other experimental settings.

[Q. 1] Clarification: In Alg 1, when sampling $(x_0,x_1)$ using $q_\eta$ or $q_\theta(x_1|x_0)$, do the authors mean $x_1$ is simulated from $x_0$ with $N$ steps? or estimated directly from the $x_1$ timestep? if it is the latter -- could you provide more details on how this is done in practice?

Your first suggestion is correct. To generate samples, we follow a standard diffusion sampling procedure. After the $l$-th step of D-IMF, we obtain $q^l_\theta(x_1 | x_{t_{n-1}})$. Moving to $(l+1)$-th D-IMF iteration, we first apply the trained model to generate $x_1$ taking as an input a dataset point $x_0$. We then perform posterior sampling using $q^{\text{ref}}(x_{t_n} | x_{t_{n-1}}, x_1)$ to obtain $x_{t_n}$. This procedure is repeated for $N+1$ steps to obtain samples from coupling $q^l_\theta(x_0, x_1)$. For the backward parametrization, we use the same scheme.

[Q. 2] Can the authors discuss ways to reduce the information loss arising from the factorization of the conditional distributions?

Please refer to our response to reviewer McJt [Q. 1].

[R. 2] However, the chosen parameters for stochasticity level (alpha) appear somewhat ad-hoc. further explanation of the choice of particular values would be helpful

The pattern of selecting $\alpha$ follows the same intuition as choosing $\epsilon$ in continuous SB methods. Specifically, lower values of $\alpha$ lead to less stochasticity in the trajectories, resulting in higher similarity to the input data but a lower-quality approximation of the target distribution. At very low values, the model may collapse due to insufficient stochasticity. Conversely, higher values of $\alpha$ introduce more variability, improving the quality of the approximation but reducing similarity to the initial data. Beyond a certain point, excessively large $\alpha$ values make the model difficult to train, leading to a drop in both quality and consistency. Unfortunately, the effective range of these behaviors is highly dependent on the dataset and the chosen reference process. Nonetheless, we provide reasonable baseline values from which one can begin and adjust as needed.

Dear reviewer MBG2, thank you for raising important questions regarding our paper.

[W. 1] Questionable Novelty of the Theoretical Contribution

We respectfully disagree. First, as highlighted in Table 1, the continuous-time frameworks DDSBM [1] and DSBM [2] rely heavily on the theoretical foundation established in [3], which does not extend to discrete time. Therefore, our work should not be viewed as a trivial extension of these frameworks but rather as a distinct theoretical foundation that shares certain similarities. Furthermore, the core theoretical contribution of our paper goes beyond simply generalizing the D-IMF procedure from [4] to discrete data: the discrete case emerges as a consequence of a broader generalization of the D-IMF procedure to arbitrary reference processes. As noted in the article (see footnote 1, page 5), our generalization enables ASBM [4] to operate with any Markov process $q^{ref}$.

[W. 2] Moreover, there is a deeper concern regarding the algorithm itself..., one must question: Is CSBM simply DDSBM with a different motivation?... For example, DDSBM has been proposed specifically for graph-structured discrete data, yet it is not included in the evaluation... In what ways does the derived algorithm differ from DDSBM?

We agree that adding practical distinctions will strengthen the comparison. Thus, we will include them in the revision.

First, let us clarify how DDSBM derives its loss. In theory, it matches the generator matrices $A$ (see [1, Equations (5, 6)]). However, in practice, authors of [1] discretize time, which leads to minimization of $D_{KL}(q^{ref}(x_{t_n}|x_{t_{n-1}}, x_1)|| q_{\theta}(x_{t_n}|x_{t_{n-1}}))$ (see the derivation in [1, Appendix E.1]). Thus, as it could be mentioned, it is indeed the same loss as presented in our article. However, since we have theoretically alternative objective $D_{KL}(q^{ref}(x_{t_n}|x_{t_{n-1}})||q_{\theta}(x_{t_n}|x_{t_{n-1}}))$ (see Prop. 3.3), we can match distributions directly by employing various loss functions such as MSE or even adversarial training as in ASBM [4]. We conducted extra experiments using only an MSE loss and observed comparable to KL (LINK: see figures with prefix 'toy'). Thus, the similarity in practical implementation reflects our design choice to use this particular parametrization. Alternative approaches could be easily used. This also explains why we did not include a comparison with DDSBM.

[W. 3] A missing baseline discrete OT with [5]

If we correctly understand, you propose first using classical discrete OT methods to get aligned data followed by the application of [5]. However, there is one crucial problem with this setup. The SB problem that lies behind [5] is still considered to be continuous, i.e., continuous Brownian motion is used as a reference process. This breaks the assumptions on discrete data. We will cite [5] in the final revision and discuss it, but we think that a comparison with it is impossible due to the reasons mentioned above.

[W. 4] Additionally, while the paper compares CSBM to ASBM and DSBM, these are methods designed for continuous spaces, meaning they are not necessarily the most appropriate baselines...

To our knowledge, there are no other discrete domain translation models for unpaired data. Thus, the only methods we compare are continuous-space methods ASBM and DSBM.

[W. 5] ...all experiments focus on image-based tasks, which, despite using vector quantization, are inherently less discrete than the originally mentioned data types...

We apologize, but the meaning of "less discrete" is unclear to us. Still, for further clarification regarding the choice of experiments (especially texts), please refer to our response to the reviewer atNH [W. 1].

[Q. 1] What values does the discrete-time perspective bring to the SB framework for categorical data?

The main theoretical advantage is the ability to consider CSBM with $N = 1$, which is guaranteed to converge to the SB. In contrast, continuous-time setups typically require to assume $N=\infty$ to achieve convergence. From a practical side, our framework also enables the flexible selection of loss functions for matching the transition distributions, which we mentioned in answers to your previous questions.

Concluding remarks. We hope that, with the above clarifications, you will kindly reevaluate our work and find it deserving of a higher rating.

[1] Kim, Jun Hyeong, et al. "Discrete Diffusion Schrödinger Bridge Matching for Graph Transformation."

[2] Shi, Yuyang, et al. "Diffusion Schrödinger bridge matching."

[3] Léonard, Christian. "A survey of the Schrödinger problem and some of its connections with optimal transport."

[4] Gushchin, Nikita, et al. "Adversarial Schrödinger Bridge Matching."

[5] Somnath, Vignesh Ram, et al. "Aligned diffusion schrödinger bridges."

Dear Reviewer McJt, thank you for pointing out the unclear parts you encountered. We will revise the text where you have highlighted, as much as possible.

[R.1] Line 110 column 2: "additional properties" sounds too vague, maybe state some examples

In line 110, by "additional properties", we refer to the reciprocal and Markovian properties, which allow us to use diffusion models (specifically, bridge matching models) as the backbone of our domain translation framework. Without these properties, we would be limited to training one-step models such as GANs, which have since been largely replaced by modern diffusion-based approaches.

[R.2] Line 235 column 2: $D=1$ seems a weird way to put it, and it confused me because I had forgotten what $D$ was.. maybe just say that you work with $\mathbb{S}$ and it'll be clear that $D=1$

In line 235, the purpose of $D = 1$ is to simplify the transition matrices defined in Equations (7, 8). Since we are using factorization, we consider feature-wise rather than data-point-wise transition matrices. So, to maintain compactness of the text, we omit introducing $Q_n$ for arbitrary $D$, as it is not used in our approach, again, due to the factorization.

[R.3] Line289 column 1: when you introduce this factorization, maybe briefly write about the limitation that this imposes (I know that this has been written in the "limitations" part, but I think it's worth putting it here too)

We agree that it is indeed important to clarify the factorization in line 289. We will include a reference to the limitations section in order not to distract the reader from the flow of the main text.

[R.4] Line 290-299 column 1 : "Since, in fact, we need N+1 neural nets to do the prediction of endpoints at each time step, we simply use a single neural network with an extra input n" this sentence is unclear to me...

In lines 290–299, we intended to explain that we use a single neural network for all time steps rather than separate networks for each step, which is a common practice. To simulate the stochastic process, one would typically require $N+1$ distinct functions $q(x_1 | x_{t_{n-1}})$ for each $n$. However, training $N+1$ neural networks is computationally expensive. Instead, we use a single neural network for all transition steps (i.e., the transition function $q(x_1 | x_{t_{n-1}})$), with additional time conditioning $q_\theta(x_1 | x_{t_{n-1}}, t_{n-1})$ for all $n \in [1, N+1]$. If you are also interested in the sampling procedure using the trained model, please refer to our response to the reviewer atNH [Q. 1].

[R.5] Line 354-355 column 2: "discrete space of images but not continuous" this looks weird.. yeah it's discrete and not continuous, what did you want to say with the "but" part? Please correct/erase whatever is needed to make this right.

In lines 354–355, we intended to highlight that image spaces are typically treated as continuous rather than discrete, i.e., the space of image pixels is commonly represented as the interval $[0, 1]^D$ rather than the discrete set {$0, 1, \dots, 255^D$}.

[Q. 1] I'd like to know more precisely some examples of what kind of dependencies will be lost in the dimensional factorization, and also if the authors have any ideas on how to "do better than" the factorization $S^D \to D\times S$ metioned in the un-numbered formula following (9). I mean, how would you increase a bit the complexity of the models to try to make it "forget less" about the dependencies between dimensions? In the case of a discrete space this is easier than in general, so it's worth giving some pointers.

Regarding factorization, to the best of our knowledge, the work [1] is the only existing work addressing this issue. Authors introduce an additional generative model that models the copula of the factorized distributions, which is, informally, the part of the joint distribution $q_{\theta}(x_1 | x_{t_{n-1}}) = q_{\theta}(x^1_1, ..., x^D_1 | x_{t_{n-1}})$ that is lost due to the factorization.

However, practically, it seems that the issue of factorizayion can also be mitigated just by increasing the number of steps, as demonstrated, for example, in our experiments with C-MNIST. The more steps that are taken during sampling, the more expressive the resulting composition of distributions becomes. As a result, repeatedly incorporating the full previous state leads to more correlated features. Even though factorization has implications for our practical implementation, it is important to recall, as stated in the article, that this issue is inherent to all discrete diffusion models.

[1] Liu, Anji, et al. "Discrete Copula Diffusion." arXiv preprint arXiv:2410.01949 (2024).