model.py: a 170-line PyTorch implementation of the neural architecture.
prior.py: a 348-line generator for synthetic supervised-learning datasets.
The striking thing is not merely that the implementation is short. It exposes the conceptual decomposition of a tabular prior-data fitted network (PFN):
turn scalar columns into label-aware cell embeddings;
learn each column as a set of training observations;
compress each row across features;
let test rows query the labeled training set;
train that entire mapping on a broad prior over synthetic datasets.
This is a faithful implementation of the core inference architecture, including repeated feature grouping, induced column attention, RoPE, and QASSMax. It is not a drop-in implementation of the complete TabICLv2 product: there are no released weights here, no pretraining loop, no scikit-learn interface, no missing-value/outlier pipeline, no many-class wrapper, and none of the full repository’s inference-time memory machinery. Those omissions are explicit in the README.
2 1. The statistical object: amortized inference over datasets
Let a supervised dataset be
\[
\mathcal D = \{(x_i,y_i)\}_{i=1}^{n_{\mathrm{tr}}},
\qquad x_i\in\mathbb R^m,
\]
and let \(x_*\) denote one or more unlabeled test rows. A Bayesian model with latent mechanism \(\phi\) would form
At the population optimum, with sufficient capacity and optimization, \(q_\theta\) approximates the posterior predictive induced by the synthetic prior. Inference on a real table is therefore amortized: the forward pass performs the work that an ordinary estimator would do through optimization, integration, or tree construction.
This makes the two source files inseparable conceptually:
prior.py specifies what kinds of data-generating mechanisms the learner should expect;
model.py specifies the computational family available for learning the resulting posterior-predictive map.
Calling the result “Bayesian” needs this qualification. It is Bayesian with respect to the chosen synthetic task distribution only to the extent that the learned approximation, finite network, training curriculum, and real-data distribution support that interpretation. Prior misspecification becomes learned inductive bias.
3 2. Tensor map of the forward pass
Write
\(B\): number of tables in a batch;
\(n=n_{\mathrm{tr}}+n_{\mathrm{te}}\): rows per table;
\(m\): columns;
\(g=3\): feature-group width;
\(d=128\): cell embedding width;
\(c=4\): row-summary tokens;
\(D=cd=512\): row embedding width;
\(r=128\): inducing tokens per column.
The forward path is:
Show code
flowchart LR X["X: B × n × m"] --> S["train-only standardization"] S --> G["circular repeated groups: B × n × m × g"] G --> E["linear cell embedding: B × n × m × d"] Y["y_train"] --> TAE["target-aware embedding"] TAE --> E E --> C["3 induced column-attention blocks"] C --> R["prepend c CLS tokens per row"] R --> RT["3 RoPE row-attention blocks"] RT --> Z["concatenate CLS outputs: B × n × D"] Y --> IY["ICL target embedding"] IY --> Z Z --> ICL["12 QASSMax ICL blocks; keys/values restricted to train"] ICL --> O["test-only MLP outputs"]
flowchart LR
X["X: B × n × m"] --> S["train-only standardization"]
S --> G["circular repeated groups: B × n × m × g"]
G --> E["linear cell embedding: B × n × m × d"]
Y["y_train"] --> TAE["target-aware embedding"]
TAE --> E
E --> C["3 induced column-attention blocks"]
C --> R["prepend c CLS tokens per row"]
R --> RT["3 RoPE row-attention blocks"]
RT --> Z["concatenate CLS outputs: B × n × D"]
Y --> IY["ICL target embedding"]
IY --> Z
Z --> ICL["12 QASSMax ICL blocks; keys/values restricted to train"]
ICL --> O["test-only MLP outputs"]
The default classifier contains 27,552,250 parameters; the otherwise-default 999-quantile regressor contains 28,560,215. These were counted directly from the checked-out implementation.
4 3. Input normalization and repeated feature grouping
4.1 3.1 Train-only normalization
The first operation in forward standardizes every feature using only the training prefix:
where the variance divisor is \(n_{\mathrm{tr}}\), not \(n_{\mathrm{tr}}-1\). This is the correct leakage boundary: test covariates do not influence the location or scale used to transform themselves.
It also establishes a strong input contract that is not validated explicitly: training rows must precede test rows, and y.shape[1] is taken to be \(n_{\mathrm{tr}}\).
4.2 3.2 Why group features at all?
A scalar column by itself cannot express feature interactions before row attention. NanoTabICL forms a local tuple around every feature index \(j\):
A learned linear map \(W_x:\mathbb R^3\to\mathbb R^{128}\) embeds each tuple. The implementation is the compact stack expression at model.py:L38-L40.
The powers-of-two construction is a sparse interaction design. For a general group size \(g\), offsets are \(2^\ell-1\). The paper proves that, when \(m\ge 2^g\), any pair of columns co-occurs in at most one group. For \(g=3\), therefore, \(m\ge8\) avoids redundant pairs. The operation gives the model cheap pairwise exposure while keeping exactly \(m\) tokens rather than constructing all \(O(m^2)\) pairs.
The modulo convention makes this mechanism well-defined even for very narrow tables, but its no-repeated-pair property no longer holds there.
5 4. Early target injection: turning columns into supervised tasks
For each training row, NanoTabICL adds a learned embedding of \(y_i\) to every feature token:
This is implemented at model.py:L40-L41. For classification, ClassEmbedding computes a one-hot vector and applies a biased linear layer. For regression it applies a linear map to the scalar target.
This early label injection is more than an implementation detail. It makes each column transformer a small supervised in-context learner: it can represent relationships between a feature’s empirical distribution and the associated outcomes before the columns interact.
Labels are injected again after row compression, using a separate map into \(\mathbb R^{512}\) (model.py:L53-L56). The two injections serve different coordinate systems:
late: identify labeled examples to the dataset-level in-context learner.
6 5. Column transformer: induced set attention
Each feature \(j\) is treated as a set of row embeddings. Full self-attention among \(n\) rows would cost \(O(n^2d)\) per feature. InducedTransformerBlock instead introduces \(r=128\) learned inducing vectors \(I\in\mathbb R^{r\times d}\) and performs two cross-attention steps:
The first step compresses the labeled training column into \(r\) prototypes; the second broadcasts that summary to every train and test row. Complexity falls from \(O(n^2d)\) to \(O(nrd)\) per column when \(r\ll n\).
a test row may query a representation learned from the training rows;
a test row does not enter the column summary;
one test observation cannot influence another through this stage.
The first compression attention uses QASSMax because its context length is the potentially large number of training rows. The second step attends to a fixed 128 tokens and uses ordinary softmax.
7 6. Row transformer: feature interaction and compression
After column reasoning, the model prepends \(c=4\) learned [CLS] tokens to each row and applies three ordinary transformer blocks across columns (model.py:L47-L51). The final block computes outputs only for the four summary queries. Their normalized representations are concatenated:
Using four summaries rather than one gives the compression stage several learned “views” of a row. The final q_max_idx=4 is also a useful memory optimization: the last layer need not materialize updated feature-token outputs that will immediately be discarded.
7.1 6.1 Rotary positional encoding
Row attention uses Rope. For each frequency \(\omega_k=\theta^{-k/(d_h/2)}\) and token position \(p\), pairs of coordinates are rotated:
Because \(R_p^\top R_q=R_{q-p}\), the query-key inner product depends on relative position:
\[
(R_pq)^\top(R_qk)=q^\top R_{q-p}k.
\]
Thus the transformer can distinguish column order without adding an absolute embedding to token values. The sine/cosine tables are cached and extended lazily.
This also means the architecture is not intrinsically invariant to column permutations. The paper relies on column permutations and related augmentation during pretraining to encourage robustness; those wrappers are not present here. Circular grouping itself also depends on feature order.
8 7. Dataset-level in-context learning
The compressed row vector \(z_i\) receives the second target embedding on training rows, then passes through twelve \(D=512\) transformer blocks. For the first eleven blocks, every row is a query but only training rows are keys and values:
The last block computes only test queries against training keys/values (model.py:L55-L58). A two-layer GELU MLP maps the result to class logits or quantile values.
This masking scheme has a clean causal interpretation: predictions are functions of the training set and the corresponding test row, not of other test rows. In a direct numerical check, changing one test row by a large amount changed its own logits (maximum absolute difference \(0.231\)) while another test row’s logits changed by exactly \(0\).
That property is stronger than merely avoiding target leakage. It yields batch-composition invariance among test observations and permits test rows to be evaluated together without transductive interaction.
Suppose one relevant key has logit \(s\) and the other \(n-1\) keys have logit \(0\). Its weight is
\[
a_*(n)=\frac{e^s}{e^s+n-1},
\]
which tends to zero as \(n\) grows. A model trained at one sample size can therefore become too diffuse at a much longer context even if it continues to identify the same relevant example.
9.2 8.2 Query-aware scalable softmax
QASSMax rescales every query coordinate before the dot product:
\(b(\log n)\) is a two-layer MLP producing a distinct scale for every head-coordinate pair;
\(g(q)\) is a two-layer MLP producing content-dependent coordinate scales;
the gate lies in \((0,2)\);
the final gate layer is zero-initialized, so the content modulation starts at exactly one.
The rough asymptotic intuition is that if logits scale like \(s\log n\), then \(e^{s\log n}=n^s\), allowing a sufficiently relevant key to compete with a denominator containing \(n\) terms. QASSMax is richer than a fixed temperature: it learns which head dimensions should sharpen as context grows and can adjust that sharpening to the current query.
NanoTabICL applies it precisely where context length varies with sample size:
inducing queries over training rows in each column;
dataset-level ICL over training rows.
It does not apply it to attention over the fixed inducing set or to row attention over features.
It reuses PyTorch’s packed projection weights but calls scaled_dot_product_attention directly, enabling PyTorch to select an optimized attention kernel. Explicit deletion of temporary \(Q,K,V\) references helps shorten tensor lifetimes during inference.
10 9. Output semantics
10.1 9.1 Classification
With max_classes=out_dim=C, the final \(C\) numbers are logits and imply
The nano architecture itself supports at most the number of classes used to construct ClassEmbedding, typically ten. The full TabICLv2 system adds mixed-radix label views and hierarchical classification for \(C>10\); those mechanisms are absent here.
10.2 9.2 Regression
The intended regression head has 999 outputs corresponding to quantile levels \(\tau\in\{0.001,\ldots,0.999\}\). Pretraining uses pinball loss
Predicting the whole conditional quantile function is more informative than emitting a mean: it represents heteroskedasticity, skew, and tail uncertainty. A point prediction can be approximated by averaging the quantile grid.
The network does not impose monotonicity across quantiles. Sorting, isotonic correction, tail handling, target transformations, and ensembling belong to the full inference wrapper. The nano README correctly warns that users must standardize regression targets and back-transform predictions themselves.
11 10. The nanoprior: a distribution over learning problems
The prior is best read as a probabilistic program that samples a structural causal-looking computation graph—not as a claim that real tables literally arise from these mechanisms.
where \(A,B_i,C_j\) are standard Cauchy draws. Ordering makes the graph acyclic. The shared \(A\) controls global density, while \(B_i\) and \(C_j\) induce heterogeneous out- and in-degrees. Heavy Cauchy tails generate both near-deterministic edge probabilities and exceptions, creating much greater graph diversity than one fixed Bernoulli rate.
Observed \(x\) and \(y\) columns are attached to random graph nodes, with replacement. Multiple observed variables can therefore share a latent node, while irrelevant latent nodes may remain in the sampled graph.
11.2 10.2 Node computation and feature importance
At a root, the code samples random points; at a non-root it transforms and combines parent tensors. A node width includes enough coordinates for its assigned observed variables plus a log-uniform number of latent coordinates. After transformation, each coordinate is standardized and reweighted, then the average row \(\ell_2\) norm is normalized (prior.py:L93-L106).
followed by a random permutation and norm scaling (prior.py:L317-L326). Wide log-uniform ranges for \(q\) and \(\sigma\) span nearly equal importance, smooth decay, and sparse dominance by a few coordinates.
concatenates all parent states and applies one random function, or
transforms each parent separately and combines outputs by sum, product, maximum, or log-sum-exp.
These choices generate additive, multiplicative, threshold-like, substitutable, and interaction-heavy structures without hand-writing a collection of named textbook DGPs.
non-additive interactions and heteroskedastic shapes
The random activation library adds monotone and non-monotone smooth functions, thresholds, rounding, ranks, periodicity, powers, and nonsmooth transformations. Random matrices vary Gaussian mixing, importance-weighted mixing, singular spectra, signed kernel structure, and activation-transformed weights.
11.5 10.5 Gaussian processes through random Fourier features
The most mathematically explicit component is rand_gp_func. By Bochner’s theorem, a stationary positive-definite kernel has a spectral representation. If \(\omega_s\) are drawn from the kernel’s spectral density and \(b_s\sim\mathrm{Unif}[0,2\pi]\), then
Varying \(a\) varies spectral tail decay and therefore function smoothness. In the rotationally invariant construction, a spectral density with tail \(g(\omega)\asymp\|\omega\|^{-q}\) yields GP paths with effective Sobolev smoothness approximately \((q-d)/2\) under the paper’s conditions. The alternate axis-aligned product spectrum creates different smoothness across coordinate interactions.
11.6 10.6 Numerical and categorical converters
rand_converter extracts an observed variable from latent node state and may transform what propagates downstream.
Numerical columns are either direct latent coordinates or Kumaraswamy-warped values, \[K_{a,b}(u)=1-(1-u^a)^b,\] after min-max scaling.
Categorical columns arise from nearest-center assignment or sampling from a softmax with random temperature and imbalance. Their downstream representation may remain continuous, become a center, pass through another random function, or become an integer index.
The distinction between the observed value and the state propagated through the graph lets categorization act as measurement, discretization, or an actual state transition.
11.7 10.7 Filtering for learnable tasks
rand_dataset_filtered repeatedly generates a dataset until shallow ExtraTrees out-of-bag predictions beat the mean-label baseline. For per-row improvement
it draws 200 bootstrap resamples and accepts when fewer than 5% have nonpositive average improvement. Classification labels are one-hot encoded so the same squared-error check can be used.
This filter removes effectively random or pathological tasks. It also changes the prior: the learner is trained conditional on a modest tree ensemble detecting signal. The resulting task distribution therefore favors relationships visible to shallow axis-aligned partitions, even though the generator itself contains much broader function families.
12 11. Computational scaling
Suppressing batch, head, layer, and embedding constants, the main costs are:
Stage
Time
Activation scale
Driver
repeated grouping
\(O(nmg)\)
\(O(nmg)\) before projection
gather \(g\) values per cell
induced column attention
\(O(mnr)\)
\(O(mnr)\) attention scores
\(r\) inducing summaries per column
row attention
\(O(n(m+c)^2)\)
\(O(n(m+c)^2)\)
feature interaction within rows
dataset ICL
\(O(n\,n_{\mathrm{tr}})\)
\(O(n\,n_{\mathrm{tr}})\)
train-key attention for every row
With fixed \(r,c,d\), this gives the paper’s headline scaling
\[
O(n^2+nm^2),
\]
rather than cell-level attention of order \(O(n^2m+nm^2)\). The architecture avoids multiplying the expensive row-to-row term by the number of features: it compresses features before full dataset ICL.
The code contains two practical reductions:
final row attention emits only the CLS queries;
final ICL emits only test queries.
But this educational implementation still lacks the main repository’s chunking, offloading, and inference wrappers, so peak memory on very large tables should not be inferred from the asymptotic expression alone.
13 12. Empirical sanity checks on the clone
The checked-out source was exercised in a clean isolated environment with PyTorch and scikit-learn:
finite (128, 12) feature matrix and (128, 1) target
README’s floating classification labels
fails: one_hot requires LongTensor
The first four checks support the intended tensor and leakage contracts. The last one identifies a documentation defect: README.md:L34 casts classification labels to float, but the current ClassEmbedding implementation requires integer class indices. The executable __main__ example correctly uses integer labels.
14 13. Engineering review
14.1 13.1 What is especially good
The code mirrors the conceptual blocks. The architecture can be understood from the forward method without navigating a framework of registries and wrappers.
Leakage control is structural. Train-only normalization and training-only keys/values are visible at the call sites rather than hidden in a mask assembled elsewhere.
The asymptotic design is legible. Inducing rows, CLS compression, and final query slicing correspond directly to the desired complexity reductions.
Modern PyTorch kernels are reused. Direct scaled_dot_product_attention allows optimized backends while retaining custom QASSMax and RoPE.
The prior is genuinely modular. Graphs, functions, converters, matrices, weights, points, and filtering can be studied or replaced separately.
The implementation is honest about scope. The README distinguishes the architecture from the complete pretrained estimator.
14.2 13.2 Important limitations and failure modes
No package or dependency manifest. Imports imply PyTorch, NumPy, and scikit-learn, but installation is not reproducible from repository metadata.
No automated tests or CI. The only model check is the script-level shape smoke test; the prior’s plotting entry point is not a unit test.
Classification dtype mismatch in the README. Labels must be integer tensors for one-hot encoding.
Little input validation. The model does not check ranks, equal batch dimensions, \(0<n_{tr}<n\), class range, feature count, or divisibility/evenness of head dimensions.
NaNs propagate through the model. The prior replaces nonfinite generated values, but real input preprocessing is deliberately omitted.
A private PyTorch function is used.torch.nn.functional._in_projection_packed is underscored and therefore more version-fragile than a public API.
RoPE assumes an even head dimension. Splitting into two equal halves is mathematically required but not asserted.
Only prefix masking is supported. Arbitrary train/test masks, padding masks, missing labels, and ragged tables need a wrapper or architectural changes.
Regression is not checkpoint-identical by default. The nano model uses LayerNorm bias, whereas the paper’s regression checkpoint is bias-free.
Raw quantiles may cross. No monotonicity or distributional postprocessing is included.
The prior filter has no attempt limit.while True can run indefinitely for unlucky specifications, and failed attempts have no diagnostics.
Reproducibility is global and mixed. NumPy and PyTorch global RNGs are both used; no generator is passed through the call graph. The fixed bootstrap RNG does not make generation itself deterministic.
The prior is CPU-bound in places. It calls .numpy() for scikit-learn filtering, precluding transparent device placement.
Degenerate raw draws are possible. Filtering and the full system’s postprocessing are important; rand_dataset_plain alone should not be treated as a guarantee of a useful task.
14.3 13.3 The most valuable compact test suite
A small suite could preserve the educational character while covering the contracts that matter:
output shapes for classification and regression;
integer-label requirement and class-bound errors;
train-only standardization invariance to test-set composition;
changing test row \(j\) cannot change prediction for test row \(k\ne j\);
permuting test rows permutes outputs equivariantly;
invalid/even head-dimension checks;
finite outputs on constant columns;
RoPE cache extension and dtype/device preservation;
QASSMax identity of the zero-initialized content gate;
deterministic prior draws after setting both RNG seeds;
prior key/shape/category-range invariants;
bounded-attempt behavior for filtered generation.
15 14. How to modify the system without breaking its logic
15.1 Replace the synthetic prior
Preserve the task-level interface: sample whole labeled tables, then split rows. A new prior changes the model’s implicit statistical assumptions. It should be evaluated not only for visual diversity but for coverage of structural properties—smoothness, sparsity, interaction order, categorical imbalance, noise, missingness, and sample/feature regimes.
15.2 Add a new column encoder
Keep the no-test-contamination invariant: summaries used by test rows must aggregate training rows only. A full self-attention replacement must carry an explicit key/value mask or it will silently become transductive.
15.3 Change feature grouping
Changing the offset design alters which low-order interactions are cheaply exposed and breaks checkpoint compatibility through the first projection. Analyze pair coverage and duplication, not just group width.
15.4 Add positional invariance
If column-order invariance is desired, it cannot coexist naively with the current RoPE/grouping semantics. Options include aggressive permutation ensembling, permutation augmentation during training, or a set-equivariant row encoder; each changes computational cost or inductive bias.
15.5 Add missing-value support
Do not simply replace NaNs with zero after standardization. Zero is the training mean in standardized space and confounds missingness with an observed central value. Add missingness indicators, learned missing tokens, or a preprocessing distribution matched during synthetic pretraining.
15.6 Implement pretraining
The essential unit is a batch of tasks, not a batch of iid rows. Each task must provide X_train_and_test and y_train, with held-out targets used only in the loss. Curriculum over row counts is central to long-context generalization; QASSMax does not remove the need to expose the model to longer contexts.
16 15. Bottom line
NanoTabICL succeeds as an executable architectural note. Its brevity reveals the central idea of TabICLv2: use target-aware set attention to turn columns into supervised representations, use RoPE row attention to compress feature interactions, and use a long-context transformer as an amortized dataset-level learner. QASSMax addresses the specific softmax dilution problem created by extrapolating to larger training sets, while the synthetic DAG/function prior defines the universe of problems over which this inference algorithm is learned.
The model file is close to the mathematical architecture; the repository as a whole is intentionally far from a production estimator. The missing pieces—weights, pretraining, preprocessing, validation, ensembling, distributional correction, many-class logic, tests, and scalable inference wrappers—are not cosmetic. They are what turn a compact neural mechanism into a reliable tabular learning system.