Source-linked AI summary
A JEPA Recipe for Tabular Foundation Models
Mingyu Jeon, Suwan Cho, Jae Young Suh
TL;DR
Tabular foundation models predict hidden cell values, while world-model self-supervision predicts representations; earlier JEPA runs collapsed on this prior. The paper proposes a recipe combining latent and value objectives and compares it with the value-only arm at convergence. Across the real-data suite, the JEPA arm trails the value-only arm and takes longer to reach its plateau.
Problem
Tabular foundation models learn from value prediction, while world-model self-supervision calls for predicting representations; whether a latent objective can train on this prior and what it costs remains open, especially after earlier JEPA collapse.
Source: §1 · p. 1Method
The recipe trains a JEPA latent term beside the value objective, with the value head reading the encoder field and an EMA-difference target, and compares arms to a plateau.
Source: §1 · pp. 1–2, §3 · p. 4Results
At convergence, jepa trails ds across the real-data suite; classification gives 32:70 wins to losses, and the one-entry-per-dataset-name count is 29:63.
Source: §4.3 · p. 6, §5 · p. 7Takeaways & Limitations
The latent term can train on this prior without collapse, but its case must be made on something beyond in-context accuracy.
Source: §5 · p. 7, §6 · p. 9One run per arm limits the claims, and the measured deficit cannot be attributed to the latent term alone because no matched zero-weight control was trained.
Source: §5.1 · p. 8, §5 · p. 8
Abstract
from arXiv · showhide
Tabular foundation models learn to predict cell values in context, whereas world-model self-supervision asks for prediction in representation space (LeCun, 2022; Assran et al., 2023). On a tabular foundation-model prior, the latent term of a joint-embedding predictive architecture (JEPA) collapsed in our earlier runs and took the encoder with it to a constant map. We report a recipe under which the latent term survives to convergence beside the value objective: the value head reads the encoder field rather than the predictor, and the target is an exponential moving average (EMA) difference. To bound its cost against the value-only arm, both arms train until a plateau rule stops them, with no fixed step budget. A fixed horizon had confounded a slowdown with a ceiling, since the value-only arm was still improving well past the usual budget. At convergence, in one run per arm, the JEPA arm trails the value-only arm across 147 real datasets, 32:70 wins to losses on classification (29:63 with one entry per dataset name) and 8:24 on regression, the margin small on classification and wider on regression, and the count leans the same way in each stratum and each benchmark. The JEPA arm (jepa) needs 1.42 times as many steps as the value-only arm (ds), and 1.66 times its wall-clock, to reach its plateau.
1 Introduction
The paper asks whether latent prediction can train on a tabular foundation-model prior, where earlier JEPA runs collapsed. It proposes a recipe, an open-horizon comparison, and a converged evaluation against the value-only arm.
1 Introduction
- The question is whether a latent objective can train on a tabular foundation-model prior and what it costs, given that earlier JEPA runs collapsed the encoder to a constant map.The collapse also flattened the value head when it read predictor output.Source: §1 · p. 1
- The recipe trains a JEPA term beside the value objective by having the value head read the encoder field and using an EMA-difference latent target.The predictor receives mask tokens at hidden cells; the runs isolate the head-read change while holding the other two choices fixed.Source: §1 · pp. 1–2
- The open-horizon protocol stops each arm at a plateau, addressing evidence that the earlier fixed budget stopped the value-only arm short.The comparison evaluates convergence against the value-only arm on 147 real datasets, stratified by number of classes.Source: §1 · p. 2
- The paper does not claim the recipe beats the value-only arm; its narrower claim is that the latent term trains on this prior and has a bounded convergence cost.That cost is bounded from above by one run per arm, conditional on no other difference helping the arm with the latent term.Source: §1 · p. 2
2 Related work
The related work places the study among prior-fitted tabular models and self-supervised representation learning. It distinguishes training a latent term on the prior itself from earlier tabular self-supervision settings.
2 Related work
- Prior-fitted networks cast tabular prediction as amortised Bayesian inference over a synthetic prior; TabPFN and TabICL scaled this approach to thousands of context rows.Follow-up work attributes much of the family’s behaviour to the prior rather than the architecture, and later work explored real tables, retrieval, self-supervision, and curated synthetic-prior mixtures.Source: §2 · p. 2
- JEPA predicts representations of masked content from visible context and avoids collapse with an EMA target encoder or an explicit regulariser.This paper reports collapse on a tabular foundation-model prior and a wiring change that removes it, rather than a new regulariser.Source: §2 · p. 2
- Earlier tabular self-supervision includes row corruption and reconstruction, single-dataset JEPA variants, generated few-shot tasks, and cross-table pre-training.Unlike those lines, this work trains the latent term on the prior itself with 128 to 1,024 context rows and evaluates on the family’s real-data suite.Source: §2 · p. 2
3 Method
The method compares a value-only arm with a JEPA arm using a shared cell-level transformer and value head. The JEPA arm adds latent prediction, while both arms train on the same prior stream and stop under a plateau rule.
3 Method
- Both arms share a six-layer, eight-head, width-256 cell-level transformer and an MLP value head over 32 bins for hidden-cell values.The value objective is cross-entropy against each hidden cell’s binned value.Source: §3 · p. 2
- The JEPA arm’s value head reads the encoder field, while its predictor reads that field with hidden cells replaced by mask tokens.Its latent target is the detached, layer-normalised difference between EMA target-encoder embeddings of the full and masked tables.Source: §3 · p. 3, Figure 1 · p. 3
- The JEPA objective is trained in one stage beside the value loss, with one weight each; the target encoder follows the online encoder with decay .996.Hidden cells follow a mixed masking policy after a 4,000-step warm-up on target-column masking.Source: §3 · p. 4
- Both arms train on the same TabICL graph-SCM prior stream, with tables of 128 to 1,024 rows and up to 100 features under a 65,536-cell budget.Batches mix classification and regression tasks, and 30% of tables carry missing feature cells.Source: §3 · p. 4
- Each arm stops after its best validation MSE fails to improve by more than .002 for 80 consecutive validations, subject to a 1,000,000-step safety cap.Validation occurs every 250 steps; checkpoints are scored on the real-data suite throughout the open horizon.Source: §3 · p. 4
4 Experiments
The experiments compare a JEPA arm with a value-only arm at value-error plateaus and on real-data benchmarks. The JEPA arm trails at convergence and reaches its plateau later, while both arms end between linear models and gradient-boosted trees.
4.1 Setup
- The 147-dataset suite contains 115 classification and 32 regression datasets; counts compare paired dataset scores, drop exact ties, and use a same-seed ds rerun as a noise floor.Scores use accuracy or R2, with up to 1,024 context rows and 64 features; the sign test weighs each dataset equally, and benchmark breakdown tests are uncorrected.Source: §4.1 · pp. 4–5
4.2 Convergence under the plateau rule
- Both arms were stopped using squared error of the value head on hidden cells of held-out prior tables, not real-data scores; jepa’s latent loss was recorded but never used to stop runs.Jepa’s latent loss changed by +.0094 over the final 20,000-step patience window.Source: §4.2 · p. 5
- jepa stopped at 253,750 steps and 41.3 hours with validation MSE .4360, versus ds at 178,500 steps and 24.8 hours with .4179.That is 1.42 times the steps and 1.66 times the wall-clock; the latent term survived to the plateau, and the earlier fixed budget had stopped ds before its plateau.Source: §4.2 · p. 6
4.3 Converged comparison
- 32:70 classification wins to losses and 8:24 regression wins to losses favor jepa over ds in the losing direction at convergence; the classification margin is small, while regression gaps are wider.With one entry per dataset name, the counts are 29:63 for classification and 8:22 for regression; both arms exceed linear models and trail trees on every benchmark block.Source: §4.3 · pp. 6–7
4.4 Suite score along training
- Neither arm beats HistGB, while both beat linear models; jepa crosses the linear-model mean later, at step 140,000 on classification and 60,000 on regression.The baseline models use each dataset’s full training half, unlike the arms’ maximum 1,024 context rows, so this is not a like-for-like bound on what the objective could buy.Source: §4.4 · p. 7
5 Discussion
The latent term trains without collapse, but jepa does not improve real-data accuracy over ds at convergence. The measured deficit is bounded by a comparison with several differences between arms, while seed count and evaluation scope limit the conclusions.
5 Discussion
- At convergence, jepa trails ds on every significant stratum and benchmark, with the margin small on binary tasks and larger on many-class and regression tasks.The remaining accuracy deficit is small compared with the slowdown seen at a fixed budget.Source: §5 · p. 7
5.1 Threats to validity
- The measured deficit bounds the latent term’s cost from above, because masking policy, context split, and warm-up also differ between arms.The study does not test whether those differences help or hurt jepa; the shared encoder field is shaped by both losses.Source: §5 · pp. 7–8
- One seed per arm means the sign test compares datasets, not seeds; the same-seed noise floor does not measure seed-to-seed variation.Both arms used the same prior stream and suite, so the comparison is between two training runs, not yet two recipes.Source: §5.1 · p. 8
5.2 Future work
- A zero-latent-weight matched control would turn the upper bound into an estimate; robustness at short contexts or under noise remains an open test.The proposed two-stage form could also be tested by probing the frozen jepa field, and chained latent-space fills remain unevaluated.Source: §5.2 · pp. 8–9
6 Conclusion
The recipe lets latent prediction train beside value prediction on a tabular foundation-model prior without collapse, but it does not improve real-data accuracy at convergence. The conclusion is to compare arms at their plateaus and next use a matched control to estimate the latent term’s cost.
6 Conclusion
- At convergence, jepa trails ds by a small margin on binary tasks and by more on many-class and regression tasks.The conclusion describes the comparison as one run per arm across 147 real datasets.Source: §6 · p. 9
- Moving the value head to the encoder field lets the latent objective train beside the value loss and reach a plateau without collapse.The target and mask tokens are held fixed in the reported recipe.Source: §6 · p. 9
- Stopping each arm at its plateau avoids scoring a slowdown as a ceiling; the next step is a matched control to estimate the bound.A fixed budget had stopped the value-only arm before its plateau.Source: §6 · p. 9
A.1 Configurations
The two arms share a plateau rule but differ in several configuration and validation details, so the comparison is not limited to the latent-term configuration.
A.1 Configurations
- The same plateau-rule function stops both trainers under the same criterion.It is one function shared by the two trainers.Source: §A.1 · p. 11
- The configurations differ in latent-term and target keys, context split, masking policy, and warm-up.These are configuration differences between the two arms.Source: §A.1 · p. 11
- By default, ds validates on 32 tables and jepa on 64 held-out batches.The trainers also differ in the number of held-out batches used for validation.Source: §A.1 · p. 11
A.2 Collapse across priors
With the value head reading the predictor, runs across three priors collapse or fail to learn the value task; the encoder-read reference performs much better on the paper’s prior. The reported sequence puts predictor-output flattening before target-field rank loss.
A.2 Collapse across priors
- Predictor-read runs across three priors scored .570 to .585 mean accuracy and at most -.003 mean R2, versus .750 accuracy and .490 R2 for the encoder-read reference.The predictor-read runs ended within 1.3% of their constant map, with target rank at most 26.7.Source: §A.2 · p. 11, Table 2 · p. 12
- On two priors where the value head learned, predictor-output cosine crossed .9 before value error returned to .98 of the constant map, and target rank fell later.The output crossed .9 after the value-error minimum; target effective rank fell below 10 only afterward.Source: §A.2 · p. 11, Figure 5 · p. 12
- On TabICL v1, value error never improved over the constant map, although predictor-output effective rank stayed at most 15.4.Its minimum value error was 1.000 of the map.Source: §A.2 · p. 11, Table 2 · p. 12
A.3 Evaluation protocol
Evaluation averages scores over random halves of three OpenML suites and tests paired per-dataset outcomes; Table 3 tracks JEPA snapshots against two value-only checkpoints.
A.3 Evaluation protocol
- Each model receives at most 1,024 context rows and 64 F-test-selected features, with scores averaged across three random halves of up to 512 test rows each.An exact two-sided sign test drops ties.Source: §A.3 · p. 12
- Table 3 reports per-dataset wins and losses for JEPA snapshots against both same-step and final ds checkpoints across classification, its many-class stratum, and regression.A dash marks a snapshot step past ds stopping.Source: Table 3 · p. 13
A.4 Snapshots along the trajectory
The appendix catalogs snapshot comparisons, sign-test counts by stratum and benchmark, and per-dataset scores for the two arms and linear baselines.
A.4 Snapshots along the trajectory
- Table 3 lists two win-loss counts for each scored JEPA snapshot: against same-step ds and against the final ds checkpoint.The section quotes the first and last snapshot comparisons.Source: §A.4 · p. 13
A.5 Sign tests by stratum and benchmark
- Table 4 reports JEPA-versus-ds final-checkpoint wins and losses with exact two-sided sign-test p-values, grouped by class stratum and benchmark.Its final block repeats counts with one entry per dataset name.Source: Table 4 · p. 13, §A.5 · p. 13
A.6 Per-dataset scores
- Table 5 lists final-checkpoint scores for both arms and each task's linear baseline, grouped by benchmark and dataset, with means and arm win counts.Classification entries include accuracy and class count k; regression entries use R2.Source: Table 5 · pp. 14–16, §A.6 · p. 14