PaperlayerView exact revision on arXiv

Source-linked AI summary

Latent JEPA: Abstract Future Prediction for Latent Reasoning in Chemistry

Xinjian Zhao, Yaoyao Xu, Xuemin Chen, Xiaozhuang Song, Tianshu Yu

arXiv:2610.01947v1cs.LGcs.CL

TL;DR

Chemical reasoning models may benefit from anticipating plausible outcomes before a solution is fully articulated, complementing explicit multistep reasoning. LATENT-JEPA trains continuous latent thoughts to predict textual and molecular future views alongside autoregressive learning. On ChemCoTBench, it improves molecular optimization and several editing and reaction metrics, while its representations become more informative about molecular outcomes and more closely aligned with structure.

  • Problem

    The paper asks whether predicting future-solution representations can improve the information retained in continuous latent thoughts and the reasoning they support.

    Source: §1 · p. 2

  • Method

    LATENT-JEPA combines autoregressive learning with textual and molecular representation-prediction objectives for future reasoning and outcomes.

    Source: §1 · p. 2

  • Results

    LATENT-JEPA improves molecular optimization and several editing and reaction metrics, while making molecular outcomes more recoverable from latent thoughts and representations more aligned with molecular structure.

    Source: §1 · p. 2, §5.2 · p. 6, §6 · p. 10

  • Takeaways & Limitations

    The findings support abstract future prediction as a learning objective for connecting continuous latent reasoning with scientific outcomes.

    Source: §1 · p. 2, §6 · p. 10

Abstract

from arXiv · show

Large language models offer a promising foundation for chemical reasoning, bringing together chemical knowledge and multistep problem solving. Chemical intuition can provide an initial sense of plausible outcomes before the details of a solution are fully worked out. Inspired by how such expectations complement explicit analysis, we study how continuous latent thoughts can be trained to anticipate informative aspects of future solutions without verbalizing every intermediate step. We introduce Latent JEPA, a framework that combines autoregressive learning with joint-embedding prediction of one or more future views. For chemical reasoning, we develop textual and molecular prediction objectives that connect latent thoughts to both subsequent reasoning and molecular outcomes. Experiments on ChemCoTBench show gains in molecular optimization and on several editing and reaction metrics. Representation analyses show that future prediction makes latent thoughts more informative about molecular outcomes and strengthens their correspondence with chemical structure. These findings support abstract future prediction as a learning principle for connecting continuous latent reasoning with scientific outcomes.

1 INTRODUCTION

The introduction frames chemical intuition as an early sense of plausible outcomes that complements explicit analysis, then asks whether continuous latent reasoning can learn such anticipation. LATENT-JEPA trains latent thoughts to predict textual and molecular views of future solutions, with reported gains on chemical tasks and more structure-linked representations.

1 INTRODUCTION

  • Chemical intuition can suggest plausible transformations or structural effects before a solution is worked out, complementing explicit analysis and later testing.The introduction connects this motivation to accounts of tacit scientific knowledge and sensitivity to coherence before explicit identification.Source: §1 · p. 1
  • The paper asks whether predicting future-solution representations can improve the information retained in continuous latent thoughts and the reasoning they support.Continuous reasoning allows internal thought updates without verbalizing every step, unlike language-based chain-of-thought.Source: §1 · pp. 1–2
  • The approach draws on abstract prediction: latent thoughts anticipate meaningful aspects of subsequent reasoning and outcomes without reconstructing every observable detail.The motivation is to complement autoregressive generation, not replace the explicit realization of the solution.Source: §1 · p. 2
  • LATENT-JEPA combines autoregressive learning with prediction of future representations, training shared latent thoughts on textual reasoning and molecular outcome views.The two views capture complementary aspects: different explanations can describe the same molecule, while similar reasoning can yield different structures.Source: §1 · p. 2
  • Future prediction improves molecular optimization and several editing and reaction metrics, while making molecular outcomes more recoverable from latent thoughts and representations more aligned with molecular structure.Editing results are more mixed, while molecular optimization gains are most consistent across the six objectives.Source: §1 · p. 2

