Özer Cem KELAHMETKELAHMET
Tools · Tabular foundation models

Axial attention

Axial attention applies self-attention along one axis at a time. Learn how it works, why it suits tables, and how it compares with full attention.

Overview

Axial attention is a way to run self-attention over two-dimensional data without paying the full quadratic cost. Instead of letting every position look at every other position, the model attends along one axis at a time. First each row looks at itself, then each column looks at itself. After two such passes, information from any cell can reach any other cell. The price drops from the square of the total number of cells to something close to the number of cells times the sum of the two side lengths.

This matters a great deal for tables. A table is a grid with rows (records) and columns (features). A model that wants to reason about both directions needs a cheap and honest way to do it. Axial attention gives that: attention across columns inside a row captures how features relate, and attention across rows inside a column captures how one feature behaves over many records. It is one of the core building blocks behind modern tabular foundation models, and it pairs naturally with masked cell modeling for pre-training.

What it is

Axial attention is an attention pattern, not a new kind of network. It reuses the same query, key and value machinery found in transformers. The only change is which positions are allowed to attend to which.

Think of a spreadsheet. Standard self-attention would flatten the sheet into one long list of cells and let every cell compare itself with every other cell. With 1,000 rows and 50 columns that is 50,000 cells, and 50,000 times 50,000 is 2.5 billion pairs. Axial attention instead says: compare each cell only with the other cells in its own row, then only with the other cells in its own column. That is 50,000 times (50 + 1,000), about 52 million pairs. The saving is large, and it grows as the table grows.

The two axes

  • Row axis (feature axis): attention among the cells of one record. It answers questions such as "given this person's age and income, what does the occupation column mean here?"
  • Column axis (sample axis): attention among the values of one feature across many records. It answers questions such as "where does this value sit in the distribution of this column?"

What stays the same

  • Queries, keys and values are still linear projections.
  • Softmax weights still sum to one.
  • Multi-head attention still splits the work into several heads.
  • Residual connections, normalization and feed-forward blocks still wrap each attention step.

Where it came from

Attention itself rose to fame with the transformer in 2017. Researchers soon noticed that images and other grids are expensive to handle with full attention, because the cost grows with the square of the number of pixels. Several lines of work looked for sparse patterns that cut the cost while keeping a wide reach.

Axial attention was proposed for images and other multidimensional data around 2019, in a paper often cited as "Axial Attention in Multidimensional Transformers". The idea was to factor attention over a grid into a sequence of one-dimensional attentions along each axis. A related vision model, Axial-DeepLab (2020), used the same factoring for image segmentation, and showed that stacking height-axis and width-axis attention can match the reach of full two-dimensional attention.

The same trick appeared in protein modeling, where multiple sequence alignments form a grid of sequences by positions. There, attention along rows and columns of the alignment became a standard design. Table models borrowed it later, for the same reason: a table is also a grid, with a different meaning for each axis.

Why tables adopted it

Early deep tabular models such as TabNet and the FT-Transformer treated one row at a time. They attended across features but not across rows, or they mixed rows only in a limited way. Later models, including SAINT, added attention across rows. In-context tabular models such as TabPFN go further and use alternating attention over features and over samples so that the model can read an entire training set in one forward pass. Axial attention is the clean name for that alternating pattern.

Axial attention belongs to a larger family of ways to make attention cheaper. Knowing the family helps you pick the right tool.

  • Sparse attention. Only a chosen set of pairs is computed. Axial attention is one fixed sparse pattern.
  • Low-rank attention. The score matrix is approximated by a smaller one.
  • Linear attention. The softmax is replaced by a kernel so the cost grows linearly with length.
  • Latent attention. Inputs are summarized by a small array, as in Perceiver IO.
  • Exact but fused attention. The result is identical to full attention, but memory is saved by smarter kernels, as in FlashAttention.

Axial attention stands out because it is exact within each line and because its pattern lines up with a real property of the data.

Everyday picture of the whole idea

Imagine a large office with desks in a grid. A manager wants every desk to learn what the others know.

  • Full attention: every worker walks to every other desk. Chaos.
  • Axial attention: first each worker talks to the people at the same row of desks, then to the people in the same column. News still crosses the whole office, but with far fewer conversations.

