Architecture design

Overview

Architecture design is the step of deciding how a neural network is built before any training begins. It covers how many layers to use, how wide they are, which attention pattern connects the inputs, how raw data is turned into vectors, which output heads produce the final answers, and how the total parameter budget is spent. It also ties the structure to the training objective, since the two must fit together.

It is the first and most consequential step in building a foundation model, because later choices such as data, optimiser and scale all work inside the limits that the architecture sets. A poor design cannot be rescued by more compute, and a good one can make training cheaper and results stronger. For models that learn from raw tables, the design problems are special: rows and columns both matter, types are mixed, values are missing and column names may be absent. This page explains the main choices and the reasoning behind them, with a focus on tables. See also /model/architecture.

What it is

An architecture is the blueprint of a model. It fixes the shape of computation, while training later fills in the numbers.

The main design decisions

• Input encoding. How raw values become vectors the network can process.

• Backbone. The repeated block that does most of the work, such as a transformer layer.

• Attention pattern. Which positions may look at which other positions.

• Depth and width. The number of blocks and the size of vectors.

• Output heads. The small networks that turn internal vectors into predictions.

• Parameter budget. The total number of weights and how they are split between parts.

• Objective coupling. How the loss and the structure fit together.

An everyday analogy

Designing a building is a fair comparison. Before laying a brick, an architect chooses the floor plan, the number of storeys, where corridors connect rooms, and where the plumbing runs. Builders can work hard, but they cannot make a corridor appear where none was planned. In a network, attention patterns are the corridors, layers are the storeys, and heads are the exits to the street.

What it is not

• It is not hyperparameter tuning. Learning rate and batch size are training settings, not structure, although they interact.

• It is not data design. The data pipeline matters just as much, as covered under synthetic data generation.

• It is not fixed forever. Many designs are revised after experiments.

Where it came from

Architecture design has gone through several eras.

• Hand-built feature pipelines and shallow models dominated early machine learning. The designer shaped features, and the model was simple.

• Convolutional networks in the 1990s and 2010s encoded the idea that nearby pixels matter. See CNN and ResNet, where skip connections made very deep networks trainable.

• Recurrent networks such as LSTM and GRU handled sequences one step at a time.

• The transformer (2017, Attention Is All You Need) replaced recurrence with attention, allowing parallel training and long-range links.

• Scaling laws (2020) suggested that loss falls smoothly with model size, data and compute, which turned architecture and budget choices into a planning problem.

• Efficiency work produced mixture of experts, FlashAttention and state-space models such as Mamba.

• Tabular deep learning moved from TabNet and FT-Transformer to SAINT, which used attention across rows as well as columns.

• Prior-data fitted networks (2022) and their successors such as TabPFN and TabICL designed architectures specifically for in-context learning on whole tables.

For the lineage that leads to tabular foundation models, see /model/research.

How it works

Designing an architecture is a sequence of decisions. This section follows the order in which a designer typically thinks, using tables as the running case.

Step 1: State the task and the data shape

Ask what goes in and what comes out.

• Input: a table with N rows, M columns, mixed types, possible gaps.

• Output: a prediction per row, a filled-in cell, an anomaly score, or a label for the whole table.

• Constraint: variable N and M, no fixed column order, possibly no names.

These facts drive every later choice. For example, a model that must accept any number of columns cannot use a fixed-size input layer.

Step 2: Choose the input encoding

The encoder turns each cell, row or column into vectors.

• Numeric values. Scale them, then pass through a small network or a set of learned bins. Robust scaling and clipping reduce outlier damage.

• Categorical values. Map to learned embeddings, or to an index plus column-specific embedding. High-cardinality columns may need hashing or text-based embedding.

• Text and dates. Use a text encoder, or split dates into parts such as year, month and weekday.

• Missing values. Add a mask flag, and give the model a learned embedding for the unknown state.

• Column identity. Use a column embedding learned from the data rather than a fixed name, so the model works with no names. If names exist, they can add extra signal.

The analogy is an interpreter. Each cell speaks its own language, such as numbers, labels or dates, and the encoder translates all of them into the one language the network understands.

Step 3: Choose the backbone and attention pattern

The backbone is the repeated block. In most modern models it is a transformer block with attention and a feed-forward layer.

For tables, the key question is which cells may attend to which.