2 RELATED WORK

Related work situates LATENT-JEPA at the intersection of predictive representation learning, continuous latent reasoning, and chemical representations. The paper applies embedding prediction to future reasoning and molecular outcomes as complementary views.

Predictive representation learning

  • JEPA predicts embeddings of withheld views rather than reconstructing low-level details; related work extends predictive learning to future text segments and combines embedding prediction with language modeling.The cited examples include LLM-JEPA (Huang et al., 2026), Future Summary Prediction (Mahajan et al., 2026), and Semformer (Yin et al., 2024).Source: §2 · p. 2

Continuous latent reasoning

Chemical reasoning and representations

  • Chemical reasoning connects linguistic explanations with molecular structures and task outcomes, while learned molecular and language-model representations provide ways to describe both.The paper uses language and molecular representations as targets for future prediction.Source: §2 · p. 3

3 LATENT-JEPA

LATENT-JEPA combines autoregressive learning with prediction of encoded future views from continuous latent thoughts. Its framework allows different views, encoders, and predictors, while this work uses prediction to supervise latent thoughts during training.

3 LATENT-JEPA

  • LATENT-JEPA combines autoregressive learning with prediction of subsequent reasoning and outcome representations, aiming for useful expectations before a solution is fully expressed.A continuous latent thought is read from the reasoner’s current workspace.Source: §3 · p. 3
  • The framework predicts encoded representations of selected future views from a latent thought, using a view selector, target encoder, and predictor for each supervised step.A view may represent a reasoning step, answer, or domain-specific outcome; the discrepancy between predicted and target representations supplies the learning signal.Source: §3 · p. 4
  • Future views may use different encoders and representation spaces, and the framework leaves target design and component use after training open.In this work, predictive branches supervise latent thoughts during training without changing the reasoner’s generation procedure.Source: §3 · p. 4
  • The predictive loss is combined with the existing training objective, which may include architecture-specific auxiliary losses.The final latent thought is selected for supervision, allowing it to integrate preceding latent steps without access to the future continuation.Source: §3 · p. 4

3.1 TARGET ASYMMETRY AND GEOMETRY

  • Stop-gradient makes prediction asymmetric, while target centering and L2 normalization expose variation across examples when embeddings share a dominant direction.Pairwise similarity, effective rank, and shuffled-target prediction are used as diagnostics; the training objective remains non-contrastive.Source: §3.1 · p. 4

4 LATENT-JEPA FOR CHEMICAL REASONING

LATENT-JEPA adds textual and molecular future-prediction objectives to Coconut-Chem’s recurrent latent reasoning. The section describes how latent thoughts are trained to anticipate continuation and molecular outcomes.

4.1 RECURRENT CHEMICAL REASONING

  • Coconut-Chem progressively replaces an initial prefix of explicit reasoning with latent thoughts, then autoregressively generates the remaining reasoning and answer.LATENT-JEPA also supervises the final latent thought with available textual and molecular targets.Source: §4.1 · p. 5

4.2 DECODER-NATIVE CONTINUATION VIEW

  • The decoder-native textual view predicts representations of the remaining reasoning and answer from the final latent thought, using continuation-only targets and horizon-specific queries.Targets are detached, centered separately for reasoning and answers, and ℓ2-normalized; the cosine-distance objective weights reasoning and answer prediction equally when both are available.Source: §4.2 · p. 5

4.3 STRUCTURE-NATIVE OUTCOME VIEW

  • The structure-native view predicts a frozen encoder representation of the reference molecule or molecular state; text-only answers receive no molecular prediction objective.Molecular targets are mean-centered and ℓ2-normalized, and separate predictors map the same latent thought into textual and molecular target spaces.Source: §4.3 · p. 5

4.4 OBJECTIVE AND INFERENCE

  • Training retains Coconut-Chem’s autoregressive and counterfactual objectives, adds eligible predictive losses, and uses prediction heads only for supervision, not inference.Each predictive loss is averaged over eligible examples and uses unperturbed inputs.Source: §4.4 · p. 6