How it works

Here is the step-by-step flow, with an everyday picture for each step.

Step 1: Turn cells into vectors

Each cell becomes an embedding vector. A numeric cell is passed through a small encoder that maps the number to a vector. A categorical cell uses a lookup table. A missing cell gets a special learned vector. Now the table is a three-dimensional array: rows by columns by embedding size.

Analogy: imagine every cell of a spreadsheet printed on a small card. Each card now carries a short list of numbers that describes its content.

Step 2: Attend along the row

Pick one row. Its cards are laid out in a line. Each card asks, "which other cards in my line matter for me?" It scores them, turns the scores into weights, and mixes their information into itself. This is done for every row, in parallel.

Analogy: a family sitting around a dinner table. Everyone listens to everyone else at that one table, but not to other tables.

Step 3: Attend along the column

Now pick one column. Its cards are laid out in a line from top to bottom. Each card looks at the other cards in the same column and mixes in what it finds useful. This is done for every column, in parallel.

Analogy: a school class photo sorted by height. Each child looks at the others in the same line and learns where they stand among peers.

Step 4: Repeat in layers

A single row pass followed by a single column pass lets one cell hear from any other cell in two hops: first along its row to a cell in the same row as the target, then along the column. Real models stack many such pairs of layers, so the signals get refined again and again.

Step 5: Read out the result

At the end, the vector of each cell, or the vector of a special target cell, goes through a small head. The head produces a class score, a regression value, a reconstructed value for a masked cell, or an anomaly score.

The cost, in plain numbers

Let a table have R rows and C columns.

Method Pairs compared per layer Example with R=1000, C=50
Full attention over all cells (R x C) squared 2,500,000,000
Row attention only R x C squared 2,500,000
Column attention only C x R squared 50,000,000
Axial (row then column) R x C squared + C x R squared 52,500,000

Axial attention is not free, but it is about fifty times cheaper than full attention in this example, and the gap widens with size.

A sketch in code

python
import torch
import torch.nn as nn

class AxialBlock(nn.Module):
    def __init__(self, dim, heads=4):
        super().__init__()
        self.row_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.col_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.ff = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(),
                                nn.Linear(4 * dim, dim))
        self.norm3 = nn.LayerNorm(dim)

    def forward(self, x):
        # x has shape (rows, cols, dim)
        R, C, D = x.shape
        # attend across the columns inside each row
        h = self.norm1(x)
        h, _ = self.row_attn(h, h, h)          # batch = rows, length = cols
        x = x + h
        # attend across the rows inside each column
        h = self.norm2(x).transpose(0, 1)      # now (cols, rows, dim)
        h, _ = self.col_attn(h, h, h)          # batch = cols, length = rows
        x = x + h.transpose(0, 1)
        return x + self.ff(self.norm3(x))

x = torch.randn(200, 12, 32)   # 200 rows, 12 columns, 32-dim cells
print(AxialBlock(32)(x).shape) # torch.Size([200, 12, 32])

The key trick is reshaping. For row attention, the row index acts like the batch and the column index acts like the sequence. For column attention, the roles swap.

Inside one attention head

A single head works in three small moves.

  1. Project. Each cell vector is multiplied by three learned matrices to make a query, a key and a value.

  2. Score. The query of one cell is compared with the key of every other cell on the same line. A higher score means "this neighbour is relevant".

  3. Mix. The scores go through a softmax, which turns them into weights that add up to one. The values are then added together using those weights.

Analogy: a cell is a person at a meeting. The query is the question they hold in mind. Each neighbour wears a name tag, which is the key. The value is what the neighbour actually says. The person listens most to the neighbours whose tag matches the question.

Variants of the idea

Axial attention comes in a few flavours, and table models pick among them.

  • Sequential axial layers. Row attention and column attention are applied one after the other in the same block. This is the most common choice.
  • Parallel axial layers. Both attentions are computed from the same input and their outputs are summed. This is a little faster on parallel hardware but mixes the two signals less.
  • Interleaved with feed-forward blocks. A feed-forward block sits between the row and column steps, so each axis sees a transformed version of the other.
  • Gated axial attention. A learned gate controls how much each axis contributes, which can help when one axis carries more signal.
  • Masked axial attention. Some pairs are blocked, for example to keep test labels hidden from other test rows.
  • Axial plus latent. A small latent array is added so that cells can also exchange information through a shared summary. See Perceiver IO.

What a mask looks like for a table

In a table used for in-context learning, there are two kinds of rows: support rows with labels, and query rows without labels. A common rule is:

  • Support rows may attend to other support rows along the column axis.
  • Query rows may attend to support rows, but not to other query rows.
  • Nobody sees the label of a query row.

This rule keeps the prediction for one query row independent of the others, so the result does not depend on which other test rows happen to be in the batch.

Why two hops are enough for reach

Picture the grid as a city with streets and avenues. Row attention is travel along a street. Column attention is travel along an avenue. Any two blocks in the city can be linked by going along one street and then one avenue. That is why one row layer and one column layer are enough for any cell to hear from any other cell, at least indirectly.

Key concepts and terms

Term Plain meaning
Axis One direction of the grid: rows or columns
Self-attention Each position scores and mixes information from other positions
Row attention Attention among the cells of one record
Column attention Attention among the values of one feature across records
Receptive field The set of positions that can influence a given cell
Quadratic cost Cost that grows with the square of the number of items
Factorized attention Splitting one big attention into several smaller ones
Cell embedding A vector that represents the content of one table cell
Alternating attention Stacking row and column layers one after another
Permutation symmetry Output should not change if rows (or columns) are reordered
Mask A rule that blocks some attention pairs
Head One of several parallel attention computations in a layer

A worked example

Take a tiny table with three rows and three columns:

age income owns_car
25 40 0
40 90 1
NaN 85 1

We want to fill the missing age in row three.

  1. Embed. Each cell gets a vector. The missing age gets a special "missing" vector, plus a learned hint for the age column.
  2. Row attention on row three. The missing age cell looks at income (85) and owns_car (1). In training, the model learned that high income and owning a car often go with mid-life ages. The missing cell absorbs that clue.
  3. Column attention on the age column. The missing cell now looks at the other ages, 25 and 40. It learns that this column holds values in a range around the thirties and forties.
  4. Second round. A second row layer lets the age cell combine its improved column knowledge with the row clues. A second column layer refines it again.
  5. Decode. A small head maps the final age vector to a number, say 41.

The key point is that neither attention direction alone could do the job well. Row attention knows the relationships between features. Column attention knows the scale and shape of the feature. Axial attention combines both without comparing every cell with every other cell.

A numeric walk-through of one row step

To see the arithmetic, shrink the vectors to two numbers and look at the row for the missing age.

  • Cell "age" (missing) has the query vector (1, 0).
  • Cell "income" has the key vector (2, 0) and the value vector (0.8, 0.1).
  • Cell "owns_car" has the key vector (0, 2) and the value vector (0.2, 0.9).
  • Cell "age" itself has the key vector (1, 1) and the value vector (0.0, 0.0), since it is missing.

The dot products of the query with the keys are 2, 0 and 1. After dividing by the square root of the vector size (about 1.41) the scores are about 1.41, 0 and 0.71. The softmax of these scores gives weights of about 0.55, 0.14 and 0.31 for income, owns_car and age itself.

The mixed value is then 0.55 times (0.8, 0.1) plus 0.14 times (0.2, 0.9) plus 0.31 times (0, 0). That is about (0.47, 0.18). The age cell has absorbed most of its new information from the income cell, which is the pattern the model learned is useful. A real model repeats this with many more dimensions and many heads.

Why the column step adds new facts

The column step uses the same arithmetic but on a different line. In the age column the missing cell compares itself with the ages 25 and 40. Those values are not available in its own row. They tell the cell about the scale of the feature, for example that ages in this table run from the twenties to the forties. Without the column step, the cell could only guess from income and car ownership.

Counting reach

After one row layer and one column layer, the age cell in row three has been influenced by:

  • all cells in row three (through the row layer), and
  • all cells in the age column (through the column layer),
  • and, because other cells in those lines were themselves updated by their own lines, indirectly by every cell in the table after two full pairs of layers.

This is the sense in which axial attention keeps a global reach with local work.

Training signals that pair well with axial layers

Axial attention is only the engine. What it learns depends on the training signal.

  • Masked reconstruction. Hide cells and predict them. See masked cell modeling.
  • Label prediction in context. Predict hidden labels from labeled rows in the same table.
  • Contrastive objectives. Pull together two views of the same row and push apart different rows. See contrastive learning and InfoNCE.
  • Denoising. Corrupt cells with noise and learn to restore them. See denoising autoencoders.
  • Column type prediction. Predict the semantic type of each column from its values.

A strong foundation model often mixes several of these.

Reading attention maps

Row attention maps are small, one value per pair of columns, so they can be shown as a heat map. They often reveal which features the model links together. Column attention maps are large, one value per pair of rows, so they are usually summarized. A common summary is the share of attention that each query row gives to support rows with the same label. High values suggest the model is acting like a learned nearest neighbour rule.

Strengths and limitations

Strengths

  • Much cheaper than full attention. The cost grows with the sum of side lengths, not the product.
  • Matches the structure of tables. Rows and columns have clear and different meanings, and the two passes respect that.
  • Works with variable shapes. Attention has no fixed width, so tables with different numbers of rows and columns fit the same network.
  • Supports in-context learning. Column attention lets test rows read the training rows directly.
  • Good fit for masking. Hidden cells can be recovered from both their row neighbours and their column neighbours.
  • Easy to implement. It is built from standard attention layers plus a reshape.

Limitations

  • Still quadratic along each axis. A table with a million rows still makes column attention expensive. Long tables need sampling, chunking, or a latent bottleneck such as Perceiver IO.
  • Indirect interactions. A cell reaches distant cells only after several layers. Very precise cross-cell patterns may need depth.
  • Order handling needs care. Rows and columns of a table usually have no natural order. The model must avoid leaning on position, or use permutation-friendly design.
  • Memory pressure. Attention matrices for long columns can exceed device memory without fused kernels such as FlashAttention.
  • Two axes only by default. Data with more structure, such as multiple linked tables, needs extra mechanisms.

How it compares with alternatives

Approach Reach Cost Fits tables? Notes
Full attention over flattened cells Every cell sees every cell Very high Poor beyond small tables Loses explicit row and column roles
Row-only attention (feature mixing) Within a record Low Partial Used by FT-Transformer; no sample-to-sample learning
Axial attention Row then column Moderate Strong Natural for grids
Latent bottleneck (Perceiver IO) Through a small latent array Linear in input size Strong for huge inputs Compresses information
Windowed or local attention Nearby cells Low Weak Tables have no useful notion of "nearby"
Trees (XGBoost, LightGBM) Learned splits Low Strong baseline No pre-training, no in-context use
Row-to-row attention (SAINT intersample) Between records Moderate Good Only one of the two axes

When to prefer which

  • If the table is small to medium and you want one network that handles both feature interactions and sample context, use axial attention.
  • If the table is very large, add a latent bottleneck or subsample rows.
  • If you only need a quick strong baseline on one dataset, gradient boosting is hard to beat.

Role in tabular foundation models

A tabular foundation model is a network pre-trained on a very large number of tables so that it can handle a new table without training from scratch. Axial attention is one of the best matches between a network design and the nature of this problem.

Rows and columns have different jobs

In a table, a row is one entity, such as a customer, a machine reading or a transaction. A column is one property measured for all entities. A foundation model has to learn both kinds of structure: how properties relate inside a record, and how a property is distributed across records. Axial attention gives each of these its own layer, so the model does not have to untangle them from a flat sequence.

Mixed types and missing values

Cells can be numbers, categories, dates or text, and some are missing. Each cell type is embedded into the same vector space, and missing cells receive a dedicated vector. After embedding, axial attention treats every cell alike. This is why one architecture can serve many table shapes.

No column names needed

Many real tables have meaningless or absent headers. Because axial attention works from the content of the cells and their pattern across rows and columns, a model can still understand a column by what its values look like and how they move with other columns. Column names, when present, become extra input but are not required.

Pre-training on synthetic and real tables

Foundation models for tables are often pre-trained on millions of generated tables that cover many causal structures, plus real tables. Axial attention is size-agnostic, so the same network sees tables of varied width and length during pre-training. Used with masked cell modeling, it learns to recover hidden cells from the rest of the grid. See /model/research for the research direction and /model/architecture for how such blocks are combined.

In-context learning

In in-context learning, the training rows and their labels are placed in the same table as the query rows. The query rows have an empty label cell. Column attention lets each query row look at the labeled rows and copy the pattern, much like a student scanning solved examples before answering. No weight update is needed at prediction time. This is the mechanism behind TabPFN and TabICL.

Continual learning

When a model must absorb new tables over time, the attention layers themselves can stay general while small parts are updated. Methods such as elastic weight consolidation and experience replay protect the knowledge in the attention weights. Because axial layers are shared across all table shapes, they are a natural place to keep stable knowledge.

Imputation

Imputation is filling missing cells. Row attention supplies clues from the same record; column attention supplies the distribution of the feature. Together they give a better estimate than either one alone, and the head can output a spread as well as a point value. See imputation networks.

Anomaly detection

An anomaly is a row or cell that does not fit its row context or its column context. A model trained to reconstruct cells will have a large error on cells that are unusual in either direction. Axial attention exposes both views, which helps separate "odd for this record" from "odd for this column". See anomaly detection and isolation forest for classic baselines.

Sector and semantic classification

When a model must decide what a table is about, such as which industry or which semantic type each column holds, it needs to read the pattern of values across the whole grid. Column attention pools evidence over many rows, and row attention links columns together. A classification head reading the pooled result can predict a sector or a column type. See semantic and sector heads. The results of this kind of work are tracked at /model/benchmarks.

Putting the pieces together in a table model

A typical tabular foundation model that uses axial attention follows this recipe.

  1. Encode each numeric column with a shared numeric encoder and each categorical column with an embedding.
  2. Add a column identity vector, so the model can tell columns apart without relying on order.
  3. Add a learned missing vector wherever a cell is empty.
  4. Run a stack of axial blocks, each with a row layer, a column layer and a feed-forward layer.
  5. Attach one or more heads for the task: classification, regression, masked cell reconstruction, anomaly score or column type.

The same stack handles every task. Only the head and the masking pattern change.

Scaling behaviour

Because the layers are shared across any table shape, the number of parameters does not depend on the table size. Memory and time do. Doubling the number of rows roughly quadruples the column attention cost, while doubling the number of columns roughly quadruples the row attention cost. In practice, columns are few and rows are many, so the column axis is the one to watch. Techniques that help include subsampling rows, grouping rows into chunks, using fused kernels and, when needed, replacing full column attention by a latent bottleneck.

Interaction with positional information

Language has a natural order and so positional encodings help. Tables usually do not. Shuffling rows or columns should not change what the table means. Axial attention is friendly to this because it uses no position by default. If a table does have order, such as time stamps in a series, the order should be given as a feature rather than as hidden position. This keeps the model honest about what it knows.

Sparse and dense hardware effects

Attention along a short axis such as the columns is cheap and fits easily on a device. Attention along a long axis is the heavy part. Fused kernels compute the softmax in blocks without storing the full score matrix, which can make the long axis many times lighter on memory. Mixed precision also helps. See mixed precision and gradient checkpointing for related training tricks.

Practical guidance and best practices

  • Alternate the axes. Use row then column in every block, rather than many row layers followed by many column layers. Mixing early helps both signals.
  • Normalize before attention. Pre-norm layers train more stably in deep stacks.
  • Use fused attention kernels. Memory-efficient kernels make long columns practical. See FlashAttention.
  • Cap the sample axis. For very long tables, sample a subset of rows per forward pass, or use chunking.
  • Treat padding with masks. When batching tables of different sizes, mask the padded rows and columns so they do not leak into attention.
  • Respect permutation symmetry. Do not add positional encodings along row order unless order carries meaning, as in time series. For columns, use identity embeddings only if you want the model to know which column is which.
  • Standardize numeric inputs. Scale each column robustly before embedding, so a column with huge values does not dominate.
  • Mark missing cells explicitly. A learned missing vector is better than a fake zero.
  • Keep heads small and several. Four to eight heads with moderate size per head is a stable default.
  • Test on shifted shapes. After training, check that the model still works on tables wider or longer than it saw in training.
  • Compare against boosting. Always include gradient boosting as a sanity baseline.

A short checklist before training

  1. Do all cell types map to the same embedding size?
  2. Is the missing-value path tested?
  3. Are padded cells masked in both axes?
  4. Does the output stay the same when rows are shuffled?
  5. Is memory use measured on the largest expected table?

Debugging tips

  • Feed a table with all identical rows and check that column attention gives uniform weights.
  • Shuffle the rows and confirm that predictions for a given row do not change.
  • Shuffle the columns and confirm the same, if column identity is meant to be learned from content.
  • Remove one axis and measure the drop in accuracy to see which axis matters for your data.
  • Plot the attention weights of one head for one row to see whether it focuses on a few informative columns or spreads evenly.

Choosing sizes

Setting Small model Medium model Large model
Embedding size 32 to 64 128 to 256 512 and above
Heads 4 8 8 to 16
Axial blocks 4 to 6 8 to 12 12 to 24
Typical use Single dataset Multi-dataset Foundation model

These are rough guides rather than rules. Start small, check that the loss falls, and grow only when validation data shows underfitting.

Common mistakes

  • Flattening the table and using full attention. It looks simple but blows up in memory and hides the row and column roles.
  • Using only one axis. Row-only models miss sample context; column-only models miss feature interactions.
  • Adding row-position encodings to unordered data. The model then learns the order of the file, which is noise, and results change when rows are shuffled.
  • Ignoring padding masks. Padded cells then take part in attention and distort the output.
  • Letting one huge column dominate. Without scaling, a large-valued numeric column swamps the embedding.
  • Forgetting that depth is needed. One pair of layers has limited reach; stack several.
  • Evaluating on one table shape. A model that works for 10 columns may fail at 200 if never trained across widths.
  • Leaking labels. In in-context setups, query labels must be masked, or the model simply reads the answer.
  • Treating attention weights as explanations. Weights show where mass went, not why a prediction was made. Use tools such as SHAP for attribution.

Rules of thumb

  • If your table has fewer than a few hundred columns and a few thousand rows, plain axial attention is usually fine on one modern GPU.
  • If it has many thousands of rows, plan for subsampling or chunked column attention.
  • If it has thousands of columns, add column grouping or a latent bottleneck for the row axis.
  • If you need very fast inference on small tables, a gradient-boosted model may still win on speed.
  • If you need one model for many unrelated tables, axial attention with pre-training is a strong choice.

What to log during training

  • Loss for each task head separately.
  • Gradient norm per layer, to catch exploding values in deep stacks.
  • Attention entropy per head, to spot heads that collapse to a single position.
  • Peak memory against table size.
  • Accuracy on tables of unseen shapes.

Frequently asked questions

What is axial attention in simple terms?

Axial attention is self-attention applied to one direction of a grid at a time. A model first lets every cell talk to the other cells in its row, then lets every cell talk to the other cells in its column. By alternating the two, information spreads across the whole grid while the work stays far smaller than comparing every cell with every other cell.

Why is axial attention a good match for tables?

A table is a grid whose two axes mean different things. Rows are records and columns are features. Axial attention gives each axis its own layer, so the network can learn feature relationships inside a record and the distribution of a feature across records without mixing the two roles. The shape of the computation mirrors the shape of the data.

How much cheaper is axial attention than full attention?

For a table with R rows and C columns, full attention over all cells costs about the square of R times C. Axial attention costs about R times C times the sum of R and C. For a table of a thousand rows and fifty columns that is roughly fifty times fewer pairs per layer, and the gain grows as the table gets larger.

Does axial attention lose information compared with full attention?

A single axial layer sees only one line of the grid, so it has a smaller view than full attention. After a row layer and a column layer, every cell can influence every other cell through an intermediate cell, so the reach is complete in two steps. Fine cross-cell patterns may need more depth, but in practice stacked axial blocks work very well.

How does axial attention support in-context learning on tabular data?

Training rows with known labels and query rows with hidden labels are placed in the same table. Column attention lets each query row read the labeled rows and copy their pattern into its own representation. Because this happens inside the forward pass, the model can make predictions for a new dataset without any weight update.

Can axial attention handle missing values?

Yes. Missing cells are given a special learned embedding, and the layers treat them like any other cell. Row attention supplies clues from the same record and column attention supplies the typical values of the feature, so the network can estimate the missing cell, or simply carry the missingness signal forward for other predictions.

Is axial attention the same as the attention used in the FT-Transformer?

Not exactly. The FT-Transformer attends across the features of one row, which is the row axis only. Axial attention adds the second axis, so rows can also exchange information. Models that mix both axes, such as those that use inter-sample attention, are closer to the axial idea.

How long a table can axial attention handle?

The cost along the sample axis is still quadratic in the number of rows, so very long tables become heavy. Common remedies are to sample a subset of rows, to chunk the table, to use fused kernels that save memory, or to add a small latent array such as in Perceiver IO that compresses the rows before reasoning.

Do I need column names to use axial attention?

No. The mechanism works from cell values and their patterns, not from header text. Column names can be added as extra embeddings when they exist and help, but a model built around axial attention can still understand a column from how its values behave and how they relate to other columns.

Where else is axial attention used besides tables?

It first appeared for images and other multidimensional arrays, where height and width are the two axes. It is also used in video, in protein alignment models where sequences and positions form a grid, and in other settings where data is naturally arranged along two or more axes.

Glossary

  • Attention: a mechanism that lets each item weigh and combine information from other items.
  • Axis: one dimension of an array, such as rows or columns.
  • Axial attention: attention applied along one axis at a time.
  • Batch: a group of examples processed together.
  • Cell embedding: a vector standing for the content of a table cell.
  • Column attention: attention among the cells in one column.
  • Factorization: splitting a big operation into several smaller ones.
  • Head: an independent attention computation inside a multi-head layer.
  • In-context learning: adapting to a task from examples placed in the input, without weight updates.
  • Latent array: a small set of learned vectors that summarizes a large input.
  • Mask: a rule that hides some positions from attention.
  • Permutation invariance: the property that output does not depend on the order of items.
  • Quadratic cost: cost that grows with the square of the input size.
  • Receptive field: the set of positions that can influence an output.
  • Row attention: attention among the cells in one row.
  • Self-attention: attention in which queries, keys and values come from the same input.
  • Support row: a row with a known label used as an example in context.
  • Query row: a row whose label is hidden and must be predicted.
  • Softmax: a function that turns scores into weights that add up to one.
  • Residual connection: a shortcut that adds a layer input to its output.
  • Layer normalization: a step that rescales values to keep training stable.
  • Feed-forward block: a small network applied to each vector after attention.
  • Padding: filler cells added so tables of different sizes fit in one batch.
  • Sparse attention: attention that computes only a chosen subset of pairs.
  • Kernel (fused): a low-level routine that does several steps at once to save memory.
  • Missing vector: a learned embedding that marks an empty cell.
  • Column identity vector: a learned vector that tells the model which column a cell belongs to.
  • Entropy (attention): a measure of how spread out the attention weights are.
  • Tabular foundation model: a model pre-trained on many tables that can serve new tables with little or no training.

Summary

In one paragraph

Axial attention applies self-attention along one axis of a grid at a time, first across the cells of a row and then across the cells of a column. This keeps global reach while cutting cost from the square of the total number of cells to roughly cells times the sum of the side lengths.

For tables it is a natural fit, because rows and columns carry different meanings and each gets its own layer. It supports mixed types, missing cells, tables without headers, in-context learning, imputation, anomaly detection and semantic classification inside one architecture. Its main limit is that each axis is still quadratic, so long tables need sampling or a latent bottleneck such as Perceiver IO. Combined with masked cell modeling, it forms a strong core for modern tabular foundation models. See also /model/architecture, /model/research and /model/benchmarks.

Key takeaways

  • Attend along rows, then along columns, instead of over all cells at once.
  • The cost drops from the square of the cell count to cells times the sum of the side lengths.
  • Rows carry feature relations, columns carry feature distributions, and both are needed.
  • Missing cells, mixed types and absent headers fit naturally.
  • Long tables still need sampling, chunking or a latent bottleneck.
  • Pair it with masked pre-training for imputation, anomaly scoring and classification.

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

Type to search.
    ↑↓ navigate↵ open