• Full attention over all cells. Simple but expensive, since cost grows with the square of the number of cells.

• Attention across columns within a row. Lets each row's features interact. This is the idea in FT-Transformer.

• Attention across rows within a column. Lets rows compare against each other, which is the key to in-context learning. SAINT uses intersample attention.

• Alternating row and column attention. Often called axial attention. It cuts cost while letting information travel across the whole table over several layers.

• Cross-attention to a small latent set. The Perceiver IO idea: compress many inputs into a few latent vectors to keep cost low.

The analogy is a spreadsheet read two ways. Reading along a row tells you about one record. Reading down a column tells you about one feature across records. A good design reads both, one after the other.

Step 4: Treat rows as a set

A table has no natural row order, and often no natural column order. Good designs respect this.

• Make the model permutation-invariant over rows, so that shuffling rows gives the same answer.

• Make it permutation-equivariant over columns, so that reordering columns reorders outputs but does not change meaning.

• Do not add fixed positional encodings to rows or columns.

This follows from the data. If you break the symmetry, the model wastes capacity learning that order is irrelevant.

Step 5: Decide how to condition on labelled rows

For in-context learning, labelled rows and query rows are processed together.

• Add the label as an extra input to each labelled row. For query rows, use a blank or mask.

• Let query rows attend to labelled rows, but prevent labelled rows from attending to query rows, so that predictions do not interfere with each other.

• Keep the pattern so that one forward pass produces predictions for all query rows.

This is where the design directly implements few-shot learning.

Step 6: Set depth, width and the parameter budget

Depth and width trade off.

• More layers allow more steps of reasoning, such as multi-hop feature interactions.

• Wider layers allow richer representations per step.

• Attention heads let the model track several kinds of relation at once.

The parameter budget is the total weights. Choose it from the compute you can spend, the data you have, and the speed you need at inference. Scaling studies suggest balancing model size and training tokens rather than growing one alone. Sparse designs such as mixture of experts give more capacity without a matching increase in compute per input.

Step 7: Design the output heads

Heads are small networks on top of the backbone.

• Classification head. Outputs class probabilities, often through a softmax.

• Regression head. Outputs a number, or a distribution over bins, which handles skew and uncertainty better than a single value.

• Imputation head. Predicts a masked cell, matched to its type.

• Anomaly head. Outputs a score for how unusual a row is.

• Semantic or sector head. Classifies a column or a table into a type or sector. See semantic sector heads.

Multiple heads can share one backbone, which shares learning across tasks.

Step 8: Tie the objective to the structure

The training goal should match what the architecture can express.

• For prediction by in-context learning, train on many tasks where the model must predict query labels from context.

• For imputation, mask random cells and predict them. See masked cell modeling.

• For representation learning, use self-supervised learning or contrastive learning.

An architecture and its objective are a pair. A design with row-to-row attention is wasted if the objective never needs rows to compare.

Step 9: Plan for efficiency

• Use memory-efficient attention such as FlashAttention.

• Use mixed precision to cut memory and time.

• Cache the labelled-row representations if the same context is reused for many queries.

• Limit the number of columns per attention call by grouping.

Step 10: Test small, then scale

Build a small version. Check that it can learn simple tasks, then scale width, depth and data together. Track the learning curve to see where returns diminish.

Key concepts and terms

A worked example

This example builds a small tabular transformer in PyTorch with alternating column and row attention, a mask token for missing values, and a head for classification. It is a teaching sketch, not a production model.

import torch
import torch.nn as nn

class CellEncoder(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.num = nn.Linear(1, d)
        self.missing = nn.Parameter(torch.randn(d))
        self.col_embed = nn.Embedding(512, d)   # column identity, no names needed

    def forward(self, x, col_ids):
        # x: (rows, cols) with NaN for missing
        mask = torch.isnan(x)
        v = self.num(torch.nan_to_num(x).unsqueeze(-1))
        v = torch.where(mask.unsqueeze(-1), self.missing.expand_as(v), v)
        return v + self.col_embed(col_ids)       # (rows, cols, d)

class AxialBlock(nn.Module):
    def __init__(self, d, heads=4):
        super().__init__()
        self.col_attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.row_attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.ff = nn.Sequential(nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d))
        self.n1, self.n2, self.n3 = nn.LayerNorm(d), nn.LayerNorm(d), nn.LayerNorm(d)

    def forward(self, h):
        R, C, d = h.shape
        # attention across columns within each row
        a, _ = self.col_attn(h, h, h)
        h = self.n1(h + a)
        # attention across rows within each column
        ht = h.transpose(0, 1)                    # (C, R, d)
        b, _ = self.row_attn(ht, ht, ht)
        h = self.n2(h + b.transpose(0, 1))
        return self.n3(h + self.ff(h))