5 EXPERIMENTS

Experiments show that future prediction improves chemical-task performance and makes molecular outcomes more recoverable from latent thoughts, while target views and target construction affect the supervision. Analyses also find stronger correspondence between recurrent representations and molecular structure.

5.1 EXPERIMENTAL SETUP

  • The evaluation uses ChemCoTBench tasks spanning optimization, editing, molecular understanding, scaffold abstraction, and forward-reaction prediction, with three sampled responses averaged per problem.Source: §5.1 · p. 6

5.2 MAIN RESULTS

  • Molecular prediction is particularly effective for editing, while adding textual prediction improves most property objectives and reaction metrics relative to molecular prediction alone.The combined objective extends benefits across more task types without leading on every metric.Source: §5.2 · p. 6
  • Future prediction also improves most evaluated metrics in LatentChem, while Vanilla Coconut favors textual prediction and different views benefit different task subsets.Source: §5.2 · pp. 6–7
  • LATENT-JEPA improves most reported metrics, with gains on all six molecular-optimization objectives, broader improvements in structural analysis and reaction prediction, and more mixed editing results.Source: §5.2 · p. 6, Table 1 · p. 7, Table 2 · p. 7

5.3 LEARNING PREDICTIVE LATENT THOUGHTS

  • On held-out examples, final-update prediction cosine rises from 0.186 to 0.397, with higher within-task retrieval and greater similarity to corresponding than shuffled targets.Examples were excluded from both reasoner training and probe fitting.Source: Figure 2 · p. 8, §5.3 · p. 8
  • Molecular-only and dual prediction show stronger agreement with SMI-TED geometry than Coconut-Chem or text-only prediction; molecular-only also leads on Morgan fingerprints.The analysis compares pairwise distances without fitting a mapping.Source: Figure 3 · p. 8, §5.3 · p. 8
  • Probe cosine, Hit@5, and structural agreement generally rise with recurrent depth, though structural gains vary by task and explicit continuations do not generally shorten.Source: Figure 2 · p. 8, Figure 3 · p. 8, §5.3 · p. 9
  • Centering and normalizing molecular targets makes outcome differences more visible to cosine prediction, and retaining both centers yields the strongest average task-ablation rank.The online textual encoder with dropout has the strongest average rank among tested target-update variants, each of which improves a majority of reported metrics over Coconut-Chem.Source: §5.3 · p. 9, Table 3 · p. 9

6 CONCLUSION

LATENT-JEPA connects continuous reasoning to predictions of future representations, instantiated in chemistry through textual and molecular views. Experiments report chemical-task gains and stronger molecular information in latent thoughts.

6 CONCLUSION

  • LATENT-JEPA connects a continuous reasoner to one or more future views, allowing different encoders and predictors to define supervision.Source: §6 · pp. 9–10
  • In chemistry, the predicted views represent the remaining reasoning and the molecular outcome, and experiments show gains in task performance and molecular information in latent thoughts.Source: §6 · p. 10

A ADDITIONAL RELATED WORK

  • The paper frames target construction around what the encoder observes and represents, using textual continuation and the molecular object as future-solution views.Source: A ADDITIONAL RELATED WORK · p. 14

B IMPLEMENTATION DETAILS

The implementation trains Coconut-Chem and its predictive variant under a shared latent-thought curriculum, with textual and molecular targets supervising the final latent thought. A supervised-adaptation experiment also tests answer-representation prediction in LatentChem.

