TL;DR: I built a hybrid where each "Chain-of-Thought" step is a JEPA prediction unit (dense MLP blocks), chained so the tail output of block i feeds the head input of block i+1. On grid-world world-modeling it beats naive Dense/CoT baselines by a mile. But the interesting part is the ablation: removing the chain (independent parallel units) wins in a simple environment, yet overfits and loses in a harder one — the chain acts as a regularizer, not just an error-accumulator. All experiments are small (CPU-trainable), fully reproducible, with param-matched baselines and collapse diagnostics.

The architecture  
Every CoT step is a JEPA unit: a Context Encoder (dense MLP) + a Predictor (dense MLP). The chain passes the predicted latent forward: z\_i = Predictor\_i(ContextEncoder\_i(z\_{i-1}), a\_{t+i-1}). Targets come from an EMA copy of the encoder (tau=0.99, stop-grad), and the loss is per-step 1 - cosine(z\_i, EMA(obs\_{t+i})), summed over K=4 steps. The final latent z\_K is the "world model embedding" that a policy head can consume.  
(see attached diagram: JEPA-CoT World Model Architecture (Chained Hybrid))  
Setup  
Grid-world navigation world model. Obs = obstacle map + agent one-hot (2 channels), 4 actions, random-policy data. 200 train maps / 50 held-out test maps (zero overlap, hash-verified). 3 seeds, 40 epochs, Adam 1e-3, batch 256. All models param-matched (<5% spread, asserted in code). Metrics on unseen maps: 1-step latent cosine, open-loop multi-step rollout MSE, frozen-latent linear probes (coordinate regression R\^2, per-action "will I get blocked" accuracy).  
Baselines: naive Dense (joint encoder, no stop-grad), naive CoT (same chain skeleton, per-step BCE reconstruction, no EMA/similarity loss), plus two "fair" strengthened baselines I added after suspecting the naive ones were strawmen: Dense-SG (stop-grad target = frozen random encoder) and CoT-LN (chain with LayerNorm).  
Results — Setting A: 8x8 grid, K=4 (mean±std, 3 seeds, unseen maps)  
model                params   1-step cos up  rolloutMSE@4 dn  coord R2 up   blocked acc up  
\-------------------  -------  -------------  ---------------  ------------  --------------  
JEPA-CoT (chained)   192,960  0.951±0.012    0.1328           0.955±0.012   0.883±0.005  
JEPA-CoT (parallel)  192,512  0.985±0.001    0.0706           0.988±0.000   0.884±0.002  
Pure Dense (naive)   199,864  1.000±0.000    0.0000           -1.945±2.379  0.772±0.000  
Pure CoT (naive)     201,192  -0.018±0.064   69.7069          0.293±0.036   0.800±0.012  
Dense-SG (fair)      193,411  0.928±0.006    0.0123           0.655±0.022   0.888±0.002  
CoT-LN (fair)        193,320  -0.013±0.006   7.5254           0.003±0.010   0.771±0.001

Results — Setting B: 12x12 grid, 18% obstacles, K=6 (harder)  
model                params   1-step cos up  rolloutMSE@6 dn  coord R2 up     blocked acc up  
\-------------------  -------  -------------  ---------------  --------------  --------------  
JEPA-CoT (chained)   297,536  0.993±0.000    0.1614           0.791±0.018     0.780±0.006  
JEPA-CoT (parallel)  297,344  0.999±0.000    0.1576           0.673±0.073     0.765±0.002  
Pure Dense (naive)   297,575  0.998±0.000    0.0000           -12.710±17.782  0.764±0.000  
Pure CoT (naive)     304,724  0.033±0.017    6.0415           0.038±0.015     0.764±0.000  
Dense-SG (fair)      296,937  0.836±0.007    0.0127           0.247±0.002     0.785±0.000  
CoT-LN (fair)        305,492  -0.018±0.050   2.8617           -0.015±0.010    0.764±0.000

The finding that surprised me  
I expected the parallel (unchained) variant to win everywhere — no error accumulation, and indeed in Setting A its rollout error decreases with horizon (0.115 -> 0.071) while the chained version's grows (0.069 -> 0.133). But in the harder Setting B the ranking flips on representation quality:  
metric (Setting B)         chained        parallel       reading  
\-------------------------  -------------  -------------  ---------------------------  
coord probe R2 (unseen)    0.791±0.018    0.673±0.073    chained generalizes better  
final train loss           0.049          0.005          parallel fits 9x better...  
1-step NMSE (unseen)       0.05-0.08      1.09-1.45      ...but is worse than mean-prediction

Parallel fits the training targets \~9x better yet generalizes worse — a textbook overfitting signature. The chain forces information through a sequential bottleneck: in a simple task that bottleneck is pure cost (error accumulation), in a complex task it's a useful inductive bias. Chain vs. parallel looks like a bias-variance knob, not a free lunch.  
Baseline sanity checks (the "these scores can't be right" section)  
The naive baselines looked suspiciously extreme, so I built diagnostics:  
model (Setting A)   latent std   neg-pair cos   1-step NMSE   pred/tgt norm  
\------------------  -----------  -------------  ------------  -------------  
Pure Dense (naive)  0.0005       0.9999          1.06          1.0  
Pure CoT (naive)    0.32         0.01            \~364          \~12x  
JEPA-CoT (chained)  0.37         0.35            0.47          0.47  
JEPA-CoT (parallel) 0.37         0.50            0.75          0.37

 • Naive Dense collapses to a near-constant latent (std 0.0005; any two latents have cosine 0.9999). Its "perfect" MSE is a trivial-solution artifact — a nice live demo of why JEPA needs the EMA target.  
 • Naive CoT's latent space is healthy but its predictions are \~300x worse than predicting the mean, with 12x norm explosion down the chain.  
 • The fair versions confirm the story: Dense-SG predicts fine (NMSE 0.25) but its representation is far weaker (coord R\^2 0.66/0.25 vs 0.96-0.99/0.67-0.79 for the JEPA variants — random features don't scale with obs dimension). CoT-LN stops exploding (NMSE \~2) but still learns nothing predictive (cosine \~0) — the CoT failure is mechanistic, not numerical.  
Limitations (being honest)  
 • Toy scale: 8x8/12x12 grids, MLPs, random-policy data. The chain-vs-parallel law needs replication at vision scale.  
 • The chained units use independent weights; I haven't yet ablated vs. weight-shared (autoregressive-style) chains — that's the strongest neighboring design and the next experiment.  
 • All metrics are open-loop + probes; no closed-loop control yet.  
 • dense probe R\^2 varies wildly between eval runs (-43 vs -1.9) because ridge probes on a collapsed latent are ill-conditioned; the conclusion (collapse) is unaffected.  
Reproducibility  
Everything ran on 2 CPU cores (\~19 min for the main sweep, \~50 min for the 12x12 sweep). Code (PyTorch), all 27 checkpoints, metrics, and plotting scripts are packaged; happy to share a link if there's interest. One-command repro: pip install -r requirements.txt && bash run\_all.sh.  
Curious whether the chain-as-regularizer effect holds in your world-model setups — and whether anyone has seen the same "parallel overfits, chain generalizes" flip at larger scale.l

JEPA-CoT: chaining JEPA units tail-to-head for latent world models — plus an ablation that surprised me (chain vs. parallel flips with task complexity)