TabMosaic: Unified In-Context Diffusion for Tabular Generative Tasks

1 INTRODUCTION

• Existing tabular generative models typically require training or fine-tuning for each new table, while in-context models predict individual cells rather than arbitrary target sets jointly.This leaves multi-cell generation and imputation assembled from separate predictions rather than modeled directly.Source: §1 · p. 1
• The paper unifies generation, imputation, and conditional generation as learning a joint conditional distribution over arbitrary target sets on unseen tables.The tasks differ by which cells are observed and which are queried.Source: §1 · p. 1, Figure 1 · p. 2
• TabMosaic uses in-context diffusion to model mixed numerical and categorical targets jointly, without imposing an autoregressive order.Its two-stage pretraining extends from short, complete tables to long, partially observed contexts; multiple corruptions share one context encoding in stage two.Source: §1 · p. 2
• Under the pretraining prior, the diffusion objective is minimized by the posterior predictive and its bound controls conditional negative log-likelihood and divergence from that predictive.The divergence guarantee holds up to a model-independent constant.Source: §1 · p. 2
• TabMosaic ranks first on generation, imputation, and class-conditional generation across 48 datasets against 20 baselines.The baselines span dataset-specific generators and imputers, in-context predictors, and training-free conditional samplers.Source: §1 · p. 2

2.1 IN-CONTEXT LEARNING AS POSTERIOR PREDICTION

• At optimum, sequence models trained on mixtures of latent tasks approximate Bayesian posterior prediction by conditioning on prompts and marginalizing over tasks.For tables, this account extends to arbitrary target sets when the pretraining objective is their likelihood.Source: §2.1 · pp. 2–3

2.2 GENERATIVE MODELS FOR TABLES

• Per-dataset diffusion generators and imputers model target conditionals through inpainting or autoregression, but can restrict target types, impose column order, or require complete contexts.The passage describes these approaches alongside TabNAT and GReaT (Borisov et al., 2022b; Zhang et al., 2025a).Source: §2.2 · p. 3
• Cross-table models CTSyn (Lin et al., 2025) and TabFORGE (Jiang et al., 2026) fine-tune on each new table and target unconditional generation.Neither answers an unseen-table query without a gradient update.Source: §2.2 · p. 3
• In-context predictors chain single-cell predictions to generate or impute, so their objective does not model a target set jointly and generation follows a chosen order.The passage names Hollmann et al. (2025) and Zhang et al. (2025c).Source: §2.2 · p. 3
• Classifier-energy samplers TabEBM (Margeloiu et al., 2024) and TabPFGen (Ma et al., 2024) require a class label; their classification objective leaves the cell distribution unconstrained.TabPFGen also imputes by sampling missing cells while holding observed cells fixed.Source: §2.2 · p. 3

3.1 ONE FAMILY OF TASKS: CONDITIONAL QUERIES OVER CELL SETS

• A task is a conditional joint distribution over target cells given the query’s observed cells and context, with tasks differing only in the chosen target set.The setting allows multiple query rows, handled independently given the context.Source: §3 · p. 3, §3.1 · pp. 3–4

3.2 IN-CONTEXT JOINT DIFFUSION OVER CELL SETS

• The reverse process jointly models mixed numerical and categorical targets: it corrupts only targets, then conditions each reverse factor on the full mixed state.Numerical targets use an EDM Gaussian kernel; categorical targets use GIDD, whose uniform-noise component allows already unmasked cells to be revised (von R¨utte et al., 2025).Source: §3.2 · p. 4
• The joint distribution is normalized for every target set, and iterative reverse steps create dependencies across cells and modalities; one step yields a conditionally independent product.The joint is factorized across time, not across cells or modalities.Source: §3.2 · p. 4

3.3 TRAINING OBJECTIVE AND GUARANTEE

• Target sets are sampled independently of cell values and noise level, placing every target pattern and noise level within the training distribution.Consequently, inference can start the reverse process at any noise level and for any fraction of target cells.Source: §3.3 · p. 4
• The objective upper-bounds conditional negative log-likelihood, and the posterior predictive minimizes it; as reverse steps increase without bound, the model equals that predictive distribution.The bound and minimizer results assume target sets are drawn independently of cell values.Source: §3.3 · p. 5

4.1 ARCHITECTURE

• Each context or query cell is a token, with state embeddings distinguishing targets, observations, and missing cells; the shared per-cell head predicts numerical values or categorical logits.No token has a column index or label role, and query rows attend only to context, allowing independent parallel answers.Source: §4.1 · p. 5

4.2 PRETRAINING: PRIOR, ESTIMATOR, AND STAGES

• Pretraining widens from short, complete tables to long, incomplete contexts; Stage 2 uses four corruptions per table, sharing one context encoding to reduce computation and gradient variance.Stage 1 trains on 32M complete tables; Stage 2 trains on 2.56M tables with missing context cells and a cell budget of 250,000.Source: §4.2 · p. 6