class TinyTabModel(nn.Module):
    def __init__(self, d=64, layers=3, n_classes=2):
        super().__init__()
        self.enc = CellEncoder(d)
        self.blocks = nn.ModuleList([AxialBlock(d) for _ in range(layers)])
        self.head = nn.Linear(d, n_classes)

    def forward(self, x):
        R, C = x.shape
        col_ids = torch.arange(C).expand(R, C)
        h = self.enc(x, col_ids)
        for blk in self.blocks:
            h = blk(h)
        row_repr = h.mean(dim=1)                  # pool across columns
        return self.head(row_repr)

model = TinyTabModel()
x = torch.randn(50, 8)
x[torch.rand_like(x) < 0.1] = float("nan")
print(model(x).shape)   # (50, 2)

Walking through it

• The encoder gives each cell a vector. Missing cells receive a learned mask vector, so they carry information.

• A column embedding gives each column an identity without any name. This is the hook for working on raw exports.

• Each block first mixes information across columns in a row, then across rows in a column, then applies a feed-forward layer. Residual connections and normalisation keep training stable.

• The head reads a pooled row vector and outputs class scores.

What this sketch leaves out

• It does not separate labelled rows from query rows, so it cannot do in-context learning as written. A real design adds label inputs for labelled rows and an attention mask that lets queries see them but not the reverse.

• It uses mean pooling and does not handle categories. A real design embeds categories and may use a special token per row.

• Its column embedding is tied to a column index. A real design might randomise indices during training to avoid memorising order.

A numeric look at the parameter budget

For hidden size d equal to 64, one block holds roughly these weights.

• Two attention modules, each with about four times d squared weights: 2 times 4 times 4,096, about 32,800.

• One feed-forward layer with two matrices of d by 4d: 2 times 64 times 256, about 32,800.

• Total per block, about 65,600 before biases and norms.

Three blocks give around 200,000 weights. Doubling d to 128 multiplies each block by four, to roughly 262,000. This shows why width is expensive: parameters scale with the square of d, while they scale only linearly with the number of blocks.

A cost look at attention

For a table with R rows and C columns, full attention over all cells costs about (R times C) squared operations. Axial attention costs about R times C squared plus C times R squared. With 1,000 rows and 20 columns, full attention is about 400 million pair scores, while axial attention is about 400 thousand plus 20 million, near 20 million. That is a twenty-fold saving, and it grows with table size.

Strengths and limitations

Strengths of careful architecture design

• Efficiency. A design that matches the data's structure learns faster and needs fewer parameters.

• Generalisation. Correct symmetries such as permutation invariance remove whole classes of errors.

• Flexibility. Variable-size inputs and optional names make a model usable on many tables.

• Reuse. One backbone with several heads serves many tasks.

• Predictable scaling. Clear structure lets teams plan compute and data.

Limitations and risks

• Design is hard to verify in advance. Many ideas only show their value after costly experiments.

• Interactions with data. A design that works on one prior may fail on another.

• Cost of attention. Even with axial patterns, very large tables strain memory.

• Fixed capacity. Too small a budget underfits. Too large a budget overfits small tasks or wastes compute.

• Brittleness at the edges. Extreme column counts, huge cardinalities or very long text cells can break assumptions.

• Complexity. More components mean more places for bugs and silent failures.

How it compares with alternatives

Choosing a path

• If you have a single large table and a fixed task, a tree ensemble is often the best first step.

• If you have many small tables or need to transfer knowledge, a pre-trained neural design with in-context learning is attractive.

• If the table is very wide or long, use designs with latent bottlenecks or grouped attention.

• If speed at inference matters, favour smaller depth, cached contexts and efficient kernels.

Role in tabular foundation models

Architecture design decides whether a model can learn from raw tables at all. Every special property of tables becomes a concrete design requirement.

Rows and columns both carry meaning

