PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zinan Lin, Enshu Liu, Xuefei Ning, Junyi Zhu, Wenyu Wang, Sergey Yekhanin

arXiv:2509.15591v2cs.LGcs.AIcs.CVstat.ML

TL;DR

The paper asks whether one principle can unify generative modeling, representation learning, and classification, whose state-of-the-art approaches use distinct methods. It proposes LZN, a shared latent-space framework, and reports results across generation, representation learning, and joint generation and classification.

  • Problem

    The paper asks whether a single principle can unify generative modeling, representation learning, and classification, which use distinct state-of-the-art methods.

    Source: §1 · p. 2

  • Method

    LZN connects data-type encoders and decoders through a shared latent space, using latent computation and soft alignment to form and align disjoint zones.

    Source: §1 · pp. 2–3

  • Results

    Across three application settings, LZN improves existing-model generation, outperforms MoCo and SimCLR in ImageNet linear classification, and improves generation while achieving state-of-the-art classification accuracy jointly.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The results support the paper's motivating intuition that seemingly distinct tasks can benefit from shared representations.

    Source: §1 · p. 3

    Whether LZN can scale gracefully as encoder–decoder pairs are added remains to be seen.

    Source: §1 · p. 3

Abstract

from arXiv · show

Generative modeling, representation learning, and classification are three core problems in machine learning (ML), yet their state-of-the-art (SoTA) solutions remain largely disjoint. In this paper, we ask: Can a unified principle address all three? Such unification could simplify ML pipelines and foster greater synergy across tasks. We introduce Latent Zoning Network (LZN) as a step toward this goal. At its core, LZN creates a shared Gaussian latent space that encodes information across all tasks. Each data type (e.g., images, text, labels) is equipped with an encoder that maps samples to disjoint latent zones, and a decoder that maps latents back to data. ML tasks are expressed as compositions of these encoders and decoders: for example, label-conditional image generation uses a label encoder and image decoder; image embedding uses an image encoder; classification uses an image encoder and label decoder. We demonstrate the promise of LZN in three increasingly complex scenarios: (1) LZN can enhance existing models (image generation): When combined with the SoTA Rectified Flow model, LZN improves FID on CIFAR10 from 2.76 to 2.59-without modifying the training objective. (2) LZN can solve tasks independently (representation learning): LZN can implement unsupervised representation learning without auxiliary loss functions, outperforming the seminal MoCo and SimCLR methods by 9.3% and 0.2%, respectively, on downstream linear classification on ImageNet. (3) LZN can solve multiple tasks simultaneously (joint generation and classification): With image and label encoders/decoders, LZN performs both tasks jointly by design, improving FID and achieving SoTA classification accuracy on CIFAR10. The code and trained models are available at https://github.com/microsoft/latent-zoning-networks. The project website is at https://zinanlin.me/blogs/latent_zoning_networks.html.

1 Introduction

LZN asks whether one principle can unify generative modeling, representation learning, and classification, and proposes a shared latent space connecting data-type encoders and decoders. Its initial applications span enhancing generation, learning representations, and jointly performing generation and classification.

1 Introduction

  • The paper questions whether separate state-of-the-art methods for generation, representation learning, and classification are necessary, asking whether one principle can unify them.The motivation is that these tasks differ in formulation but are fundamentally related and may benefit from one another.Source: §1 · p. 2
  • LZN connects data-type-specific encoders and decoders through a shared latent space, with tasks formed from their combinations.Encoders map data into latent zones, and corresponding decoders map latents back to data.Source: §1 · p. 2
  • Latent computation creates disjoint sample zones with a Gaussian aggregate, while soft latent alignment matches zones across encoders for cross-data-type translations.The alignment approximation operates midway through flow matching to make training effective and tractable.Source: §1 · p. 3
  • Across three application levels, LZN improves existing-model generation, outperforms MoCo and SimCLR in ImageNet linear classification, and improves generation while achieving state-of-the-art classification accuracy jointly.For joint generation and classification, performance on both tasks exceeds training each task in isolation.Source: §1 · p. 3
  • The authors leave open whether LZN can scale gracefully as encoder–decoder pairs are added, despite the framework's projected growth in supported applications.They state that the range could grow at least quadratically in principle, but whether this potential can be realized remains to be seen.Source: §1 · p. 3

2 Latent Zoning Network (LZN)

LZN frames diverse machine-learning tasks as mappings between data and a shared latent space designed to capture task information while retaining a simple generative distribution. Its encoders, decoders, and alignment operations support both within- and cross-data-type tasks.

2.1 Overall Framework

  • Existing generative, representation-learning, and classification methods impose different latent constraints and can discard information needed by other tasks, motivating a unified mapping framework.Generative models can be difficult to invert, while contrastive and classification objectives may discard details or lack distributional constraints.Source: §2.1 · p. 4
  • LZN assigns each observed sample a latent zone in a shared space, enforces a Gaussian latent distribution, and provides encoders and decoders for mappings in both directions.Different data types can share the same underlying latent entity even when their observations differ.Source: §2.1 · p. 4
  • An encoder or decoder can perform single-module tasks, while encoder–decoder pairs compose cross-module tasks such as label-conditional generation and image classification.Image encoding yields embeddings; image decoding supports unconditional generation, and text encoders and decoders support text generation.Source: §2.1 · pp. 4–5

2.2 Implementation of Atomic Operations

  • LZN’s two atomic operations, latent computation and latent alignment, are presented as sufficient to implement a variety of tasks.The framework also treats existing generative models as possible LZN decoders.Source: §2.2 · p. 5, §2.3 · p. 7

2.2.1 Latent Computation

  • Latent computation maps encoder-produced anchor points through flow matching to latent zones intended to follow a Gaussian prior while remaining disjoint across samples.The computation is differentiable, allowing gradients to reach the encoders; numerical solvers approximate its integral.Source: §2.2.1 · p. 5

2.2.2 Latent Alignment

  • Because flow matching assigns each latent deterministically to an anchor, LZN uses a soft assignment approximation and maximizes assignment probability over later trajectory steps to align zones.The early-step cutoff addresses the lack of training signal that can arise when the maximum is attained at the initial uniform assignment.Source: §2.2.2 · p. 6

2.3 Decoder

  • The decoder maps an LZN latent back to its corresponding sample and can be implemented using a generative model or flow matching.The framework describes tasks through compositions of the corresponding data-type encoders and decoders.Source: §2.1 · p. 5, §2.3 · p. 7

2.4 Relationships to Alternatives

  • Unlike generation-centric autoregressive approaches that use surrogate methods for representation learning, LZN directly formulates generation, representation learning, and classification within one framework.The paper presents image-domain case studies spanning generation, unsupervised representation learning, and joint classification and generation; it describes LZN as complementary to generative models.Source: §2.4 · p. 7, §2.5 · p. 7

3 Case Study 1: Unconditional Generative Modeling

Adding LZN latents as a conditioning input to rectified flow improves image-generation quality and reduces reconstruction error across four evaluated datasets. On CIFAR10, it reduces the FID gap between unconditional and conditional generation by 59%.

Approach

  • LZN latents are supplied as an additional generator input, with the encoder and decoder trained jointly using the original generative loss.The encoder computes target-image latents during training; at inference, Gaussian-prior samples are passed to the decoder, so the encoder is not used.Source: §3 · p. 7

Why it helps

  • Unique image latents encourage the generator to depend more directly on the latent, making the conditional image distribution more deterministic and the objective easier to optimize.The paper reports that experiments confirm these latents capture useful features for generation.Source: §3 · p. 7

Related work

  • Unlike RCG and Diffusion Autoencoder, LZN’s Gaussian-constrained representations avoid requiring an auxiliary generative model for unconditional generation.The paper describes LZN as an end-to-end, single-stage approach.Source: §3 · p. 8

Results

  • 59%: LZN reduces the CIFAR10 FID gap between unconditional and conditional generation, while also outperforming conditional RF in sFID and reconstruction error.Across CIFAR10, AFHQ-Cat, CelebA-HQ, and LSUN-Bedroom, LZN improves image quality and reduces reconstruction error; the table reports unconditional image-generation quality scores across the four datasets.Source: §3 · p. 8, Table 1 · p. 8

4 Case Study 2: Unsupervised Representation Learning

LZN trains an image encoder by aligning latent zones for two random augmentations of each image, without contrastive losses or decoders. On ImageNet linear evaluation, it outperforms MoCo and SimCLR but remains below state of the art.

Approach

  • The encoder is trained by aligning latent zones computed for two randomly augmented views of each image, using the same encoder for both views.The approach needs neither decoders nor contrastive losses that push dissimilar samples apart; distinct latent zones are intended to prevent collapse by design.Source: §4 · pp. 8–9

Related work

  • Contrastive methods address collapse by separating dissimilar samples or using memory banks or specialized architectures, whereas LZN assigns different images distinct latent zones.The related-work discussion also identifies masked image modeling and auxiliary-task training as other unsupervised approaches.Source: §4 · p. 8

Approach

  • At inference, representations can be LZN latents, encoder anchor points, or intermediate encoder outputs; the latter two avoid costly latent computation.The encoder outputs used for training come from two sets of augmentations of each image batch.Source: Figure 6 · p. 9

Results

  • LZN outperforms MoCo by 9.3% and SimCLR by 0.2% on ImageNet linear classification, but remains below state-of-the-art performance.Representations are learned on unlabelled ImageNet with ResNet-50; the linear classifier is trained in a supervised manner, without pretraining or external data.Source: §4 · p. 9

5 Case Study 3: Conditional Generative Modeling and Classification

LZN adds label encoders and decoders to the image encoder-decoder framework, placing label inputs and outputs in a shared latent space. On CIFAR10, it jointly supports conditional and unconditional generation and classification, with reported gains in generation quality and classification accuracy.

Approach

  • The label encoder maps one-hot labels to shared anchor points, while the label decoder maps latents to class IDs using flow matching and the nearest anchor.The training objective adds Align(X, Y); after training, the model can perform conditional and unconditional generation as well as classification.Source: §5 · p. 10

Related work

  • Unlike approaches that treat generator inputs and classifier outputs as separate components, LZN unifies label inputs and outputs in a shared latent space.The label encoder-decoder pair is added alongside the image encoder and decoder.Source: §5 · p. 10

Results

  • Joint training improves both generation and classification: RF+LZN improves generation quality over § 3, and beats RF+LZN (no gen) in classification accuracy.LZN also improves image quality and reconstruction error, achieves classification accuracy on par with state of the art, and trails the best CIFAR10 result by 1%.Source: §5 · p. 10

6 Limitations and Future Work

The paper identifies training cost, remaining performance gaps, and limited modality and task coverage as open directions. Its appendix explains how latent computation approximately achieves a Gaussian prior and disjoint zones.

6 Limitations and Future Work

  • LZN training is computationally expensive because it backpropagates through the flow-matching trajectory; performance still trails state of the art, and experiments focus mainly on images and at most two simultaneous tasks.The authors leave efficiency, performance improvements, standalone generative modeling, and broader multimodal or multitask applications for future work.Source: §6 · p. 10

A.1 The Desired Properties

  • Latent computation yields a Gaussian prior, while the probability of assigning a sample to an incorrect zone can be made arbitrarily small by choosing a sufficiently small noise scale.The appendix illustrates the error probability with a one-dimensional, two-sample case.Source: §A.1 · p. 18

A.2 More Implementation Details

The appendix describes a concentrated latent option and efficiency optimizations for scaling LZN. Minibatch approximation preserves the Gaussian prior but sacrifices globally disjoint latent zones.

A.2 More Implementation Details

  • A scaling factor α in [0, 1) controls the noise applied during latent computation to obtain a more concentrated latent distribution.The appendix notes similar techniques in prior generative models for improving sample quality.Source: §A.2 · p. 18

A.3 Efficiency Optimizations

  • Minibatch approximation reduces memory and computation costs, preserves the Gaussian prior, but violates disjointness across batches because zones depend on batch membership.The authors report that it performs well: cross-batch inconsistency is irrelevant for representation alignment and classification within a batch, while larger batches can improve label-zone accuracy.Source: §A.3 · pp. 18–19
  • Custom gradient checkpointing reduces latent-computation memory complexity from O(n^2qr) to O(nqr) by retaining latent trajectories and recomputing velocity computations.The trade-off is additional computation during backpropagation.Source: §A.3 · p. 19
  • Latent parallelism partitions samples across GPUs, then broadcasts anchor points so each GPU can compute latents using the full set of anchors.Gradient flow to the originating GPUs is maintained with torch.distributed.nn.functional.all_gather.Source: §A.3 · p. 19

B.1 More Implementation Details

These implementation details describe how LZN is integrated into Rectified Flow (RF) training and generation, including changes to the decoder and encoder architectures.

B.1 More Implementation Details

  • The alignment objective may optionally apply a logarithm to the assignment probability to resemble standard cross-entropy.The passage introduces the resulting objective but does not provide its equation here.Source: §B.1 · p. 19

B.2 Efficiency Optimizations

  • The latent-alignment procedure uses the same efficiency optimizations as the earlier method.The referenced optimization details are not included in this passage.Source: §B.2 · p. 19

C.1 Algorithm Pseudocode

  • The RF and RF+LZN training and generation processes are illustrated side by side.The training illustration emphasizes computing LZN latents; the generation illustration shows additional sampled latents as decoder inputs.Source: §C.1 · p. 19, Figure 7 · p. 20, Figure 8 · p. 20
  • RF+LZN training computes LZN latents for samples and supplies them as an additional input to the RF decoder.The training-process comparison highlights this as a key difference from standard RF.Source: Figure 7 · p. 20
  • During iterative generation, RF+LZN samples extra LZN latents from Gaussian noise and uses them as additional inputs to the RF decoder.The figure contrasts this with standard RF, which starts from Gaussian noise.Source: Figure 8 · p. 20

C.2 More Implementation Details

  • The RF decoder concatenates the LZN latent with the timestep embedding, while its encoder adds latent transformation blocks to each ResNet block.Each block uses a 1×1 convolution projecting to 20 channels, then an MLP with a 200-dimensional hidden layer.Source: §C.2 · p. 20

C.3 More Experimental Settings

The experiments compare RF and RF+LZN on four image datasets, using shared evaluation samples and specified training and sampling settings.

C.3 More Experimental Settings

  • Experiments use CIFAR10, AFHQ-Cat, CelebA-HQ, and LSUN-Bedroom, with the stated image sizes and dataset counts.CIFAR10 has 50,000 training images; AFHQ-Cat has 5,153 images; CelebA-HQ has 30,000; and LSUN-Bedroom has 3,033,042.Source: §C.3 · pp. 20–21
  • FID, sFID, IS, precision, recall, and CMMD compare generated images with real training-set images, while reconstruction error measures image reconstruction.Metrics use consistent random subsamples for RF and RF+LZN; reconstruction error averages 20 reconstructions per sampled image.Source: §C.3 · p. 21
  • For FID, sFID, IS, precision, and recall, each model generates as many samples as the selected training-set subsample.The counts are 50,000 for CIFAR10, 5,120 for AFHQ-Cat, and 29,952 for each of CelebA-HQ and LSUN-Bedroom.Source: §C.3 · p. 21
  • The main experiments use the adaptive RK45 sampler, with Euler sampling steps varied in a separate analysis.The reported RF and RF+LZN settings specify batch sizes, Adam optimization, learning rates, and gradient clipping; RF+LZN uses a 200-dimensional latent.Source: §C.3 · pp. 21–23
  • RF+LZN training takes 8 hours on 16 A100 GPUs for CIFAR10 and 10, 58, and 341 hours on 32 A100 GPUs for AFHQ-Cat, CelebA-HQ, and LSUN-Bedroom, respectively.These costs exclude periodic evaluation.Source: §C.3 · p. 23

C.4 More Results

Additional results examine RF+LZN image generation across datasets, FID implementations, and sampling-step settings, alongside implementation and representation-learning details.

C.4 More Results

  • Generated-image examples for RF and RF+LZN are provided in Figures 9–16.The passages identify the image sets but do not describe a specific visual quality comparison.Source: §C.4 · p. 23
  • RF+LZN achieves the best FID on three of four datasets, and its relative ranking remains consistent across three FID implementations.The table covers unconditional image generation; it identifies consistency models, RF, and clean FID implementations.Source: §C.4 · p. 23, Table 5 · p. 24
  • RF+LZN generally achieves better FID than RF across most tested Euler sampling-step settings, with its only stated underperformance at the highest step count.Figure 17 presents FID against the number of Euler sampling steps.Source: §C.4 · pp. 23–24, Figure 17 · p. 28

D.1 Algorithm Pseudocode

  • The RF+LZN training pseudocode is provided, and the trained encoder can produce image representations after training.The passage points to additional representation-extraction strategies without specifying them.Source: §D.1 · p. 29

D.2 More Implementation Details

  • The representation-learning setup uses a ResNet-50 encoder with group normalization and a one-hidden-layer MLP projection head.The ResNet-50 backbone, without the projection head, supplies representations for linear-classifier training without LZN latent computation.Source: §D.2 · p. 29
  • The representation-learning objective uses the version with a logarithm, and data augmentation follows the cited prior method.The cited passage reports no significant performance difference from replacing batch normalization with group normalization in early experiments.Source: §D.2 · p. 29

D.3 More Experimental Settings

  • LZN is trained on ImageNet’s training set and evaluated by classification accuracy on its validation set.The experiment uses 1,281,167 training images and 50,000 validation images; training takes 1,800 hours on 128 A100 GPUs, excluding periodic evaluation.Source: §D.3 · p. 29

D.4 More Results

Further results examine representation choices, training duration, and visualized ImageNet embeddings, while the accompanying sections describe joint generation and classification settings.

D.4 More Results

  • Table 6 reports ImageNet classification accuracy from linear classifiers trained on unsupervised representations, including contrastive-learning methods.The caption says its horizontal line separates baselines performing worse or better than LZN, but the supplied text gives no cell values.Source: §D.4 · p. 29, Table 6 · p. 30
  • In a t-SNE projection of embeddings from 20 randomly selected ImageNet validation classes, samples from different classes are well clustered.The authors say this suggests LZN learns meaningful image representations; the validation set was chosen to avoid training-set overfitting.Source: §D.4 · pp. 29–30, Figure 18 · p. 31
  • The LZN latent representation has the lowest prediction accuracy in the feature-extraction comparison, while setting α = 0 yields significantly higher accuracy.The authors attribute latent underperformance to minibatch approximations that enlarge zones and can cause overlap; the “without head” features also substantially outperform “with head.”Source: §D.4 · pp. 30–31
  • Classification accuracy continues improving rapidly at the end of training, suggesting more training might further reduce the gap to the SoTA.Figure 20 plots accuracy against training iteration.Source: §D.4 · p. 31, Figure 20 · p. 32

E.1 Algorithm Pseudocode

  • The paper provides pseudocode for joint training, generation, and classification.The passage names these procedures without describing their steps.Source: §E.1 · p. 32

E.2 More Implementation Details

  • For conditional generation, the class label is supplied as a condition and LZN latents are sampled from the corresponding class zone.For unconditional generation, the class label is decoded from the LZN latent; batch label imbalance is handled with a weighted mixture during training, then a uniform prior is used at testing.Source: §E.2 · pp. 32–33

E.3 More Experimental Settings

  • The joint-generation experiments use CIFAR10 test accuracy and the generation metrics described earlier, with training taking 31 hours on 16 A100 GPUs.The objective uses the version without a logarithm; the authors report no large difference when they tried the log version.Source: §E.3 · pp. 33–34, §E.2 · p. 33
  • Table 7 reports conditional CIFAR10 generation FID across consistency-model, RF, and clean-FID implementations.The caption identifies the metric and implementations but supplies no result values.Source: Table 7 · p. 34

E.4 More Results

E.4 tests how generation and classification results respond to implementation and hyperparameter choices. The related-work discussion distinguishes LZN’s broader aims and design from prior latent-space flow matching.

E.4 More Results

  • RF+LZN achieves the best FID under all three FID implementations, despite differences in the reported numbers.Source: §E.4 · p. 34
  • RF+LZN generally achieves better FID than RF across most Euler sampling-step settings on CIFAR10.The comparison is shown in Fig. 24.Source: §E.4 · p. 34, Figure 24 · p. 37
  • 2.9%: the best classification setting improves on the default batch size=2000 and α=1.0, with larger batches and lower α generally improving accuracy.Fig. 25 shows the hyperparameter trends.Source: §E.4 · p. 35, Figure 25 · p. 37
  • Reducing latent-alignment u from 20 to 5 has little effect, and performance remains better than baseline RF across most metrics.The ablation results are in Table 8; u is intended to help avoid local optima rather than change the optimal alignment.Source: §E.4 · pp. 35–36, Table 8 · p. 37

F Extended Discussions on Related Work

  • Unlike [16], LZN targets generation, representation learning, and classification, using discrete-to-continuous matching, an adaptive latent space, direct numerical flow solving, and cross-type latent alignment.[16] instead applies flow matching to a fixed, pretrained Stable Diffusion autoencoder latent space.Source: F Extended Discussions on Related Work · p. 38

G Extended Discussions on Limitations and Future Work

The discussion notes a quadratic batch-size training bottleneck, while inference is often comparable to existing approaches. It also identifies a failure on complex datasets for one explored approach and potential beneficial and harmful uses.

G Extended Discussions on Limitations and Future Work

  • For image generation, inference samples from the Gaussian prior without latent computation; representation inference passes images through the encoder, with final-layer dropping improving performance.These inference paths are described as comparable in speed to the base model or traditional contrastive learning methods.Source: G Extended Discussions on Limitations… · p. 38
  • LZN training has a quadratic batch-size cost, also seen in MoCo and SimCLR; its latent computation has structural parallels to LLM attention that may inform efficiency improvements.The stated latent-computation complexity is O(n^2qr), and the paper expects LLM training advances could be adapted to LZN.Source: G Extended Discussions on Limitations… · p. 38
  • Empirical-distribution matching performs well on MNIST but produces blurry CIFAR10 images; the authors hypothesize minibatch approximation or gap-free latent zones may contribute.They leave deeper exploration of this direction to future work.Source: G Extended Discussions on Limitations… · p. 39
  • LZN may support beneficial applications such as creative content generation and information retrieval, but also harmful creation of fake or misleading content.Source: G Extended Discussions on Limitations… · p. 39

Related papers

Loading 2509.15591v2…