4.3 INFERENCE: ONE PROCEDURE FOR EVERY TASK

• One pretrained procedure handles generation, imputation, and class-conditional generation by changing the observed and target cells, reverse-chain start, step count, and number of trajectories.The context is encoded once and reused across reverse steps, query rows, and trajectories.Source: §4.3 · p. 6

5.3 RQ2. POSTERIOR PREDICTIVE FIDELITY

• At n = 1,024, GP NLL moves from −0.53 to −0.71 versus −0.77 for the exact posterior predictive; GP+Dirichlet categorical-cell NLL is 0.114 versus 0.103.The NLL approaches the exact posterior predictive on every tested prior, including beyond training context sizes.Source: Figure 3 · p. 8, §5.3 · pp. 8–9

5.2 RQ1. GENERATIVE PERFORMANCE

• Generation has the best mean utility rank (4.1 vs. TabNAT’s 4.6); at 10,240 rows and 32 columns, it takes 53 s versus 6–7 min for TabPFN-3 and LimiX.Per-table models sample in seconds once trained, but incur training cost on every table.Source: §5.2 · p. 8, Figure 4 · p. 8
• Imputation Elo is 1458 versus 1201 for TabICLv2 and 1108 for tuned ReMasker; mean error is 0.129 versus 0.146 and 0.157.The categorical-cell error is 0.314 versus 0.375 for TabICLv2.Source: §5.2 · p. 8
• Class-conditional generation ranks first by Elo (1163), ahead of TabPFN-3 (1144) and TabNAT (1122), with label consistency 0.983 of real rows’ versus TabPFGen’s 1.093.TabMosaic performs this task by observing the label cell.Source: §5.2 · p. 8

5.3 RQ2. POSTERIOR PREDICTIVE FIDELITY

• For |S| ≥ 2, TabMosaic stays within 1.0–1.3 times the sampler floor as the factorized oracle degrades; it is on par with per-context-trained TabNAT.TabMosaic produces the joint distribution in one pass without an ordering.Source: Figure 5 · p. 8, §5.3 · p. 9

5.4 RQ3. TARGET-SET PERFORMANCE

• TabMosaic has the lowest relative hidden-cell error in every tested setting; its error rises from 0.62 to 0.84 with missing rate and is flat across MCAR, MAR, and MNAR.The scenarios vary missing rate, mechanism, and random-cell versus whole-column patterns.Source: Table 1 · p. 9, §5.4 · p. 9

5.5 RQ4. STAGED PRETRAINING

• Stage 2 lowers imputation error on 45/48 tables (mean 0.156 to 0.129), including 10/11 naturally incomplete contexts; generation and class-conditional utility improve on 40/48 and 26/36 tables.Gains concentrate on tables above 10k rows, beyond lengths seen in Stage 1.Source: Figure 6 · p. 9, §5.5 · p. 9

6 CONCLUSION

• TabMosaic ranks first on all three tasks across 48 datasets against 20 baselines and serves arbitrary target sets without per-table training.Its diffusion objective defines a valid joint for mixed-type target sets without an ordering, upper-bounds conditional NLL, and is minimized by the posterior predictive under the pretraining prior.Source: §6 · p. 9
• The objective’s population optimum is the posterior predictive distribution, extending the joint conditional modeling formulation to arbitrary target sets.The objective upper-bounds conditional NLL and defines a valid joint for mixed-type target sets without an ordering.Source: §6 · p. 9

AI USE STATEMENT

• The authors report using generative AI tools for experimental design, code, figure creation, proof-checking, and drafting, while retaining research decisions and responsibility.They state that authors reviewed every AI-assisted output.Source: AI USE STATEMENT · p. 10

A PROOFS

• The appendix sets c = (xobs, D), defines S as the numerical and categorical target cells, and specifies their mixed state space Z.Categorical states include each column’s classes and a MASK token.Source: A PROOFS · p. 14

A.1 PROOF OF PROPOSITION 1 (VALIDITY)

• The reverse process composes normalized cellwise Gaussian or categorical transitions into a normalized joint law over target cells.The endpoint is obtained by marginalizing intermediate states from a normalized terminal law.Source: §A.1 · p. 14
• The loss decomposes over cells, but the endpoint model retains cross-cell and cross-modality dependencies through the coupled reverse steps.The factorization of each step does not imply that the final joint factorizes across cells.Source: §A.1 · p. 14

A.2 PROOF OF THEOREM 1 (CONDITIONAL ELBO)