Rows are records. Columns are features. A single flat sequence loses this structure. Axial or alternating attention restores it, with column attention capturing feature interactions and row attention enabling comparison across records. FT-Transformer, SAINT and TabICL each take a different route to the same goal.

Mixed types

Input encoders must treat numbers, categories, dates and text in a single shared vector space. The choice of encoding determines how well the model handles skew, rare categories and ordinal scales. Text-aware designs such as TabSTAR allow semantic information from names or text cells to enter the model.

Missing values

Missingness is a signal. A learned mask embedding and masked-cell training let the model reason about gaps. This connects to masked cell modeling and imputation networks, where the output head predicts the missing value with an appropriate type.

No column names

If names are absent, the model must identify column roles from data. The design uses learned or random column identifiers, permutation-equivariant layers and data-driven statistics, so that nothing depends on header text. Names can still be used when present, through an optional text channel.

Pre-training on synthetic or real tables

The architecture has to be paired with a prior that produces diverse tables. See synthetic data generation. The design must accept the full range of sizes, types and noise that the prior produces, and the training loop must sample tasks efficiently. Curriculum choices such as curriculum scheduling can move training from small, simple tables to larger, harder ones.

In-context learning

The model must separate context rows from query rows with the right attention mask, take labels as input for context rows, and output predictions for all queries in one pass. This is the structural basis of few-shot learning and, at table level, of zero-shot learning. Because context length limits how many labelled rows the model can read, efficient attention and compression of the context are central design topics.

Continual learning

A foundation model will need updates as new data and skills arrive. Design helps here. Modular heads, adapters and expert layers let new skills be added without overwriting old ones. Training-side methods such as experience replay, EWC and LoRA reduce forgetting. A shared backbone with swappable heads is easier to extend than a monolith.

Imputation, anomaly detection and sector classification

One backbone can serve several jobs through different heads. An imputation head fills cells, an anomaly head scores rows, and a semantic or sector head tags columns and tables. See anomaly detection and semantic sector heads. Training these jointly can improve each, since they share what the model knows about table structure.

Scale and evaluation

Parameter count alone does not tell quality. Compare designs on held-out real benchmarks, not just on the synthetic prior. See /model/benchmarks for how such comparisons are framed, and /model/research for the background.

Practical guidance and best practices

Begin with symmetries

• Write down which symmetries the data has. For tables, these are row permutation invariance and column permutation equivariance.

• Choose layers that respect them from the start.

Keep the encoder simple and robust

• Scale numbers robustly and clip extremes.

• Give every type a clear path to a vector.

• Always include a missing indicator.

Use proven building blocks

• Residual connections, layer normalisation and standard feed-forward layers are well understood.

• Add novelty only where the data demands it.

Make attention affordable

• Prefer axial or grouped attention for large tables.

• Use efficient kernels such as FlashAttention.

• Limit the number of columns processed together and chunk when needed.

Plan the parameter budget

• Fix a compute budget first, then pick depth and width inside it.

• Avoid very narrow and very deep or very wide and very shallow extremes.

• Consider mixture of experts when capacity matters more than per-input compute.

Design for several tasks early

• Give the backbone a clean interface so that heads can be added.

• Keep heads small, and share as much as possible.

Run ablations

• Remove one component at a time and measure the loss.

• Keep only the parts that help on held-out data.

• Record all results so that decisions can be revisited.

Search wisely

• Use Bayesian optimization or Optuna for training settings, but treat structural choices as hypotheses to test with ablations.

Test robustness

• Evaluate with shuffled column order, dropped names and added noise columns.

• Test with extreme row counts and column counts.

• Check behaviour on all-missing columns and constant columns.

Think about deployment from day one

• Measure latency at the target table size.

• Check memory use at peak context.

• Plan for quantised or distilled variants if needed.

Common mistakes

• Copying an architecture from images or text without asking what symmetries tables have. A flat sequence over cells wastes capacity and cost.

• Using fixed positional encodings for columns. This makes the model depend on column order.

• Requiring column names. Many real files have none.

• Ignoring missing values at the encoder. Zero-filling confuses a real zero with an absent value.

• Letting query rows influence each other or the context. This leaks information and breaks in-context learning.

• Over-investing in depth before fixing data. Poor data cannot be fixed by deeper layers.

• Judging by parameter count. Bigger is not automatically better.

• Skipping ablations. Without them, you cannot know which parts matter.