B IMPLEMENTATION DETAILS

  • Coconut-Chem and its LATENT-JEPA variant use the same curriculum schedule, with cthought = 2 continuous thoughts per replaced reasoning step and latent lengths Ks ∈ {0, 2, 4, 6}; each phase lasts three epochs.The implementation uses λreason = 0.025 and λchem = 0.01; the initial, fully textual phase has no latent thoughts and receives neither predictive loss.Source: B IMPLEMENTATION DETAILS · p. 14
  • The textual target is an independently encoded residual reasoning-and-answer sequence, and predictions from the final latent thought target each remaining reasoning step and the answer.The target pass uses current decoder parameters with gradients disabled; all predictions use the same final latent thought.Source: B IMPLEMENTATION DETAILS · p. 15
  • Molecular targets are fixed, centered, ℓ2-normalized SMI-TED-light encodings of reference SMILES, while separate running centers normalize reasoning and answer targets.The reasoning and answer centers use an EMA with β = 0.99; molecular targets and their training mean remain fixed.Source: B IMPLEMENTATION DETAILS · p. 15
  • Molecular counterfactual loss applies only when molecular features exist, while predictive losses use the unperturbed input.Examples without molecular inputs receive no molecular counterfactual loss.Source: B IMPLEMENTATION DETAILS · p. 16

C TRANSFER TO LATENTCHEM

  • The LatentChem transfer experiment tests answer-representation prediction during supervised adaptation, before reinforcement learning, with both arms receiving the same training budget and trainable modules.The molecular updater is frozen; the table marks improvements over the adapted LatentChem baseline in dark red.Source: C TRANSFER TO LATENTCHEM · p. 16, Table 4 · p. 16

D TARGET ENCODER ABLATION

The text-target encoder ablation varies target dropout and parameter updates while retaining the molecular branch and the other training settings. The variants trade task-specific gains rather than showing a uniform advantage for momentum updates.

D TARGET ENCODER ABLATION

  • Each text-target encoder variant improves a majority of reported metrics over Coconut-Chem, and the online encoder with dropout has the strongest average rank.The comparison retains both predictive views, predictors, loss coefficients, and the training schedule.Source: D TARGET ENCODER ABLATION · pp. 16–17, Table 5 · p. 17
  • Disabling dropout improves LogP but reduces deletion and substitution accuracy, while EMA improves substitution, QED, and DRD2 but reduces deletion accuracy and major-product prediction.The EMA alternative uses LoRA parameters with momentum 0.995 and disables target dropout.Source: D TARGET ENCODER ABLATION · pp. 16–17

E TRAINING AND TARGET DIAGNOSTICS

The diagnostics examine how target centering affects molecular target geometry, task performance, and source-specific prediction. They report that calibration separates molecular targets and that retaining both centers gives the strongest average rank.

E TRAINING AND TARGET DIAGNOSTICS

  • Retaining both centers gives the strongest average rank and outperforms either center removal on most task metrics; removing text centering lowers all six property-improvement scores.The centering comparison evaluates both target geometry and the numerical scale of the cosine loss.Source: E TRAINING AND TARGET DIAGNOSTICS · p. 17, Table 6 · p. 18
  • 0.989 to 0.010: centering and ℓ2 normalization reduce mean random-pair cosine among molecular targets, whose raw representations share a dominant direction.Figure 4 also compares spectral energy before rowwise normalization.Source: E TRAINING AND TARGET DIAGNOSTICS · p. 17, Figure 4 · p. 18
  • At the final curriculum phase, the dual model’s autoregressive, weighted textual, and weighted molecular losses are approximately 0.0665, 0.0144, and 0.0021, respectively.Predictive losses are lower than their shuffled-source counterparts.Source: E TRAINING AND TARGET DIAGNOSTICS · p. 17
  • Without molecular centering, original- and shuffled-source prediction distances are both small (0.003 and 0.009), consistent with relying less on the corresponding input.Removing text centering also decreases measured source separation, while the dual model benefits on more metrics from retaining both centers.Source: E TRAINING AND TARGET DIAGNOSTICS · p. 18, Table 7 · p. 19

F PREDICTION VIEWS AND CHEMICAL TASK PERFORMANCE

Prediction views contribute differently across chemical tasks: dual prediction has the strongest average rank in Coconut-Chem, while textual prediction leads the Vanilla Coconut comparison. Molecular target relationships also cross task-family boundaries.

F PREDICTION VIEWS AND CHEMICAL TASK PERFORMANCE

  • Dual prediction has the strongest average rank over shared metrics, while molecular prediction is particularly effective for editing and adding the textual view improves most property objectives and all reported reaction metrics over molecular-only prediction.The comparison varies textual supervision strength as well as the set of prediction views.Source: F PREDICTION VIEWS AND CHEMICAL TASK… · p. 19
  • Textual prediction improves 10 of 17 metrics over Vanilla Coconut, with four regressions and three ties, and has the lowest average rank in that comparison.Molecular and dual prediction improve seven and eight metrics, respectively; the source interfaces, initialization, and training histories are not otherwise matched.Source: F PREDICTION VIEWS AND CHEMICAL TASK… · pp. 19–20, Table 9 · p. 20

G CHEMICAL RELATIONSHIPS IN THE PREDICTION TARGETS

  • Molecular target relationships do not simply reproduce task taxonomy: major-product and byproduct targets have lower affinity than byproduct and ring-system targets.Murcko and ring-system targets form a close pair; the figures show calibrated target representations and cross-task cosine relationships.Source: G CHEMICAL RELATIONSHIPS IN THE… · p. 21, Figure 6 · p. 21, Figure 5 · p. 21
  • The reported affinities describe this encoder and target cohort; identifying their specific chemical drivers would require additional molecular annotations.The molecular target view supplies similarities among solutions beyond task membership.Source: G CHEMICAL RELATIONSHIPS IN THE… · pp. 21–22

H ORGANIZATION OF LATENT THOUGHTS

The choice of predictive view changes which information is most accessible in latent thoughts. Retrieval, molecular task-geometry agreement, and local organization favor different views or reveal distinct patterns.

H ORGANIZATION OF LATENT THOUGHTS

  • The comparison analyzes final-layer hidden states at the sixth latent thought across 16,943 training examples and 13 tasks, using a separate probe split from Figure 2.The PCA basis is shared within the four-variant comparison; dispersion and task-centroid agreement use full-dimensional calibrated states.Source: H ORGANIZATION OF LATENT THOUGHTS · p. 22
  • Dual prediction raises outcome retrieval from 30.8% to 75.0% over Coconut-Chem, while molecular prediction gives the highest retrieval and dual prediction the strongest molecular task-centroid agreement.Dispersion does not follow retrieval: textual prediction has the lowest within-family dispersion, while dual prediction is close to Coconut-Chem (0.695 versus 0.694).Source: H ORGANIZATION OF LATENT THOUGHTS · p. 22
  • In Vanilla Coconut, dual prediction raises retrieval from 5.8% to 78.5%, while molecular prediction gives the strongest agreement with molecular task relations.The models differ in initialization and training history, so their absolute scores do not isolate the role of ChemTokens.Source: H ORGANIZATION OF LATENT THOUGHTS · pp. 22–23
  • Figure 9 qualitatively shows predictive supervision changing local neighborhood organization: textual prediction produces clearer task-family separation, while molecular and dual supervision yield different local arrangements accompanying stronger molecular-outcome retrieval and geometry agreement.Source: H ORGANIZATION OF LATENT THOUGHTS · p. 23

I PREDICTING MOLECULAR OUTCOMES FROM LATENT THOUGHTS

The analyses test whether molecular outcomes can be recovered from latent thoughts and whether their structure tracks molecular relationships across tasks and recurrent depth. They also examine explicit reasoning length and specify the representations, probes, geometry measures, and uncertainty intervals used.