• For numerical targets, the continuous-time bound uses a likelihood-weighted squared denoising error, and positive reweighting preserves its posterior-mean minimizer.The likelihood weight is wlik(t) = σ′(t)/σ(t)^3; replacing it with the EDM weight changes optimization geometry, not the target.Source: §A.2 · p. 14
• Categorical stepwise KL terms telescope into the continuous-time weighted GIDD form, while product-factor KLs make each step’s bound additive over cells.The categorical terms are grouped by modality, as are the numerical terms.Source: §A.2 · p. 14

A.3 PROOF OF COROLLARY 1 (THE OBJECTIVE TARGETS THE POSTERIOR PREDICTIVE DISTRIBUTION)

• The expected objective upper-bounds expected KL divergence to the posterior predictive when target patterns are independent of the sampled table.Conditioned on the target pattern, observed cells, and context, the clean target set follows the posterior predictive.Source: §A.3 · p. 15

A.4 PROOF OF COROLLARY 2 (THE POSTERIOR PREDICTIVE MINIMIZES THE OBJECTIVE)

• The posterior predictive minimizes the objective, and with its denoiser the reverse process recovers that predictive distribution as the number of steps tends to infinity.Numerical denoising is minimized by the posterior mean and categorical denoising by the true posterior PMF.Source: §A.4 · p. 15

B KERNEL DETAILS: EDM AND GIDD

• The EDM numerical kernel and GIDD categorical kernel support mixed-type diffusion; GIDD’s uniform component allows decoded tokens to be revised.The EDM weight shares the likelihood-weight minimizer, while MASK remains a legal categorical state throughout diffusion.Source: B KERNEL DETAILS: EDM AND GIDD · pp. 15–16

C MULTI-CORRUPTION ESTIMATOR OF THE OBJECTIVE

• Averaging K independent corruptions for the same table, time, and target pattern leaves the estimator and gradient unbiased, while reducing corruption-conditional variance by 1/K.The minimizer is unchanged by K; the estimator averages losses rather than outputs.Source: C MULTI-CORRUPTION ESTIMATOR OF THE… · p. 16

D MISSINGNESS IN THE IN-CONTEXT SET

• Stage 2 masks context cells and represents each as observed, target, or missing-in-context; at inference, missing query cells are targets to impute.Training samples MCAR, MAR, and MNAR mechanisms, with numerical transforms fit only on observed context cells.Source: D MISSINGNESS IN THE IN-CONTEXT SET · p. 16

E PRETRAINING DETAILS

• Synthetic tables come from the TabICL v2 structural-causal-model prior, with numerical columns quantile-transformed using observed context cells.The label is treated as an ordinary column and may be numerical or categorical.Source: E PRETRAINING DETAILS · p. 17
• Target patterns include full query rows for generation or independently sampled subsets of query cells, with the pattern independent of cell values and diffusion time.All query cells are targets with probability 0.4; otherwise the target rate is sampled uniformly from 0 to 1.Source: E PRETRAINING DETAILS · p. 17
• Stage 1 uses 1,024-row complete tables; stage 2 samples 400–10,240 rows and up to 200 features under a 250,000-cell budget.Stage 2 uses a context fraction of 0.8, while stage 1 samples its context fraction from 0.3 to 0.9.Source: E PRETRAINING DETAILS · p. 17
• Stage 2 warm-starts from stage 1 and adds long tables, missing contexts, and K = 4 corruptions; the evaluated checkpoint receives no per-table training, tuning, or selection.The denoiser has 29M parameters, and both pretraining stages use 8 GPUs.Source: E PRETRAINING DETAILS · p. 17, Table 2 · p. 17

• The training split supplies context, capped at 8,192 rows; generation samples it randomly, while imputation uses a label-stratified subset and preserves missing cells as missing-in-context.Source: F INFERENCE DETAILS · p. 17
• Imputation starts the reverse process at tstart = 0.5 for 16 ancestral steps, while generation starts at tstart = 1 for 128 steps.The table defines tstart as the reverse-process start time and N as the number of steps.Source: F INFERENCE DETAILS · p. 18, Table 3 · p. 18
• With one step, numerical targets return conditional means and categorical targets are drawn independently; more than one step couples cells through their shared current states.Source: F INFERENCE DETAILS · p. 18
• For imputation point estimates, TabMosaic averages eight numerical trajectories and takes a categorical majority vote; distribution-level metrics instead use the trajectories as samples.Source: F INFERENCE DETAILS · p. 18
• Generation runs the full reverse process jointly over targets; class-conditional generation additionally observes each query row’s label and fixes it to the requested class.Categorical sampling uses temperature 1; lowering it trades diversity for marginal sharpness and degrades distribution-level metrics.Source: F INFERENCE DETAILS · p. 18
• The context is encoded once per query batch and reused, yielding a 25-second median generation time per table on one GPU, versus 1,696 seconds for TabNAT.Source: F INFERENCE DETAILS · p. 18

G EVALUATION PROTOCOL

• The benchmark uses the first official split of 48 tables, excluding three tables with more than 1,000 features; class-conditional generation uses 36 classification tables.Source: G EVALUATION PROTOCOL · p. 18
• Utility and label consistency average scores from Light-GBM, CatBoost, and TabPFN at default settings, with no predictor tuned on synthetic data.Source: G EVALUATION PROTOCOL · p. 18
• Generation is scored by utility, α-precision, β-recall, and authenticity, with utility comparing predictors trained on synthetic rows against real-row training.Source: G EVALUATION PROTOCOL · p. 19
• Imputation injects 30% missingness completely at random into feature cells and scores only injected cells, using scaled RMSE or categorical misclassification rate.Source: G EVALUATION PROTOCOL · p. 19
• Class-conditional generation preserves the training split’s class proportions; label agreement is 1 for every method except CTGAN, which uses a soft constraint.Source: G EVALUATION PROTOCOL · p. 19
• Elo treats each table–metric pair as a task and fits ratings to pairwise win rates, using 200 bootstrap resamples for 95% intervals.Source: G EVALUATION PROTOCOL · p. 19
• TabSyn and TabDiff are excluded from some tasks because their papers do not claim the evaluated forms, while CTSyn and TabFORGE lack public code at writing.Source: G EVALUATION PROTOCOL · p. 20

H GENERATION AND IMPUTATION WITH IN-CONTEXT PREDICTORS

• In-context baselines pretrained on single cells form generation joints by chaining column conditionals, while their imputers fill cells separately or sequentially rather than sampling a joint.For multiple masked cells, LimiX’s product of per-cell distributions parameterizes no joint; TabMosaic couples target cells at each reverse step and bounds their joint likelihood.Source: H GENERATION AND IMPUTATION WITH… · p. 20

I DISTANCE TO THE POSTERIOR PREDICTIVE: PROTOCOL

• The synthetic test priors include Gaussian-process numerical relationships, a four-column Dirichlet categorical model, and a mixed GP–Dirichlet model with class-dependent numerical correlations.Source: I DISTANCE TO THE POSTERIOR… · pp. 20–21
• A 6-block, 6.7-million-parameter model trained on an equal mixture of the three priors is evaluated across context sizes and checkpoints using sampled-value negative log-likelihood.Source: I DISTANCE TO THE POSTERIOR… · p. 21
• The protocol tests TabMosaic on the Dirichlet prior with target sets of one to four cells, measuring total variation to the exact joint PMF alongside sampling and factorized-oracle references.This pretrained model had never seen those priors, so its distances include a prior-mismatch term not bounded by the corollaries.Source: I DISTANCE TO THE POSTERIOR… · pp. 20–21

J TARGET-SET SCENARIOS: PROTOCOL

• Held-out imputation scenarios vary missingness mechanism, rate, and the number of wholly hidden columns; Table 1 compares nine methods on the 24 tables where all completed.The training split remains the context, while only test-row target cells vary; larger tables were sometimes unfinished within the compute budget.Source: J TARGET-SET SCENARIOS: PROTOCOL · p. 21

K WHAT THE SECOND STAGE ADDS

• Stage 2 adds long tables and missing-context handling; on 11 naturally incomplete tables, Stage 1 uses prefilled context whereas the final model reads incomplete context directly.Source: K WHAT THE SECOND STAGE ADDS · p. 21
• Imputation error is lower after Stage 2 on 45 of 48 tables (mean 0.129 against 0.156), while generation and class-conditional utility improve on 40 of 48 and 26 of 36 tables, respectively.Both utility gains concentrate above 10k rows: +0.04 on 9 of 9 generation tables and 4 of 4 class-conditional tables, lengths absent in Stage 1.Source: K WHAT THE SECOND STAGE ADDS · p. 22

• Tables 5–7 report every method’s generation, imputation, and class-conditional generation results across the 48 and 36 tables, including reference baselines and in-context methods.Source: L ADDITIONAL RESULTS · p. 22
• Table 5 reports generation Elo with 95% bootstrap intervals, per-metric means, and median per-table runtime, with Elo tasks spanning utility, α-precision, β-recall, and authenticity.Source: Table 5 · p. 22
• Table 6 reports imputation Elo and errors on 30%-masked feature cells, including a second Elo for the 41 tables completed by every method.Source: Table 6 · p. 23
• Table 7 reports class-conditional generation Elo on 36 classification tables, including utility, class-macro precision and recall, and label-consistency gap.Source: Table 7 · p. 23
• Table 8 extends target-set scenarios to nine methods and reports held-out per-cell error relative to mean/mode imputation across missingness rates, mechanisms, and hidden-column counts.It covers 24 tables where every method completed, or 23 for four hidden columns; values below 1 indicate improvement over mean/mode.Source: L ADDITIONAL RESULTS · p. 24, Table 8 · p. 24