PR-JEPA: Decoupling Latent Planning from Few-Step Discrete Realization for Non-Autoregressive Generation of Code, Legal Documents, and Control
- 1SLING AI Inc.
Abstract
Autoregressive (AR) decoders perform one forward pass per token, so sequential depth grows with output length and the key-value (KV) cache with context. Both costs limit closed-loop control and long structured outputs in which late content constrains early content: the imports and signatures of a program, the pragmas and array declarations of High-Level Synthesis (HLS) C code, the defined terms and cross-references of a contract.
We propose PR-JEPA (Plan-Realize JEPA), which separates deciding what to generate from emitting it. A Joint-Embedding Predictive Architecture (JEPA) planner, trained with an exponential-moving-average target encoder and VICReg, samples a discrete mode and predicts a continuous plan over the output's structure, such as an abstract syntax tree or a clause outline. A cross-attention bridge conditions a few-step discrete flow-matching realizer that refines all positions in parallel in up to 8 steps, using a closed-form jump kernel tilted by plan-derived constraint energies and clipped to remain valid; a Logit-KL variant integrates a straight path in logit space instead.
We prove that small VICReg penalties exclude rank-one collapse, that the K-step realizer amplifies plan errors by at most a bounded Lipschitz factor, that plan-mediated hard constraints survive fully parallel updates, that re-anchoring bounds latent error over any horizon, and how to calibrate an uncertainty gate for a target risk level. A roofline model predicts 9 to 21 times lower single-request latency than AR decoding at 8 steps, at about one-eighth of AR's saturated throughput. We pre-register evaluation on general-purpose code, HLS-C, contract drafting, and robot manipulation; end-to-end measurements are pending.
Related
- Structure-First Text Generation: A Unified Framework for Hierarchical Planning, Content Realization, Ordering, and Composition in Autoregressive Language Models
- Design and realization of robust disturbance rejection decoupling control strategy for time delayed multi-variable process
- LinkerGen: Guided Autoregressive Generation of Organic Linkers for Reticular Chemistry
- Limitations on Autoregressive Token Generation via Sequential Sampling
- Meta Non-Autoregressive Models: A Framework for Generating, Selecting, and Adapting Non-Autoregressive Generation Mechanisms
BibTeX
@misc{lee2026prjepa,
title = {{PR-JEPA: Decoupling Latent Planning from Few-Step Discrete Realization for Non-Autoregressive Generation of Code, Legal Documents, and Control}},
author = {Yeonseok Lee},
howpublished = {Zenodo},
year = {2026},
month = oct,
doi = {10.5281/zenodo.23498811},
url = {https://zenodo.org/records/23498811},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - PR-JEPA: Decoupling Latent Planning from Few-Step Discrete Realization for Non-Autoregressive Generation of Code, Legal Documents, and Control
AU - Yeonseok Lee
PY - 2026
DA - 2026/10/20
PB - Zenodo
DO - 10.5281/zenodo.23498811
UR - https://zenodo.org/records/23498811
AB - Autoregressive (AR) decoders perform one forward pass per token, so sequential depth grows with output length and the key-value (KV) cache with context. Both costs limit closed-loop control and long structured outputs in which late content constrains early content: the imports and signatures of a program, the pragmas and array declarations of High-Level Synthesis (HLS) C code, the defined terms and cross-references of a contract. We propose PR-JEPA (Plan-Realize JEPA), which separates deciding what to generate from emitting it. A Joint-Embedding Predictive Architecture (JEPA) planner, trained with an exponential-moving-average target encoder and VICReg, samples a discrete mode and predicts a continuous plan over the output's structure, such as an abstract syntax tree or a clause outline. A cross-attention bridge conditions a few-step discrete flow-matching realizer that refines all positions in parallel in up to 8 steps, using a closed-form jump kernel tilted by plan-derived constraint energies and clipped to remain valid; a Logit-KL variant integrates a straight path in logit space instead. We prove that small VICReg penalties exclude rank-one collapse, that the K-step realizer amplifies plan errors by at most a bounded Lipschitz factor, that plan-mediated hard constraints survive fully parallel updates, that re-anchoring bounds latent error over any horizon, and how to calibrate an uncertainty gate for a target risk level. A roofline model predicts 9 to 21 times lower single-request latency than AR decoding at 8 steps, at about one-eighth of AR's saturated throughput. We pre-register evaluation on general-purpose code, HLS-C, contract drafting, and robot manipulation; end-to-end measurements are pending.
ER -