• Evaluating only on the training prior. Real tables behave differently.

• Forgetting efficiency. A design that cannot run at realistic table sizes has little use.

• Making it impossible to extend. A rigid design makes continual learning and new heads costly.

Frequently asked questions

What is architecture design in machine learning?

It is the process of choosing the structure of a model before training. It covers the input encoding, the repeated blocks, the attention pattern, the depth and width, the output heads and the total parameter budget. These choices set what the model can learn and how much it costs to train and run.

Why does architecture matter so much for foundation models?

A foundation model is trained once at great cost and then reused for many tasks. A weak design limits every later use, and fixing it means retraining. A strong design lets the same compute produce a better and more flexible model, so the early choices have large downstream effects.

How is a tabular architecture different from a text architecture?

Tables have two axes, rows and columns, with mixed types and no natural order. Text is a single ordered sequence of one type. A tabular design therefore needs separate attention along rows and columns, permutation-aware layers, type-specific input encoders and a way to treat missing cells as information.

What is axial attention and why is it used for tables?

Axial attention applies attention along one axis at a time, such as across columns within each row, then across rows within each column. It lets information move across the whole table over several layers while costing far less than attention over all cells at once. This makes larger tables practical.

How can a model work without column names?

The design gives each column an identity through learned or randomised column embeddings and uses layers that do not depend on column order. The model then infers what each column means from its values and its relation to the target. If names exist, an optional text channel can add extra signal.

How are missing values handled in the architecture?

The encoder adds a mask flag and a learned embedding for the unknown state, so a missing cell is a distinct input rather than a zero. Training with random masks teaches the model to use or fill gaps. An imputation head can then predict the missing cell in the right type.

How does the architecture support in-context learning?

Labelled rows are given their labels as input, while query rows get a blank. An attention mask lets query rows read the labelled rows but not the other way, and the model outputs predictions for all queries in one forward pass. Pre-training on many tasks teaches the model to use this setup.

How do I choose the parameter budget?

Start from the compute and time you can spend on training and the speed you need at inference. Then pick depth and width so that the model neither underfits nor wastes compute, guided by small experiments and scaling trends. Sparse designs such as mixture of experts can add capacity without proportional cost.

Should one model handle many tasks or should I build separate models?

A shared backbone with separate heads is a good default. It shares knowledge of table structure across prediction, imputation, anomaly scoring and tagging. Separate models can still be better when tasks have very different needs or when strict isolation is required.

How do I know my design is good?

Test it on held-out real data, with ablations to show which parts matter, and with robustness checks for shuffled columns, missing values and different table sizes. A design that holds up across these tests, and runs at realistic sizes, is more trustworthy than one that only scores well on its training prior.

Glossary

• Ablation: removing a component to measure its effect.

• Adapter: a small trainable module added to a frozen network.

• Attention head: one of several parallel attention computations.

• Axial attention: attention applied along one axis at a time.

• Backbone: the main repeated part of a network.

• Bottleneck: a narrow stage that compresses information.

• Context length: how many items a model can process together.

• Equivariance: output order follows input order.

• Feed-forward layer: a per-position network inside a transformer block.

• Head: a small output network for a specific task.

• Inductive bias: built-in assumptions that guide learning.

• Invariance: output unchanged by reordering the input.

• Mask token: a learned vector standing for a hidden or missing value.

• Parameter budget: the total count of trainable weights.

• Residual connection: a shortcut that adds input to output.

• Sparse activation: using only part of the network per input.

Summary

Architecture design chooses the structure of a model: how inputs are encoded, which attention pattern links them, how deep and wide the network is, which heads produce outputs, and how the parameter budget is spent. It is the first and most consequential step in building a foundation model, because it sets the limits for everything that follows. Good designs match the symmetries of the data, stay efficient, and tie the structure to the training objective.

For tabular foundation models, the design must respect rows and columns, mixed types, missing values and missing column names. Axial attention, learned mask and column embeddings, permutation-aware layers, and attention masks for context and query rows make pre-training on synthetic or real tables and in-context learning possible. Shared backbones with separate heads support imputation, anomaly detection, sector and semantic classification, and later extension by continual learning. Related reading includes transformers, axial attention, TabPFN, /model/architecture, /model/research and /model/benchmarks.

Related

It belongs to Foundation models from scratch in my stack, under Foundation models and applied ML. Others in the same group: