PaperlayerView exact revision on arXiv

Source-linked AI summary

Cola2: Co-Adaptive Hierarchical Language Modeling with Continuous Semantic Compression and Balanced Discrete Generation

Hongcan Guo, Yian Zhao, Feng Wang, Shen Nie, Guoqiang Wei, Yan Zeng, Hengshuang Zhao

arXiv:2610.31305v1cs.CL

TL;DR

Hierarchical language generation must determine what semantic information latents retain and how discrete decoders use it, while coordinating representation and prior learning. Cola2 combines dynamic semantic compression and balanced discrete generation with co-adaptive training; experiments report competitive, unsaturated scaling and a 60.10 average across eight tasks in a separate large-token run.

  • Problem

    Hierarchical modeling must determine what information latents compress and how decoders use it, since shorter sequences or reconstruction-only objectives do not ensure predictive semantic representations.

    Source: §1 · p. 2

  • Method

    Cola2 jointly learns dynamically compressed semantic latents and a discrete masked-diffusion decoder balanced between reconstruction and autonomous generation, while decoupling prior fitting from encoder updates.

    Source: §1 · p. 2

  • Results

    With a 346.03B-source-token run, Cola2 reaches a 60.10 average across eight tasks and surpasses Llama-3-8B on HellaSwag and StoryCloze despite over 43× fewer reported tokens.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Experiments support Cola2's design principles and indicate potential for continuous-latent hierarchical generation.

    Source: §1 · p. 3, §6 · p. 20

    The system remains far from convergence, especially the latent DiT, and current training scale is limited by available resources.

    Source: §5 · p. 20

Abstract

from arXiv · show

Hierarchical language generation separates continuous semantic modeling from discrete text realization, but its effectiveness depends on what the latent space compresses and how the discrete generator uses it. We propose Cola2, a framework that jointly studies continuous semantic compression, the balance between reconstruction and generation, and continuous-discrete co-adaptation. Cola2 dynamically segments text into semantic spans, controls the compression rate through auxiliary-loss-free negative feedback, and combines reconstruction and generation supervision to learn compact, predictive latent representations. A block-causal masked diffusion decoder is trained with clean, noisy, and dropped latent conditions, while classifier-free guidance adjusts the balance between semantic reconstruction and autonomous generation. Token-prediction gradients update the encoder through the latents, whereas Flow Matching prior learning is decoupled from representation learning by stop-gradient. During inference, generated text is re-encoded under a fixed span layout to keep the latent history consistent with the actual text context. Experiments on five downstream tasks with about 210 billion training tokens validate the core designs and show competitive, unsaturated scaling. With 346.03 billion source tokens, an unconverged checkpoint reaches a 60.10 average across eight tasks and surpasses Llama-3-8B on HellaSwag and StoryCloze despite its reported pretraining budget being over 43 times larger. These results demonstrate the potential of continuous-latent hierarchical language generation.

1 Introduction

Cola2 addresses the gap between shortening latent sequences and learning compact semantic representations that predict missing or future text. It jointly designs semantic compression, discrete reconstruction and generation, and continuous–discrete co-adaptation.

1 Introduction

  • The latent space must retain information useful for reconstruction and prediction, while the discrete generator completes residual uncertainty when latent conditions are incomplete or imprecise.Shorter latent sequences alone do not ensure semantic compression, and reconstruction-only training does not guarantee prediction of missing or future text.Source: §1 · p. 2
  • Cola2 frames hierarchical modeling as deciding which objectives should jointly shape representations and which learning processes should remain decoupled.A fixed encoder can limit adaptation to decoding needs, while direct prior-fitting gradients may favor latents that are easier to fit rather than more useful for discrete generation.Source: §1 · p. 2
  • Cola2 dynamically segments text into semantic spans, controls compression through auxiliary-loss-free negative feedback, and combines reconstruction and generation supervision.Its block-causal masked diffusion decoder trains with clean, noisy, and dropped latent conditions; classifier-free guidance adjusts the reconstruction–generation balance at inference.Source: §1 · p. 2
  • Token-prediction gradients flow through latents to the encoder, while Flow Matching prior fitting is stop-gradient-decoupled from representation learning.At inference, generated text is re-encoded block by block under the fixed span layout to keep latent history consistent with actual text context.Source: §1 · p. 2
  • The framework analyzes latent information and residual discrete generation through conditional information decomposition.Its stated focus is semantic compression, reconstruction–generation balance, and continuous–discrete co-adaptation.Source: §1 · p. 2
  • With a 346.03B-source-token run, Cola2 reaches a 60.10 average across eight tasks and surpasses Llama-3-8B on HellaSwag and StoryCloze despite over 43× fewer reported tokens.The controlled five-task scaling experiments use about 210B training tokens and show competitive, unsaturated scaling.Source: §1 · p. 3

2 Related Work

Related work covers fixed, adaptive, and concept-level semantic compression, token-space generation paradigms, and continuous latent-space modeling. Cola2 distinguishes itself by using negative feedback for compression and jointly optimizing representations with discrete generation supervision while isolating prior-fitting gradients.

2.1 Semantic Compression and Hierarchical Language Models

  • Cola2 maintains a target compression rate through negative-feedback control without an auxiliary compression-rate loss, alongside reconstruction and generation supervision for compact, predictive semantic representations.The section contrasts fixed-granularity, adaptive compression, and concept-level approaches.Source: §2.1 · p. 3

2.2 Token-Space Language Models

  • Unlike token-space methods, Cola2 models a semantic prior in compressed latent space before conditional generation with a discrete decoder.The section contrasts autoregressive serial generation, iterative discrete diffusion, and continuous-augmented denoising methods.Source: §2.2 · p. 3

2.3 Continuous Representations and Latent-Space Language Models

  • Cola2 uses discrete reconstruction and generation supervision to shape representations while isolating prior-fitting gradients, jointly optimizing semantic compression and decoder balance.The section distinguishes token-aligned continuous modeling, autoencoder-based latent methods, and joint-learning methods.Source: §2.3 · p. 3

3 Cola2: Continuous Latent Diffusion for Hierarchical Text Modeling

Cola2 formulates hierarchical text generation as continuous latent modeling followed by discrete token realization, and develops methods for semantic compression, co-adaptation, training, inference, and cost analysis.

3.1.1 Hierarchical Probabilistic Modeling

  • The encoder allocates limited latent capacity between information retained in Z and residual uncertainty handled by the discrete generator.Cola2 factorizes text generation into a latent prior and conditional token generation.Source: §3.1.1 · pp. 3–4

3.1.2 Semantic Compression and Information Density

  • Under compression, latent quality depends on whether each position retains semantic and predictive information, not only on sequence length.The paper treats semantic-information density as a design objective, not a directly estimated training loss.Source: §3.1.2 · p. 4

3.1.3 Reconstruction--Generation Balance and Continuous--Discrete Co-adaptation

  • Cola2 jointly trains the encoder and decoder so token prediction shapes latent information allocation, while latent-prior fitting remains decoupled.This co-adaptation lets representations suit the actual decoder without letting prior fitting directly update the encoder.Source: §3.1.3 · p. 5

3.2 Workflow

  • The workflow first jointly learns the encoder and decoder, then adds a DiT trained on detached online latents.At inference, the DiT proposes latents block by block and the decoder realizes them as tokens.Source: §3.2 · p. 5, Figure 1 · p. 6

3.2.1 Model Architecture

  • The encoder dynamically segments tokens into spans and forms continuous latents; a block-causal DiT models latents and an MDM decoder realizes tokens.The compression rate is maintained near its target by an independent gradient-free controller rather than an auxiliary differentiable rate loss.Source: §3.2.1 · pp. 5, 7, §3.2 · p. 5

3.2.2 Training Procedure and Objectives

  • Stage 1 trains the decoder with clean, noisy, or dropped latent conditions, spanning conditioned reconstruction, robustness, and weak-condition generation.Token-prediction gradients update the encoder through the latents; Stage 2 adds detached-latent Flow Matching prior fitting.Source: §3.2.2 · pp. 7–8

3.2.3 Inference Procedure

  • At inference, generated text is re-encoded blockwise under the fixed span layout, aligning latent history with the actual token context.Feedback does not rerun the dynamic segmenter; the persistent latent history uses corrected latents.Source: §3.2.3 · p. 8

3.3 Computational Cost

  • The paper’s computational analysis is a forward-round proxy, not a latency guarantee.End-to-end cost also depends on CFG branches, token decoding, and feedback.Source: §3.3 · p. 9

4 Experiments

Across component analyses and scaling experiments, Cola2’s results support content-aware compression, generation-oriented supervision, and coordinated training of the latent and discrete generator. Its post-pretraining checkpoint reaches a 60.10 average across eight tasks before convergence, while exceeding Llama-3-8B on HellaSwag and StoryCloze despite using over 43 times fewer reported pretraining tokens.

4.2.1 Dynamic Segmentation and Rate Control

  • 29.4% at ρ = 1/2 and Lmax = 8 exceeds the uncompressed model’s 28.7%; dynamic segmentation at Lmax = 4 reaches 25.0%, beating random and fixed-length spans by 4.4 and 8.2 points.Feedback control reaches 25.0%, versus 18.4% for DLCM [39].Source: §4.2.1 · pp. 10–11, Figure 2 · p. 10

4.2.3 Compression and Latent Capacity

  • 32.3% is the peak at dz = 64, falling to 14.5% at dz = 512; stronger compression also lowers accuracy from 32.3% at span length 2 to 25.2% at span length 4.The effect of Lmax depends on truncation: with about 1% cap mass at span length 4 and Lmax = 18, accuracy rises to 21.7%.Source: §4.2.3 · pp. 11–12, Table 1 · p. 12, Figure 5 · p. 12

4.2.4 Generative Supervision of the Encoder

  • 58.2% with an MDM decoder and generative loss exceeds 53.5% from end-to-end co-evolution alone and 47.6% with a frozen encoder; pure reconstruction reaches only 15.7%.The BERT-style objective reaches 35.2% with an unfrozen encoder, compared with 31.1% when frozen.Source: §4.2.4 · p. 12, Figure 6 · p. 13

4.3.1 VAE Update Policy and Gradient Detachment

  • 32.8% with a trainable VAE exceeds 28.7% with a frozen VAE; detaching the DiT gradient ultimately performs better than allowing it into the encoder.Gradient injection leads early but is 1.0 point behind at 140k steps; all three noise-schedule configurations peak at loc = 2.0.Source: §4.3.1 · p. 13, Figure 7 · p. 13, §4.3.2 · p. 14

4.3.3 Block Size and Generator Capacity

  • For a 2B DiT, block size B = 4 reaches 36.4%, versus 31.8% at B = 1 and 27.3% at B = 24; larger DiTs have flatter declines above B = 8.At 8 steps, generator accuracy increases from 12.2% at γgen = 1 to 37.1% at γgen = 14; results from 8 to 128 steps differ by at most 0.3 points.Source: §4.3.3 · p. 14, Figure 9 · p. 14, §4.3.4 · p. 15

4.4.2 Decoder Data Partition

  • The selected 0.4/0.4/0.2 drop/noisy/clean mixture ties for the best LAMBADA accuracy at 44.7% while retaining more clean data; reconstruction accuracy stays above 99.5% throughout.Source: Figure 10 · p. 15, §4.4.2 · p. 16

4.5.2 Post-Pretraining Benchmark Comparison

  • 60.10 across eight tasks: the unconverged Cola2 checkpoint trails Llama-3-8B’s 66.25 average but surpasses it on HellaSwag (41.75 vs. 39.21) and StoryCloze (80.42 vs. 79.49), using over 43× fewer reported tokens.The comparison is contextual rather than compute-matched because training recipes, tokenizers, and token-counting conventions differ.Source: Table 4 · p. 19, §4.5.2 · p. 19

5 Limitations & Future Prospects

The current system is an initial demonstration, not a finished endpoint: it remains far from convergence and is constrained by available resources. Future work must establish its performance ceiling, translate theoretical compression savings into practical speedups, and investigate decoder trade-offs and multimodal uses.

5 Limitations & Future Prospects

  • The model, especially the latent DiT, remains far from convergence, and available resources limit training scale.Larger models, longer training, and broader data and compute budgets are needed to establish its performance ceiling.Source: §5 · p. 20
  • Theoretical reductions in latent-prior FLOPs do not establish faster end-to-end generation, which also depends on decoder, CFG, feedback, and hardware factors.Infrastructure co-design, including specialized kernels, cache management, and parallel scheduling, is identified as necessary to translate theoretical savings into practical speedups.Source: §5 · p. 20
  • Co-training objectives, inference trade-offs, and the optimal division of labor between autoregressive and masked-diffusion decoders remain underexplored.Source: §5 · p. 20
  • The continuous latent space may provide a shared interface with modalities such as vision, but unified cross-modal representations and generation remain prospective.Source: §5 · p. 20

6 Conclusion

Cola2 combines dynamic semantic compression with balanced discrete reconstruction and generation through continuous–discrete co-adaptation. Evaluations support its design and scaling potential, including a resource-constrained result that beats Llama-3-8B on two tasks, while the latent DiT remains far from convergence.

6 Conclusion

  • A 60.10 eight-task average on 346.03B source tokens accompanies results above Llama-3-8B on HellaSwag and StoryCloze despite over 43× fewer reported pretraining tokens.The checkpoint is unconverged; the comparison contextualizes performance across different training recipes, tokenizers, and token-counting conventions.Source: §6 · p. 20
  • Cola2 unifies dynamic semantic compression, balanced discrete reconstruction and generation, and continuous–discrete co-adaptation.Token-prediction gradients update the encoder through the latents, while latent-prior fitting is decoupled from representation learning.Source: §6 · p. 20
  • Controlled evaluations on five downstream tasks scaled to about 210B training tokens support the design principles and show competitive, unsaturated scaling.Source: §6 · p. 20
  • The results indicate transfer of the Qwen-initialized causal decoder’s generation capability and compatibility with autoregressive and masked-diffusion decoders.Source: §6 · p. 20

A Detailed Derivations of the Probabilistic Model and Information Decomposition

The derivations relate the encoder, latent prior, and decoder to conditional likelihood and information decomposition. They also explain how compression constrains latent information and how encoder masking promotes predictive information about missing spans.

A.1 Conditional ELBO

  • The conditional ELBO separates the encoder’s information in the latent, the prior’s fit to the induced conditional latent distribution, and token modeling conditioned on the latent.Flow Matching is the solver used to learn the conditional prior, not part of the generative model’s definition.Source: §A.1 · p. 24

A.2 Information Identities and Semantic Decomposition

  • Conditional mutual information decomposes target information into information about high-level semantics and residual information given those semantics.Under compression, semantic information density distinguishes concentrating semantic information from merely shortening the latent sequence.Source: §A.2 · p. 24

A.3 KL, Compression Rate, and Latent Information Capacity

  • The VAE KL term upper-bounds latent information, while normalizing it per latent element makes regularization more comparable across latent lengths and widths.Compression rate controls the number of positions; dynamic segmentation determines where those positions are allocated.Source: §A.3 · p. 25

B Information-Theoretic Interpretation of Generation-Oriented Latent Supervision

  • The generation-oriented encoder-mask objective increases predictive information about missing spans rather than merely local reconstructability.Masked targets cannot be copied directly; the objective uses visible context, with fixed spans retained when encoder observations are replaced.Source: B Information-Theoretic Interpretation… · p. 25, §B.1 · pp. 25–26

B.1 Fixed-Span Encoder-Mask / BERT-Style Objective

  • Fixed-span encoder masking applies BERT-style generation pressure by predicting masked text from latents encoded without observing those targets.The implementation is not equivalent to arbitrary span-BERT losses used in earlier implementations.Source: §B.1 · pp. 25–26

C Latent Hacking and Gradient Decoupling in Compressed Representations

The analysis distinguishes useful decoder-to-encoder learning from prior-loss gradients that can conflict with representation objectives. It gives conditional risks of semantic loss and decoder sensitivity under compression, motivating detachment of the prior gradient without claiming that training must collapse.

C Latent Hacking and Gradient Decoupling in Compressed Representations

  • The analysis distinguishes retaining discrete-generator gradients, which serve token reconstruction and generation, from allowing latent-prior Flow Matching gradients to shape the encoder.It does not claim that compression necessarily shrinks the valid region or that training generally collapses.Source: C Latent Hacking and Gradient… · p. 26

C.1 Joint Objective and Gradient Conflict

  • The DiT prior loss can decrease while the representation objective locally increases because their gradients are not inherently aligned.Source: §C.1 · p. 26

C.2 The DiT-to-Encoder Gradient Is Not a Semantic Gradient

  • Flow Matching gradients depend on denoising residuals, the DiT Jacobian, conditioning, and noise schedule, not on a reconstruction likelihood or semantic objective.Thus the prior loss makes the latent easier for the prior solver to fit but does not guarantee preservation of decodable semantics.Source: §C.2 · pp. 26–27

C.3 Lowering the FM Loss Admits a ‘‘Simpler Latent’’ Shortcut

  • With a trainable encoder, Flow Matching loss can fall through changes that reduce conditional velocity variance, potentially flattening semantically meaningful modes.A lower Flow Matching loss alone therefore cannot rule out semantic collapse under unrestricted joint training.Source: §C.3 · p. 27

C.4 Information Constraints and Local Decoder Geometry under Compression

  • Under compression, increased predictive-information requirements may reduce information slack if latent capacity does not increase accordingly, but this depends on data, context, and distortion.Local decoder sensitivity claims also require specified conditions; the perturbation scale is not a guaranteed safe radius, and adding tokens does not ensure greater curvature.Source: §C.4 · pp. 27–28

C.5 Prior-Induced Local Drift and Gradient Decoupling

  • Prior-induced latent drift may increase token loss along decoder-sensitive directions, with the effect depending on drift, the encoder Jacobian, and decoder adaptation.Detachment blocks the direct prior-loss gradient path to the encoder, but does not guarantee stability of the full training process.Source: §C.5 · pp. 28–29

D Joint Information-Theoretic Interpretation of VAE and DiT Noise

VAE posterior noise and DiT corruption jointly affect how distinguishable latent semantics are, so their calibration depends on latent sharpness. The decoder’s clean, noisy, and dropped conditions distribute learning across reconstruction, robustness to perturbation, and weakly conditioned generation.

D Joint Information-Theoretic Interpretation of VAE and DiT Noise

  • The VAE posterior noise and DiT corruption jointly determine semantic-signal discriminability, so a fixed DiT noise schedule can correspond to different semantic-information regimes across latents.The two noise sources should be calibrated jointly; no fixed logSNR or timestep is claimed to be a model-independent optimum.Source: D Joint Information-Theoretic… · p. 29

E.1 Information Allocation under Different Latent Conditions

  • Mixing clean, noisy, and dropped latent conditions trains one decoder across strong conditioning, local perturbations, and weakly conditioned generation.Training only with clean conditions may create dependence on precise latents; training only with dropped conditions may cause the decoder to ignore latents.Source: §E.1 · p. 30
  • The best mixture of clean, noisy, and dropped conditions depends on model capacity, training budget, and inference-time latent errors, and should be determined empirically.The decoder’s cross-entropy includes remaining conditional uncertainty and finite-capacity modeling error.Source: §E.1 · p. 30

E.2 Local Perturbation Path of the Noise Condition

  • Small Gaussian perturbations train adaptation to imprecise latents, but do not guarantee uniform robustness to inference-time deviations.The perturbation distribution is not assumed to match DiT output errors, and Gaussian noise has no finite support.Source: §E.2 · p. 30

F Probabilistic Derivation of Classifier-Free Guidance

The appendix interprets classifier-free guidance (CFG) as a condition-dependent adjustment to model outputs, with distinct interpretations for continuous Flow Matching and the masked decoder. Its identities are qualified: pointwise or token-level relations do not guarantee a corresponding global distribution for complete generation.

F Probabilistic Derivation of Classifier-Free Guidance

  • CFG combines conditional and unconditional branch outputs using a guidance scale, but its per-time or token-level interpretations do not guarantee a power-mixture form for the complete generated distribution.For the latent generator the operator combines velocities; the MDM decoder combines logits along prefix and latent axes.Source: F Probabilistic Derivation of… · p. 30, §F.3 · p. 33

F.1 From Bayes’ Rule to the Guided Score

  • At a fixed noise time, γ = 0 recovers the base density, γ = 1 the conditional density, and values below or above 1 attenuate or amplify the conditional likelihood ratio.This describes the instantaneous score field, not the final sampling distribution.Source: §F.1 · p. 31

F.2 Equivalence of Flow Matching Velocity CFG and Score CFG

  • Under a shared Gaussian path and optimal branch regression fields, velocity CFG and score CFG correspond pointwise, but this does not ensure the guided flow’s endpoint follows the same power-reweighted data distribution.The per-time score identity alone does not guarantee path consistency.Source: §F.2 · p. 32

F.3 Prefix--Latent Two-Axis CFG for MDM

  • The MDM’s two-axis CFG combines prefix-visibility and latent-condition branches, yet changing demasking states means the token-level rule does not automatically define a global mixture over complete texts.The coefficients can be negative, making the combination affine in logits rather than a convex mixture of probabilities.Source: §F.3 · pp. 32–33

G Implementation Details of the Overall Workflow

This section details Cola2’s overall workflow, covering segmentation, decoder, training, inference, and the refinement and commitment of block conditions during decoding. It also defines the execution-graph conventions and notation for computational-cost accounting.

G.1 Dynamic Semantic Segmentation and Compression-Rate Control

  • Dynamic segmentation maps adjacent-token semantic discontinuities to boundary probabilities, while a non-backpropagated bias adjusts boundary density toward the target rate.A maximum span length is enforced through forced cuts during a greedy scan.Source: §G.1 · p. 33

G.1.1 Convergence of Compression-Rate Control

  • With a fixed encoder and attainable target, the compression-rate controller converges locally at a geometric rate under the stated monotonicity and step-size conditions.The appendix cautions that this is not a convergence guarantee for fully non-stationary joint training; it offers a tracking interpretation for slowly varying responses.Source: §G.1.1 · pp. 33–34

G.2 Decoder Duration and Token-Level Conditioning

  • The duration head predicts span lengths, expands each latent with within-span relative positions, and refines token-level conditions through block-causal conditioning layers.This makes token positions expanded from the same latent distinguishable.Source: §G.2 · p. 35

G.3 Block-Causal Masked-Diffusion Decoder

  • The block-causal masked-diffusion decoder denoises within each current token block without reading its clean target copy, and encoder-side masking trains latents to recover missing semantics from context.The encoder-mask branch fixes spans before masking and supervises only tokens masked on the encoder side, preventing boundary changes from making the task easier.Source: §G.3 · p. 35

G.4 Stage 1: Unified Compression VAE / Discrete Generator Decoder Objective

  • Stage 1 trains the block-causal decoder end to end with the encoder under clean, noisy, and dropped latent conditions; Stage 2 retains token-loss updates to the encoder while isolating the DiT loss with stop-gradient.This is a full-state handoff, not a frozen-VAE phase; token prediction gradients can still shape encoder latents through the condition path.Source: §G.4 · p. 36

G.5 Stage 2: Detached Block-Causal Flow Matching

  • Stage 2 trains a block-causal Flow Matching DiT over latent blocks while the online encoder and discrete decoder continue updating, but the DiT loss cannot update the encoder.Noisy queries read strictly earlier clean blocks and denoise bidirectionally within their current block; partial-clean-prefix augmentation handles prompts misaligned with block boundaries.Source: §G.5 · p. 37

G.6 Complete Inference Procedure and Latent Feedback

  • At inference, Cola2 re-encodes finalized tokens under the predicted fixed spans and commits corrected latents, keeping latent, condition, and token histories aligned.The proposal remains temporary until feedback correction; the current block is not regenerated after correction, so feedback affects subsequent blocks.Source: §G.6 · p. 38

G.6.1 Execution Details of MDM Prefix--Latent CFG; Appendix F.3. The refined condition is prepared only once per block:

  • Inference uses prefix–latent CFG for token decoding, with all branches sharing the refined condition and condition-off branches zeroing its injection after refinement.The condition is prepared once per block; zeroing it does not recompute duration or the condition stack.Source: §G.6.1 · p. 38, Appendix F.3 · p. 39

Appendix F.3. The refined condition is prepared only once per block:

  • Refinement branches read the same token state without updating the cache, and finalized tokens are committed with a clean, latent-free forward pass.The committed cache reflects only finalized tokens; branch hidden states and masked states are temporary. Different CFG weights can change subsequent tokens, unmasking order, and committed history, so full rollouts cannot share one offline logits trajectory.Source: Appendix F.3 · pp. 39–40

H Detailed Derivation of the Computational Cost

  • Stage 2 cost accounting includes continued training of the compression encoder and discrete generator decoder, plus a detached DiT; inference adds fixed-span feedback before corrected latent-history commit.This replaces the earlier “Stage 1 VAE + Stage 2 DiT/decoder” and one-way “DiT →decoder” accounting.Source: H Detailed Derivation of the… · p. 40

H.1 Notation

  • The FLOP estimates count dominant matrix multiplications across actual calls, with executed attention pairs determined by the implementation rather than mask-valid pairs alone.Training counts forward and backward passes but excludes recomputation, optimizer updates, communication, and element-wise operations; duration prediction does not guarantee the realized generation ratio because of boundary effects and padding.Source: §H.1 · pp. 40–41

H.2 Training: Persistent Representation/Decoder Path and an Additional Detached DiT

The appendix accounts for training and inference costs across the persistent representation/decoder path, the detached DiT, blockwise generation, latency, and KV caching. It emphasizes that compression affects only part of the system, so end-to-end comparisons require a consistent accounting scope.

H.2 Training: Persistent Representation/Decoder Path and an Additional Detached DiT

  • The persistent representation path remains part of Stage 2, while detached DiT and masked-input requests add computation beyond the main decoder pass.Clean, noisy, and dropped latent conditions share one main request rather than tripling the backbone pass; two decoder requests process about 4T positions, with a dominant linear term of 24N_DT.Source: §H.2 · p. 41
  • Compression mainly reduces the added prior’s cost; it does not make total training cost scale automatically with the compression rate.The reported scaling ratios describe linear length and attention pairs under comparable layouts, not a fixed speedup for the complete system.Source: §H.2 · p. 42

H.3 Inference: Proposal, Discrete Decoding, Feedback, and Commit

  • Each inference block separates proposal, discrete decoding, feedback, and commit; proposed latents enter the persistent cache only after feedback.The feedback encoder and corrected-condition commit are distinct from proposal preparation, and these components are not multiplied by four for CFG branches.Source: §H.3 · pp. 42–44

H.4 Serial Call Rounds, Latency, and Block Parallelism

  • Serial call-round counts are an idealized proxy, not a sufficient condition for end-to-end speedup, and actual latency depends on module, sequence, batching, and hardware.Sequentially invoked CFG branches must be counted as separate forward-pass rounds.Source: §H.4 · p. 45

H.5 KV Cache and Temporary States

  • Persistent KV accounting covers latent- and token-axis histories but excludes temporary states, so it is not total peak memory.End-to-end comparisons should separately account for processed positions, call rounds, latency, and KV capacity under a unified FLOP scope.Source: §H.5 · p. 45

I Final-Checkpoint Pretraining Recipe

  • The final-checkpoint pretraining recipe is separate from the matched component and controlled scaling setup.The section describes the recipe used for the final checkpoint in Table 4.Source: I Final-Checkpoint Pretraining Recipe · p. 46

I.1 Shared Training Configuration

The shared configuration uses retokenized OLMo 2 data, 4096-token sequences, and AdamW. Stage I trains a dynamic-compression VAE with a causal dual-stream decoder objective and clean, noisy, and dropped latent conditions.

I.1 Shared Training Configuration

  • The corpus is retokenized with the Qwen3-4B tokenizer and uses 4096-token sequences; all stages use AdamW with β = (0.9, 0.95) and weight decay 0.01.Stages II and III use EMA decay 0.9998, and stage boundaries restore model, optimizer, scheduler, and data-cursor state.Source: §I.1 · p. 46

I.2 Stage I: Dynamic-Compression VAE

  • Stage I processes approximately 52.43B source tokens and compresses about two tokens per latent on average, using maximum span length 8 and latent dimension 64.The decoder uses causal dual-stream training with clean, noisy, and dropped latent conditions sampled at 0.4/0.3/0.3; masked positions are [MASK] inputs while self-attention remains causal.Source: §I.2 · p. 46

I.3 Stage II: Joint VAE and Latent-DiT Training

Stage II jointly trains the VAE and latent DiT with detached VAE latents; Stage III freezes the EMA VAE encoder and continues DiT optimization. The accompanying budget table reports source-token positions processed in each stage.

I.3 Stage II: Joint VAE and Latent-DiT Training

  • Stage II processes approximately 125.83B source tokens while jointly training the VAE and an approximately 4B-parameter latent DiT.The DiT reads detached VAE latents, so its diffusion loss does not backpropagate into the VAE, which is updated only by compression and decoder objectives.Source: §I.3 · p. 46

I.4 Stage III: Frozen VAE and Continued DiT Training

  • Stage III freezes the Stage-II EMA VAE encoder and optimizes only the DiT, processing approximately 167.77B source tokens.The released model uses the EMA VAE decoder associated with the frozen encoder.Source: §I.4 · p. 46

I.5 Source-Token Accounting

  • Table 5 reports final-checkpoint pretraining budget as source-token positions processed in each stage.The supplied caption identifies the table’s scope but does not state stage-level token values.Source: Table 5 · p. 47

Related papers

Loading 2610.31305v1…