Latent Zoning Network: A Unified Principle for Generative Modeling, Representation Learning, and Classification

We sincerely appreciate your thoughtful and constructive feedback. Below, we provide detailed responses to each of your questions.

Q1. Performance.

Thank you for raising this important point. We would like to clarify our perspective and provide additional context.

• In Use Case 1 and Use Case 3, LZN already outperforms the state-of-the-art in generative modeling (see Tables 1 and 4). In Use Case 3, LZN also achieves classification accuracy that is within 1% of the best-performing method, despite being designed as a unified model for both generation and classification. In contrast, the baselines are optimized specifically for classification, not for multiple tasks.

• For Use Case 2, which is the only case where LZN was initially behind by a larger margin, we are happy to report that our updated results significantly improve upon the submission. After extending training to 4.5 million iterations, the top-1 accuracy increases to 69.3% (+3.7%) and top-5 accuracy to 89.1% (+2.0%). With these results, LZN is now among the top 4 methods on this benchmark. Importantly, these competing methods are designed exclusively for representation learning, whereas LZN supports more tasks including generation. Notably, LZN also matches or outperforms seminal methods such as SimCLR. We believe these are strong results.

• Finally, we would like to emphasize our broader perspective. Even if one considers only Use Case 1, LZN already advances the state-of-the-art in image generation, which we believe is a substantial and meaningful contribution on its own. Beyond that, LZN is a general-purpose framework capable of solving both generative and discriminative tasks within a unified design—yet it still performs competitively with specialized, task-specific models. In our view, this generality combined with strong performance across domains goes above and beyond typical expectations of papers.

  We respectfully believe that emphasizing a modest performance gap in one setting, while overlooking areas where LZN achieves top results, would not offer a balanced or fair assessment.

Q2. Comparison with baselines that can deal with task unification

Thank you for the helpful suggestion. In our experiments, we have chosen to compare with the state-of-the-art methods for each benchmark, which we believe sets a higher bar than comparing against older or weaker baselines, regardless of whether they support multiple tasks.

That said, we are happy to include additional unified-task baselines in the revision—especially if the reviewer has any specific suggestions in mind. We greatly appreciate such pointers that can help us strengthen the comparisons.

Q3. Organization and formatting.

Thank you for the thoughtful suggestions. We agree that clarity and structure are important, and we appreciate your perspective.

Regarding Section 2.1, we did consider alternative placements during the writing process. Ultimately, we decided to keep it early in the paper for the following reasons:

• Motivational context: Section 2.1 explains why existing frameworks fall short of achieving our goal of a unified model for generation, representation learning, and classification. Without this context, readers may not fully appreciate the need for our proposed approach.
• Conceptual foundation: It also provides essential background—such as the roles of encoders, decoders, and latent spaces—that our method builds upon. Skipping this section would make the introduction of LZN less accessible to a broader audience.

That said, we understand the concern about length and will work on tightening the section to improve flow and focus.

We will also address the other formatting issues you kindly pointed out in the revision. Thank you for catching that!

Thank you very much for the constructive comments, which are very helpful for us to improve the paper. Below, we respond to all your questions one by one.

Q1. FID Computation

We followed exactly the setting that the reviewer suggested: evaluating the FID score on the training set with a large number of samples. In fact, we went further to ensure a fair and rigorous comparison, especially on the FID implementation, which many prior works in this domain overlooked. Please find the details below.

• We have reported the number of samples used for FID computation and the dataset split they are taken from in Appendix C.3, which we repeat here for clarity:

  • Following the convention [47,67,42,80,43], the metrics are all computed using the training set of the dataset.
  • For FID, sFID, IS, precision, and recall, we subsample the training set and generate the same number of samples to compute the metrics. The number of samples are:
    • CIFAR10: 50000 (the whole training set).
    • AFHQ-Cat: 5120, the largest multiple of the batch size (256) that is less than or equal to the training set size (5153).
    • CelebA-HQ: 29952, the largest multiple of the batch size (256) that is less than or equal to the training set size (30000).
    • LSUN-Bedroom: 29952, the largest multiple of the batch size (256) that is less than or equal to 30000. We limit the number of samples to 30000 so that the computation cost of the metrics are reasonable.

• It is known that different FID implementations can produce different results [56]. However, many works in this space neglect this fact and compare results from different FID implementations, which can lead to unfair comparisons. We take this issue seriously. To ensure fairness:

  • We evaluate all models using three different FID implementations, as shown in Table 5 and Table 7.
  • These evaluations are run by us (i.e., we do not use reported values from other papers) to ensure that the evaluation settings and code are aligned.
  • Despite differences in absolute values, all three implementations consistently support that LZN improves image generation quality over the baselines.

We hope this clarifies our careful handling of FID evaluation and ensures the fairness and validity of our reported results.

Q2. Re-training the baselines

• The baseline results reported in Table 2 are taken directly from prior papers, with the source noted next to each number. While we agree that re-running all baselines would be ideal, we would like to clarify the rationale behind our approach:

  • Following standard convention: There is a long-standing convention in the representation learning field to reuse numbers reported in prior works rather than re-running all baselines independently. This practice dates back to seminal works such as MoCo (2019) and SimCLR (2020), if not earlier. We adhere to this convention and responsibly cite the source of each number to ensure transparency and reproducibility.

  • Re-running all baselines in our case may not be useful: We believe re-running baselines is most valuable when the proposed method builds directly on top of a baseline. For example, in our image generation experiments (Tables 1 and 4), where LZN is applied on top of Rectified Flow, we re-implement Rectified Flow ourselves, ensuring identical training setups for a fair comparison.

    However, in Table 2, our method (LZN) is fundamentally different from the baselines. We only share the network architecture (from SimCLR), but not the training algorithm or objective. In such cases, hyperparameter or code alignment is not viable, and re-running the baselines would not offer additional insight, unless there are concerns about the trustworthiness of the original results, which we assume are reliable.

• Similarly, except for the “RF+LZN (no gen)” baseline, the other baseline results in Table 3 are taken from [35], as noted in the table caption. This follows the same reasoning: the baselines are standard classifiers trained with cross-entropy losses, whereas our method employs a fundamentally different objective and training procedure. Because there are no overlapping hyperparameters or training pipelines, re-running these baselines would not yield new insights.

  We do, however, re-run the “RF+LZN (no gen)” baseline ourselves, as it serves as a direct ablation of our full method. In this case, aligning experimental settings is essential to draw valid conclusions about the contribution of each component.

We hope this clarifies our design choice and assures the reviewer that our methodology is both principled and consistent with field-wide practices.

Q3. CMMD metric.

Thank you for pointing out this metric. Following the suggestion, we compute the CMMD scores for all our image generation experiments. The images used for computing the metrics are the same as described in Q1.

The results are shown below:

We can see that RF+LZN consistently outperforms RF in all cases. These results further support our claim that LZN improves generative modeling quality across a variety of datasets.

We will add the results to the revision.

Q4. Why LZN Can Improve Image Quality

Thank you for the question. Below, we explain why applying LZN on top of a generative model can improve image quality.

Given a dataset, the goal of a generative model (Rectified Flow in our case) is to learn the distribution $p(x)$ of samples $x$. The entropy (or diversity) of this distribution affects the difficulty of learning. For example,

• If $p(x)$ is highly deterministic (i.e., low entropy), it is easier for the model to learn.
• If $p(x)$ has a diverse support (i.e., high entropy), learning becomes more challenging.

LZN helps by introducing a strong conditioning signal that reduces the effective entropy of the target distribution. Specifically, LZN maps images into structured regions of a latent space (see Section 2), where each latent code $z$ corresponds to a specific image or visually similar images (when using minibatch approximation in Appendix A.3). As a result, when having this z as an condition, the generative model only needs to learn $p(x|z)$ — a much simpler, more concentrated distribution — rather than the full, high-entropy distribution $p(x)$. This makes the learning problem easier, which leads to higher-quality image generation.

Q5. How to reconstruct images using RF, and why LZN reduces reconstruction error a lot

• How to reconstruct images using RF.
  This process is described in Appendix C.3 (lines 1020–1026). Rectified Flow (RF) defines an ODE that maps between the data space (images) and a Gaussian latent space (note: this latent space is different from the LZN latent space). To reconstruct an image:

  1. We first apply the inverse ODE to map the image to its corresponding Gaussian latent representation.
  2. Then, we apply the forward ODE to map this latent back to the image space.

  This reconstruction procedure is a standard technique already described in the original Rectified Flow paper (e.g., see Figure 12 therein), not something we introduced.

• Why LZN reduces reconstruction error.
  This follows the same intuition as discussed in Q4. LZN assigns a unique latent $z$ to each image, thereby capturing the dominant variation in the data. As a result, the conditional distribution $p(x|z)$ that RF needs to learn becomes simpler and lower in entropy. This reduction in complexity allows the model to reconstruct images more accurately from their latent representations.

Q6. “The paper lacks a more complete evaluation and discussion of the results, for example between LZN and Autoregressive methods.”

We apologize, but we do not fully understand the question, as our current implementation does not rely on autoregressive methods. If the reviewer could point us to a specific paper or work related to this concern, we would be happy to review it and include it in our discussion.

Q7. Variable length latent spaces

This is a super insightful and interesting question! Our current approach does not support variable-length latent spaces. To handle token sequences of variable length, our current idea is to embed them into a fixed-dimensional vector, which can then be modeled with LZN. This is feasible given recent advances such as [arXiv:2310.06816], which show that variable-length sequences can be effectively encoded into fixed-size embeddings while preserving semantic information. We believe this is a promising direction and plan to explore it in future work.

Q8. Table/figure formats.

Thank you for catching these formatting issues and the suggestions. We will correct them in the revision.

Thank you very much for appreciating the value of our work and for the thoughtful and interesting questions! Below, we address each of your questions one by one.

Q1. Advertise against more recent methods than MoCo.

Thank you for the suggestion; we will revise it in the paper.

We would also like to take this opportunity to highlight that our results have substantially improved since the submission. By extending training to 4.5 million iterations, the top-1 accuracy improves to 69.3% (+3.7%) and the top-5 accuracy improves to 89.1% (+2.0%). With these updated results, LZN is now among the top 4 state-of-the-art methods for this benchmark. Importantly, these competing methods are designed exclusively for representation learning, whereas LZN supports more tasks including generation.

Q2. Unifying the geometry and topology of differently behaved inputs.

• We would like to first clarify the concern regarding “connected components.”

  Our encoders map images and labels to anchor points. Since the dataset is finite, the set of anchor points for both images and labels is also finite. These anchor points are then passed through the latent computation process (via flow matching), which maps them into a shared Gaussian prior space, and each image and label is associated with a latent zone in this space.

  Importantly, since flow matching is a continuous transformation, each latent zone is always a connected region. The latent alignment between zones occurs in the continuous Gaussian prior space, not directly on the discrete anchor points. Therefore, aligning zones in the Gaussian space is not mathematically infeasible. Indeed, in our experiments on CIFAR10 (Use Case 3), we observe that the training alignment accuracy exceeds 95%, demonstrating that the method can learn strong alignments in practice.

• That said, it is possible that finite-capacity encoders may not achieve perfect alignment between zones of different data types (e.g., images and labels). However, perfect alignment is not a necessity for meaningful learning. As long as the objective encourages the encoder to align zones in the correct direction, the model can learn useful and effective representations. As an analogy, consider traditional classification with cross-entropy loss: a classifier of finite size might not perfectly map images to their labels, but it can still generalize well and achieve high performance. Similarly, in LZN, even imperfect latent alignment can lead to strong performance across tasks.

Q3. Design choices of latent alignment

In Section 2.2.2, we discussed several strawman alternatives to latent alignment and analyzed why they fail:

• Directly minimizing the distance between the reconstructed anchor points and the target anchor point (Lines 206–214).
• Soft alignment across all time steps (Lines 215–227).
• Max soft alignment without cutoff (Lines 228–232).

We believe these discussions follow a natural progression from the most naive solution to the final one, as each failure motivated the refinement that led us to the final alignment strategy. In fact, these discussions honestly reflect the thought process we went through during the development of our method.

We understand the reviewer's desire to see more detailed analysis of the design choice. To address this, we conducted an ablation study on the only hyperparameter introduced in the latent alignment algorithm: u, which determines how many time steps are excluded from the latent alignment objective (Line 234).

In our main experiments, we set u = 20 (out of 100 total steps) across all tasks. Here, we perform an ablation in Case Study 3, reducing u to 5 (i.e., 4× smaller). The results are summarized in the table below:

We can see that this parameter u does not affect the results much, and the performance remains better than the baseline RF across most metrics. This is expected. Unlike common hyperparameters (such as loss weights) that influence the optimal solution, u does not alter the optimal solution, which is the perfect alignment between two latent zones. Instead, this parameter is introduced solely to help avoid getting stuck in local optima (as discussed in Lines 231–236). We expect that any small but non-zero value of u should be sufficient in practice.

Q4. Why unconditional generation does not work

Thank you for the great question! We believe there are two main reasons:

• When using the minibatch approximation for latent computation (see Appendix A.3), the latent zone assigned to an image can vary across batches. As a result, the same latent could be asked to reconstruct different images in different batches, and this inconsistency makes unconditional generation with LZN infeasible.

• Even without minibatching, generation remains challenging. To isolate this issue, we conducted an experiment where we reduced the dataset size so that all training images could fit into a single batch. This eliminates the need for minibatch approximation. However, even in this case, the generated images were not ideal. We hypothesize that the problem arises from the strict requirement that the Gaussian prior space be partitioned exactly among all images, with no gaps. This is a very strong constraint and could be difficult for the model to satisfy. In particular, regions where latent zones connect are extremely sensitive: a small movement across zone boundaries can drastically change the target image, making training hard.

We believe solving this issue is an important direction for future work.

Q5. Are LZN latents interpretable?

Thank you for the very interesting suggestion! Following your idea, we conducted the following experiment to examine the interpretability of LZN latents.

We take images from randomly selected 20 classes from the validation set of ImageNet and computed their embeddings using the LZN representation model from Use Case 2. We chose the validation set to ensure that the results are not influenced by training set overfitting. We then projected these embeddings into a 2D space using t-SNE—a widely used method for visualizing high-dimensional representations, following seminal works such as SimCLR.

The resulting t-SNE plot (not shown here due to rebuttal constraints) indicates that samples from different classes are well-clustered. To summarize this quantitatively, we present statistics below. The embeddings of different classes exhibit distinct means, and their standard deviations are small compared to the overall standard deviation across all samples. This suggests that the LZN latents are indeed clustered and show class-wise separation, indicating meaningful and interpretable structure in the learned representation.

We will add the experiment to the revision.

Q6. Related work

Thank you for pointing out the inspiring work “Masked Completion via Structured Diffusion with White-Box Transformers.” We appreciate the connection and will include a citation and discussion of it in the revised version of our paper.

Thanks for the detailed reply! Below is a point-by-point discussion of some things which I still have questions about, but I continue to recommend acceptance (hence keep my score).

By extending training to 4.5 million iterations, the top-1 accuracy improves to 69.3% (+3.7%) and the top-5 accuracy improves to 89.1% (+2.0%). With these updated results, LZN is now among the top 4 state-of-the-art methods for this benchmark.

By this benchmark you mean ImageNet classification (reported in Table 2)? There are better models (like DINOv2) not represented in the paper, so saying top-4 (without qualifying among the models in Table 2) seems like a bit of an overclaim. If you mean something else, please clarify.

Our encoders map images and labels to anchor points. Since the dataset is finite, the set of anchor points for both images and labels is also finite.

Just to clarify what I mean --- you construct the mapping in the paper by considering only the finite set of sample images/labels/etc., and thanks to your clarification and a re-read it seems rigorous to me. I am now curious about what happens if you use the learned network (obtained via finite samples) on the underlying data distribution, i.e., what happens when you push-forward your multimodal data through the LZN system. This is where my question about geometry/topology looks like (because the image distribution is infinite/continuous and labels are finite/discrete). I still don't know a real answer, and it seems out-of-scope for this paper, but it seems quite interesting to me.

In Section 2.2.2, we discussed several strawman alternatives to latent alignment and analyzed why they fail:

Thanks. I think this discussion is quite good, but a bit truncated/short (almost certainly due to format constraints). In my opinion it would be a very valuable contribution to expand on this, as few papers actually do go over their thought process.

We hypothesize that the problem arises from the strict requirement that the Gaussian prior space be partitioned exactly among all images, with no gaps. This is a very strong constraint and could be difficult for the model to satisfy. In particular, regions where latent zones connect are extremely sensitive: a small movement across zone boundaries can drastically change the target image, making training hard.

This seems like an interesting phenomenon; I would have a priori predicted otherwise, i.e., when the number of samples is small, there is lots of room at the boundaries between valid images allowing for a smooth change.

Interpretability experiment

Sort of hard to see the pattern from the table, but this is obviously due to the no-image constraint. Looking forward to seeing it in the revision.

Related work

Here is a sample of more papers that use dynamical systems as models for forward passes (hence a kind of representation learning):

1. The Neural Covariance SDE: Shaped Infinite Depth-and-Width Networks at Initialization (https://arxiv.org/abs/2206.02768)
2. Do Residual Neural Networks discretize Neural Ordinary Differential Equations? (https://arxiv.org/abs/2205.14612)

Thank you very much for these very insightful questions! Below, we address all your questions point-by-point.

---

---

Q1. Related work

Thank you very much for pointing out this interesting and relevant work. While we acknowledge its connection to our topic, we would like to clarify that both the goals and the techniques of the two papers differ significantly.

Goals

• The goal of the referenced paper is to improve generation tasks.
• In contrast, our goal is more ambitious: to develop a unified framework that supports generation, representation learning, and classification. This broader scope requires a different design philosophy and technical approach, as detailed next.

Techniques

• The referenced paper applies flow matching to the latent space of a pre-trained Stable Diffusion autoencoder, which is reasonable when focusing solely on generation. However, such a latent space is high-dimensional and retains spatial structure, limiting its suitability for classification and compact representation learning.

• To support our broader objectives, we introduce several novel techniques:

  • Match a discrete distribution (i.e., the anchors) to a continuous one, as opposed to a continuous-to-continuous distribution matching in the referenced paper.
  • Use an adaptive latent space, since our encoder and decoder are trained end-to-end, as opposed to using a fixed pre-trained autoencoder and fixed latent space in the referenced paper.
  • Numerically solve the flow directly, as opposed to training an additional model to learn the flow in the referenced paper.
  • Latent alignment between different data types (e.g., image and label), which is new in our paper.

That said, we agree that the referenced work is relevant and should be discussed. We will add the above points to the revision.

---

---

Q2. Scalability

We acknowledge the computational cost of our approach, as discussed in Section 2 and Section 6. However, we would like to highlight several key points:

• The computation cost, while non-trivial, is comparable to many existing techniques.

  • The main bottleneck stems from the quadratic cost with respect to the batch size (i.e., the $n^2$ term). Notably, this is also the case for many contrastive learning methods, including the seminal works MoCo [25] and SimCLR [8], which compute pairwise similarities between all examples in a batch.
  • As discussed in Appendix F, the computational profile of LZN is similar to that of large language models (LLMs): LZN scales quadratically with the number of samples, whereas LLMs scale quadratically with the number of tokens. Given that modern LLMs often process millions of tokens, the computational cost of LZN is relatively modest in comparison.

• We propose several optimization strategies to improve LZN’s efficiency, as detailed in Appendix A.3. With these optimizations, the training cost remains manageable across all our experiments.

• At inference time, LZN often does not incur this high computational cost.

  • For image generation (Case Studies 1 and 3), we do not need to compute the latent during inference. Instead, latents are sampled from the Gaussian prior and passed directly to the decoder, making the generation speed comparable to the base model. This is explained in Lines 269–271.
  • For representation learning (Case Study 2), we find that dropping the final encoder layers during inference—similar to several contrastive learning methods (e.g., [8])—not only improves performance but also eliminates the $O(n^2)$ cost. In this case, inference involves simply passing an image through the encoder, just like in traditional contrastive learning methods. This is discussed in Lines 1187–1192.

---

---

Q3. Mini-batch approximation

Thank you for the great question!

• Mini-batch approximation still preserves the Gaussian prior.
  LZN ensures that the latent distribution within each mini-batch is approximately $\mathcal{N}(0, I)$. As a result, the overall latent distribution becomes a mixture of Gaussians with the same parameters, which is still $\mathcal{N}(0, I)$. Therefore, the global prior remains valid under the mini-batch approximation.

• However, this approximation does violate the disjoint zone property.
  To quantify this effect, we conducted the following experiment:
  We randomly sample one instance $x$ from the dataset and construct two batches with other randomly sampled instances: $\\{x, x_2, \dots, x_n\\}$ and $\\{x, x_2', \dots, x_n'\\}$, where $n$ is the batch size. We then compute the latent for $x$ using each batch and evaluate how often the latent of $x$ from one batch lies within the latent zone of $x$ in the other batch. The result shows that only 49% of the time does the latent remain in its intended zone, implying a 51% chance that it overlaps with other zones.

• Despite this limitation, the mini-batch approximation performs well in practice.

  • In generative modeling (Case Studies 1 and 3), the latent does not need perfect zone disjointness—as long as it provides some information about the input sample, it can help reduce the variance needed to learn by the generative model (rectified flow in our case) and improve the generation quality.
  • In representation learning (Case Study 2), latent alignment occurs within a single batch. Thus, inconsistency across batches is irrelevant.
  • In classification (Case Study 3), we only need to map samples within a single batch to the latent zones of labels. Thus, inconsistency across batches is irrelevant.

Given these points, we observe strong empirical results across all three case studies despite using the mini-batch approximation.

That said, developing a better method that is both efficient and preserves the disjoint latent zone property might further improve LZN’s performance. We leave this as an exciting direction for future work.

---

---

Q4. Performance on classification tasks

Thank you for the question.

• For the classification experiments in Case Study 2, we would like to highlight that our results have substantially improved since the submission. By extending training to 4.5 million iterations, the top-1 accuracy improves to 69.3% (+3.7%) and the top-5 accuracy improves to 89.1% (+2.0%). With these updated results, LZN is now among the top 4 state-of-the-art methods for this benchmark. Importantly, these competing methods are designed exclusively for representation learning, whereas LZN supports more tasks including generation. Notably, LZN also matches or outperforms seminal methods such as SimCLR. We believe these are strong results.

• For the classification experiments in Case Study 3:

  • • Our method is only 1% below the best method, even though LZN is designed to support both classification and generation simultaneously within a unified model—in contrast to the competing methods that are tailored solely for classification.
  • We believe this gap is largely due to architecture design. The competing methods use architectures explicitly tailored for classification, while we use a standard diffusion model architecture without any classification-specific modifications (see Appendix E.2). This is supported by the fact that using our architecture solely for classification yields only 93.59% accuracy (Table 3).
  • While adopting a better architecture could potentially close this gap, the more important takeaway is that joint training for classification and generation via LZN leads to better performance than training each task in isolation (Lines 354–358).

---

---

Q5. Soft alignment

The only hyperparameter introduced in soft alignment is u, which determines how many time steps are excluded from the latent alignment objective (Line 234). In our experiments, we set u = 20 (out of a total of 100 steps) across all tasks.

Following your suggestion, we conducted an ablation study in Case Study 3 by reducing u to 5 (i.e., 4× smaller). The results are shown below:

We can see that this new parameter u does not affect the results much, and the performance remains better than the baseline RF across most metrics. This is expected. Unlike common hyperparameters (such as loss weights) that influence the optimal solution, u does not alter the optimal solution, which is the perfect alignment between two latent zones. Instead, this parameter is introduced solely to help avoid getting stuck in local optima (as discussed in Lines 231–236). We expect that any small but non-zero value of u should be sufficient in practice.

---

---

Q6. Claims about generality; more tasks.

We tried training on two tasks simultaneously in Case Study 3, where we did not observe any performance degradation. In fact, we saw an improvement compared to training on each task separately (Lines 354–358).

Applying LZN to a broader set of tasks is indeed an important direction for future work, as we also discussed in Section 6.

We appreciate the suggestion and will refine the statements in the paper accordingly.

---

---

Q7. Table alignment

Thank you! We will fix these in the revision.

---

---