I PREDICTING MOLECULAR OUTCOMES FROM LATENT THOUGHTS

  • LATENT-JEPA raises mean final-update cosine from 0.186 to 0.397, with higher cosine on 277 of 321 examples, indicating gains across most of the evaluation cohort.Corresponding targets remain more predictable than targets reassigned within the same task.Source: I PREDICTING MOLECULAR OUTCOMES FROM… · p. 24
  • Dual-model SMI-TED RSA rises from 0.190 at the first update to 0.400 at the final update; final-update scores are 0.110 for Coconut-Chem and 0.392 for molecular-only prediction.With Morgan fingerprints, final-update RSA is 0.029, 0.118, and 0.329 for Coconut-Chem, text-only, and dual prediction, respectively.Source: I PREDICTING MOLECULAR OUTCOMES FROM… · p. 24
  • Retrieval gains span tasks and recurrent depth, but their depth profiles differ: substitution gains grow later, while major-product and Murcko-scaffold gains appear at the first step.Task-level geometry gains are positive at the final update in every analyzed task, with larger later gains for Murcko and major-product prediction.Source: I PREDICTING MOLECULAR OUTCOMES FROM… · pp. 24–25, Figure 12 · p. 25, Figure 11 · p. 25

J EXPLICIT REASONING AFTER LATENT COMPUTATION

  • Explicit reasoning changes are not uniform across tasks: reductions are clearest for substitution and major-product prediction, while optimization stays near baseline and byproduct intervals are wider.Pre-answer token counts exclude final SMILES and answer-declaration lines.Source: J EXPLICIT REASONING AFTER LATENT… · p. 26

K METHODS FOR REPRESENTATION ANALYSIS

  • The appendix defines the representation cohorts and measurements for outcome recoverability, molecular relationships, and task organization, alongside the matched-view configurations.Optimization and curriculum settings are specified separately in Appendix B.Source: K METHODS FOR REPRESENTATION ANALYSIS · p. 27

K.1 REPRESENTATIONS AND EVALUATION COHORTS

  • Linear probes use final-layer hidden states, whereas RSA uses feedback embeddings; neither evaluation representation receives the gold continuation during extraction.On the fixed-length route, the hidden state at step t supplies the feedback embedding at step t + 1.Source: §K.1 · p. 27

K.2 LINEAR PROBING AND MOLECULAR RETRIEVAL

  • Separate ridge probes are fitted at each model and reasoning step, then cosine-scored against within-task molecular candidates to measure rank and Hit@k.Probe fitting excludes test examples and is performed independently of the JEPA predictors.Source: §K.2 · p. 27

K.3 REPRESENTATIONAL SIMILARITY ANALYSIS

  • For the SMI-TED comparison, cosine-distance RSA preserves model ordering: LATENT-JEPA rises from 0.176 at first update to 0.410 at final update, while Coconut-Chem falls from 0.138 to 0.118.The SMI-TED analysis uses Euclidean distances between feedback embeddings and training-mean-centered, L2-normalized molecular targets. Morgan analysis instead uses latent cosine distance and fingerprint Tanimoto distance.Source: §K.3 · p. 28

L.1 CHEMICAL TASKS

The evaluation covers molecular optimization, editing, molecular understanding and scaffold tasks, and forward-reaction prediction, using task-specific outcome measures. Results average multiple sampled responses rather than selecting the best sample, and optimization success rates are reported separately.

L.1 CHEMICAL TASKS

  • The benchmark includes property optimization, addition/deletion/substitution editing, structural counting and scaffold tasks, and major-product or byproduct reaction prediction, each scored with task-specific metrics.Optimization scores are computational property estimates; improvement is measured against the source molecule, with invalid or unscorable outputs contributing zero improvement.Source: §L.1 · pp. 29–30

L.2 EVALUATION SET AND SCORING

  • Evaluation covers 1,020 optimization, understanding, scaffold, and reaction examples plus 100 editing examples; scores average three response streams rather than selecting the best response.Editing accuracy counts correct responses across all three samples, not problems solved by at least one sample.Source: §L.2 · p. 30

M OPTIMIZATION SUCCESS RATES

  • Optimization success rate is the fraction of scored generations with positive property improvement, averaged across the three completion streams.Table 11 reports molecular optimization success rates in percent.Source: M OPTIMIZATION SUCCESS RATES · p. 30, Table 11 · p. 30
Loading 2610.01947v1…