# TabMT: Generating Tabular data with Masked Transformers

Manbir S. Gulati
AI Accelerator
Leidos Inc
<EMAIL_ADDRESS>
&Paul F. Roysdon
AI Accelerator
Leidos Inc
<EMAIL_ADDRESS>

###### Abstract

Autoregressive and Masked Transformers are incredibly effective as generative
models and classifiers. While these models are most prevalent in NLP, they
also exhibit strong performance in other domains, such as vision. This work
contributes to the exploration of transformer-based models in synthetic data
generation for diverse application domains. In this paper, we present TabMT, a
novel Masked Transformer design for generating synthetic tabular data. TabMT
effectively addresses the unique challenges posed by heterogeneous data fields
and is natively able to handle missing data. Our design leverages improved
masking techniques to allow for generation and demonstrates state-of-the-art
performance from extremely small to extremely large tabular datasets. We
evaluate TabMT for privacy-focused applications and find that it is able to
generate high quality data with superior privacy tradeoffs.

##  1 Introduction

Figure 1: Diagram of TabMT. $m$ is the mask token, $p_{i}$ is the masking
probability of the $i^{th}$ row

Generative models have attracted significant attention in the field of deep
learning due to their ability to synthesize high-quality data and learn the
underlying structure of complex datasets. Such models have been successfully
applied to various data types including images [25], text[11], and tabular
data [17]. This work concentrates on tabular data, which is prevalent in
numerous fields like healthcare, finance, and social sciences. The
heterogeneous nature of tabular data, characterized by its diverse data types,
distributions, and relationships, presents distinct challenges not present in
other domains.

The development of effective synthetic tabular data generators is crucial for
numerous reasons including: privacy preservation, data augmentation, model
interpretability, and anomaly detection. Prior work in this domain has
produced a myriad of generative models, including Generative Adversarial
Networks (GANs)[37][34], Variational Autoencoders (VAEs)[34], Autoregressive
Transformer[30] [3], and Diffusion models[17]. Although these existing models
strive to address the challenges associated with tabular data generation,
there is still room for exploration and improvement. Specifically, we
demonstrate improvements in robustness, scalability, privacy preservation, and
handling of missing data.

Transformers [31], originally designed for natural language processing (NLP)
tasks, have lead to significant advancements in a variety of applications.
Their powerful capacity for modeling complex dependencies and generalizing
across applications has spurred researchers to extend transformers to other
data types, such as images[14][6] and audio[24].

In this paper, we investigate transformers as synthetic tabular data
generators, further expanding their cross-domain applicability. We
specifically examine Masked Transformers (MT), originally developed to produce
strong text embeddings [8], and successfully generalized across many
domains[12][5][13]. We explore their utility as tabular data generators. We
show Masked Transformers make robust and scalable tabular generators,
achieving state-of-the-art performance across a wide array of datasets.

Our key contributions are as follows:

  1. 1.

We propose TabMT, see Figure 1, a simple but effective Masked Transformer
design for generating tabular data, that is general enough to work across many
tasks and scenarios.

  2. 2.

We provide a comprehensive evaluation of TabMT, demonstrating state-of-the-art
performance when compared to existing generative model families, including
GANs, VAEs, Autoregressive Transformers, and Diffusion models. We also
showcase its scalability from very small to very large tabular datasets.

  3. 3.

We highlight the applicability of our model in privacy-focused applications,
illustrating TabMTs ability to arbitrarily trade-off privacy and quality
through temperature scaling. Furthermore, our model’s masking procedure
enables it to effectively handle missing data, thereby increasing privacy and
model applicability in real-world use cases.

##  2 Related Work

#### Language modeling:

Transformer[31] based Language Modeling[8][11] learns Human language via token
prediction. Given a context of either previous tokens in time, or a random
subset tokens within a window, the model is tasked with predicting the
remaining tokens. Using either an autoregressive model or masked model,
respectively. Both these paradigms have shown success across a wide range of
tasks and domains.[12][6][24][35][5] Masked Language models are traditionally
used for their embeddings, although some papers have explored masked
generation of content[5]. Our work builds on Masked Training and demonstrates
its effectiveness for modeling and generating tabular data.

#### Deep Tabular Generators:

Deep learning models are increasingly utilized for generating synthetic
tabular data. Synthetic data is particularly important for tabular data as it
is often subject to privacy requirements. Additionally, tabular datasets are
often hard to acquire, and usually smaller than datasets in other domains.
These conditions increase the importance and challenges associated with
generating tabular data. Deep Tabular Generators have been constructed using
essentially all Deep Generative model families Autoencoders[34],
GANs[34][37][36][10], and AR transformers [30] [3] [23]. Our work uses Masked
Transformers and demonstrates superior performance when compared to prior deep
tabular generators.

#### Netflow Generators:

Netflow data is a specific type of tabular data that captures network
communication events and is commonly used for network traffic analysis,
intrusion detection, and cybersecurity applications. Netflow datasets are
typically extremely large, with complex rules between fields, and a high
number of possible values per field. Generating realistic synthetic netflow
data is crucial for developing and testing network monitoring tools and
security algorithms. There have been several models developed specifically for
netflow generation [28] [33] [21]. We demonstrate that our general tabular
data generator handily outperforms domain specific models.

##  3 Method

TabMTs structure is particularly well suited for generating tabular data, for
a number of reasons.

  1. 1.

TabMT accounts for patterns bidirectionally between fields. Tabular data lacks
ordering, meaning bidirectional learning likely produces better understanding
and embeddings within the model.

  2. 2.

A “prompt” to a tabular generator is not likely sequential. TabMT’s masking
procedure allows for arbitrary prompts to the model during generation. This is
unique as most other generators have very limited conditioning capabilities.

  3. 3.

Missing data is far more common in tabular data than in other domains. TabMT
is able to learn missing values by setting their masking probability to 1.
Other generators require that we impute data separately before we can generate
high-quality cleaned samples.

These structural advantages come from TabMT’s novel masked generation
procedure.

Figure 2: A row of data being sampled from TabMT. Fields are sampled in a
random order and field values are sampled according to the predicted
distributed.

Below we outline how we construct TabMT from the original masked training
procedure outlined in BERT[8] and the justifications behind our design
choices. We also outline the specific changes we make to allow for
heterogeneous data types. A naive adaptation of BERT’s masking procedure would
look as follows. Given an $n$ by $l$ dataset F of categorical and numerical
features, for each row $\textbf{F}_{i}$, the transformer is provided with a
set of unmasked fields $\textbf{F}_{i}^{u}$ and a set of masked fields
$\textbf{F}_{i}^{m}$. Each field in the masked set $\textbf{F}_{i}^{m}$ has
its value replaced with a mask token. The model is then tasked with predicting
the original value for all masked tokens. The row $\textbf{F}_{i}$ is
partitioned into the unmasked and masked sets by conducting a Bernoulli trial
on each field, $\textbf{F}_{i,j}$, such that
$P(\textbf{F}_{i,j}\in\textbf{F}_{i}^{m})=0.15$.

The BERT masking procedure produces a strong embedding model, but not a strong
generator. To create a strong generative model we make two key changes: sample
our masking probability from a uniform distribution and predict masked values
in a random order during generation. To understand why these changes are
effective, we can look at the distribution of masked sets. As a result of the
repeated Bernoulli trials during masking, the size of the masked set for each
row $|\textbf{F}_{i}^{m}|$ will follow a Binomial distribution. However, when
generating data one field $\textbf{F}_{i,j}$ at a time, the model will
inference on masked subset sizes $(0,\dots,l-1)$, once each. We would like the
training distribution of $|\textbf{F}_{i}^{m}|$ to be uniform, matching the
uniform distribution encountered when generating data. With a fixed masking
probability $P(\textbf{F}_{i,j}\in\textbf{F}_{i}^{m})=p_{m}$ we will instead
encounter a Binomial distribution centered around $p_{m}\cdot l$. However, if
we sample our masking probability $p_{m}$ for each row $\textbf{F}_{i}$ such
that $P(p_{m}=p)\sim U(0,1)$, we will train uniformly across subset sizes:

$P(|\textbf{F}_{i}^{m}|=k)=\int_{0}^{1}{\binom{l}{k}}p^{k}(1-p)^{l-k}dp=\frac{l!}{k!(l-k)!}\frac{k!(l-k)!}{(l+1)!}=\frac{1}{l+1}.$
(1)

Fixing this train and inference mismatch is critical to forming a strong
generator. A traditional autoregressive generator would generate fields
$(\textbf{F}_{i,0},\dots,\textbf{F}_{i,l-1})$, sequentially. However, tabular
data, unlike language, does not have an inherent ordering. Generating fields
in a fixed order introduces another mismatch between training and inference.
During training $\textbf{F}_{i}^{m}$ will take on the distribution

$P(\textbf{F}_{i}^{m}=s)=\frac{1}{\binom{l}{|s|}\cdot l}.$ (2)

When generating in a fixed order, the model will infer across $l$ distinct
subsets and no others. However, if we instead infer in a random order, then at
generation step $0\leq t<l$, the distribution of $\textbf{F}_{i}^{m}$ will be
given by

$P(\textbf{F}_{i}^{m}=s)=\frac{t!\cdot(l-t)!}{l!}=\frac{1}{\binom{l}{t}}.$ (3)

Since we encounter each $t$ exactly once, this overall distribution is
identical to the masking distribution encountered during training, fixing the
discrepancy caused by generating fields in a fixed order.

A transformer model will typically have an input embedding matrix
$\textbf{E}\in\mathbb{R}^{k\times d}$, where $k$ is the number of unique input
tokens and $d$ is the transformer width. Because tabular data is
heterogeneous, we instead construct $l$ embedding matrices, one for each
field. Each embedding matrix will have a different number of unique tokens
$k$.

For categorical fields we use a standard embedding matrix initialized with a
normal distribution. For each continuous field we construct an ordered
embedding $\textbf{O}\in\mathbb{R}^{k\times d}$ from its unordered embedding
matrix E and two endpoint vectors $\textbf{l},\textbf{h}\in\mathbb{R}^{d}$.

To construct each ordered embedding matrix O, we first we quantize the values
of the continuous field. Our default quantizer is K-Means. We consider the
maximum number of clusters a hyper-parameter. Let
$\textbf{v}\in\mathbb{R}^{k}$ be the vector of ordered cluster centers. We
construct an vector of ratios $\textbf{r}\in\mathbb{R}^{k}$ using min max
normalization such that

$\textbf{r}_{i}=\frac{\textbf{v}_{i}-\min(\textbf{v})}{\max(\textbf{v})-\min(\textbf{v})}.$
(4)

We use the ratio vector r to construct each ordered embedding in O:

$\textbf{O}_{i}=\textbf{E}_{i}+\textbf{r}_{i}\cdot\textbf{l}+(1-\textbf{r}_{i})\cdot\textbf{h}.$
(5)

This structure allows the transformer to both take advantage of the ordering
of the properties and add unordered embedding information to each cluster. The
unordered embeddings are useful in attention, multi-modal distributions, and
encoding semantic separation between close values. We use this same structure
to construct a dynamic linear layer at the output during prediction. This can
be converted to a traditional linear layer once the model is trained.

Relying too heavily on the unordered embeddings might negate the benefit of
our ordered embedding, as information isn’t effectively shared between close
values. To address this, we bias TabMT to rely on the ordering of embeddings.
For continuous fields, we zero-init the unordered embedding matrix E. Whereas,
the endpoint vectors l and h use a normal distribution with 0.05 standard
deviation. Because entries in matrix O are not independent of each other, to
sharpen the output distribution, the network must either rely on matrix E or
increase magnitude of the endpoint vectors. This can reduce use of the priors
encoded by O or cause instability. To combat this, we include a learned
temperature which can sharpen the predicted distribution using a single
parameter per field instead. Each field’s predicted distribution
$\hat{\textbf{y}}\in\mathbb{R}^{k}$ is given by

$\hat{\textbf{y}}=\frac{e^{\textbf{z}/(\tau_{l}\cdot\tau_{u})}}{\sum_{j}{e^{\textbf{z}/(\tau_{l}\cdot\tau_{u})}}},$
(6)

where $\textbf{z}\in\mathbb{R}^{k}$ is a vector of logits, $\tau_{l}$ is the
learned temperature, and $\tau_{u}$ is the user-defined temperature. See
Figure 1 for an overall diagram of these components. Figure 2 shows the
generation of a single sample. Detailed pseudocode is available in the
Appendix.

##  4 Evaluation

In this section, we present a comprehensive evaluation of TabMT’s
effectiveness across an extensive range of tabular datasets. Our analysis
involves a thorough comparison with state-of-the-art approaches, encompassing
nearly all generative model families. To ensure a robust assessment, we
evaluate across several dimensions and metrics.

#### Datasets:

For our data quality and privacy experiments we use the same list of datasets
and data splits as TabDDPM[17]. These 15 datasets range in size from $\sim
400$ samples to $\sim 150,000$ samples. They contain continuous, categorical,
and integer features. The datasets range from 6 to 50 columns. For our scaling
experiments we use the CIDDS-001[27] dataset, which consists of Netflow
traffic from a simulated small business network. A Netflow consists of 12
attributes, which we post-process into 16 attributes; see the Appendix for
more details. This dataset is extremely large with over 30 million rows and
field cardinalities in the tens of thousands. Other datasets listed all have
cardinalities below 50. Unlike our other benchmarks, we purposely do not
quantize the continuous variables here to further test the scaling of our
model. In other words, every unique value is treated as a separate category in
our prediction process.

#### Prior Methods:

We select four techniques to compare against, one from each each major family
of deep generative models.

  * •

TVAE[34] is one of the first deep tabular generation papers introducing two
models, a GAN and a VAE. We compare against their VAE because it is the
strongest VAE tabular generator we are aware of.

  * •

CTABGAN+[37], at the time of writing, is the state-of-the-art for GAN-based
tabular synthesis.

  * •

TabDDPM[17] adapts diffusion models to tabular data, with the strongest
results of all prior work.

  * •

RealTabFormer[30] is a recent work on adapting autoregressive transformers to
tabular and relational data. This method is most similar to our technique,
however, they use an autoregressive transformer which demonstrates worse
results than our masked transformer.

###  4.1 Data Quality

We use the CatBoost variant of ML Efficiency (MLE)[36][17] for evaluating the
quality of our synthetic data. This metric trains a CatBoost[9] model on the
synthetic data instead of a weak ensemble. The CatBoost model is able to
detect subtle patterns in the data, that weak classifiers cannot. This is a
holistic metric that accounts for both diversity and quality of samples. For a
fair comparison we use the the standard hyper-parameter tuning budget of 50
trials[37][34]. Our full search space is provided in the Appendix. Each data
quality experiment was conducted using a single A10 GPU each. For evaluation,
we generate scores and standard deviations111For comparison, we use the same
error reporting method as prior work, however, we plan to present a thorough
coverage of more accurate error estimation for generative models in future
work. on the test set, training a CatBoost model 10 times on 5 samples of
synthetic data.

Table 1: MLE and standard deviations across techniques. The highest MLE score for each dataset is highlighted in bold. $F_{1}$ is used for classification. $R^{2}$ is used for regression tasks. *: the source paper [30] cites a $\sim 3\%$ higher test accuracy using the real train set over what other papers achieve, likely because they used a different split or version of this dataset.  DS | TVAE | CTabGAN+ | RealTab. | TabDDPM | TabMT | Real
---|---|---|---|---|---|---
AB |  0.433$\pm$0.008 |  0.467$\pm$0.004 |  0.504$\pm$0.011 | 0.550$\pm$0.010 |  0.535$\pm$0.004 |  0.556$\pm$0.004
AD |  0.781$\pm$0.002 |  0.772$\pm$0.003 |  0.811$\pm$0.002 |  0.795$\pm$0.001 | 0.814$\pm$0.001 |  0.815$\pm$0.002
BU |  0.864$\pm$0.005 |  0.884$\pm$0.005 |  0.928$\pm$0.003∗ |  0.906$\pm$0.003 | 0.908$\pm$0.002 |  0.906$\pm$0.002
CA |  0.752$\pm$0.001 |  0.525$\pm$0.004 |  0.808$\pm$0.003 |  0.836$\pm$0.002 | 0.838$\pm$0.002 |  0.857$\pm$0.001
CAR |  0.717$\pm$0.001 |  0.733$\pm$0.001 | - |  0.737$\pm$0.001 | 0.738$\pm$0.001 |  0.738$\pm$0.001
CH |  0.732$\pm$0.006 |  0.702$\pm$0.012 | - | 0.755$\pm$0.006 |  0.741$\pm$0.005 |  0.740$\pm$0.009
DI |  0.714$\pm$0.039 |  0.734$\pm$0.020 |  0.732$\pm$0.027 |  0.740$\pm$0.020 | 0.769$\pm$0.018 |  0.785$\pm$0.013
FB |  0.685$\pm$0.003 |  0.509$\pm$0.011 |  0.771$\pm$0.004 |  0.713$\pm$0.002 | 0.798$\pm$0.002 |  0.837$\pm$0.001
GE |  0.434$\pm$0.006 |  0.406$\pm$0.009 | - |  0.597$\pm$0.006 | 0.605$\pm$0.008 |  0.636$\pm$0.007
HI |  0.638$\pm$0.003 |  0.664$\pm$0.002 | - |  0.722$\pm$0.001 | 0.727$\pm$0.001 |  0.724$\pm$0.001
HO |  0.493$\pm$0.006 |  0.504$\pm$0.005 | - | 0.677$\pm$0.010 |  0.619$\pm$0.004 |  0.662$\pm$0.003
IN |  0.784$\pm$0.010 |  0.797$\pm$0.005 | - |  0.809$\pm$0.002 | 0.811$\pm$0.003 |  0.814$\pm$0.001
KI |  0.824$\pm$0.003 |  0.444$\pm$0.014 | - |  0.833$\pm$0.014 | 0.876$\pm$0.011 |  0.907$\pm$0.002
MI |  0.912$\pm$0.001 |  0.892$\pm$0.002 | - |  0.936$\pm$0.001 | 0.938$\pm$0.001 |  0.934$\pm$0.000
WI |  0.501$\pm$0.012 |  0.798$\pm$0.021 | - | 0.904$\pm$0.009 |  0.881$\pm$0.009 |  0.898$\pm$0.006

MLE scores are presented in Table 1; note that we match or exceed state-of-
the-art on 11 of 15 datasets. To gain a qualitative understanding of data
quality we visualize the distribution of correlation errors; see Figure 3. To
calculate the correlation errors, we first compute the correlation
$\textbf{r}_{i,j}$ between each pair of fields $(i,j)$. To compute
correlations involving categorical columns, we convert them to one-hot
vectors. We then compute the correlation between columns on the synthetic data
$\hat{\textbf{r}}_{i,j}$. The correlation error is the absolute difference
between these two values $|\textbf{r}_{i,j}-\hat{\textbf{r}}_{i,j}|$. These
errors should approach zero because we expect the correlations between columns
in the synthetic data to be the same as those in the real data. Erroneous and
missing correlations appear as non-zero values in the histograms of Figure 3.

Figure 3: A comparison of Correlation Error Histograms between TabMT (Green)
vs TabDDPM (Purple) and TabMT (Green) vs CTabGAN+ (Purple). A good generator
should have correlation errors distributed close to zero (the left of the
plot). We can see TabMT’s correlation errors are consistently distributed
closer to zero than either TabDDPM or CTabGAN+.

###  4.2 Privacy and Sample Novelty

Maintaining privacy of the original data is a key application for synthetic
tabular data. Machine learning is increasingly being used across a wide range
of areas to produce valuable insights. Simultaneously, there is a rapid rise
in both regulation and privacy concerns that need to be addressed. In Section
4.1 we demonstrated that data produced by TabMT is high enough quality to
produce strong classifiers. Now we evaluate our model for privacy. This
evaluation complements our quality evaluation and verifies that our model is
generating novel data. Novelty means data is not substantially similar to
samples encountered during training. A high-quality non-private model can
trivially be formed by directly reproducing the training set exactly. None of
this data is novel, but it is high quality. By ensuring our model is both
private and high quality, we verify that our model has learned the
distribution of the data, and not simply memorized the training set.
Memorization is a larger issue in tabular data due to smaller dataset sizes
and increased privacy concerns.

To evaluate privacy and novelty we adopt the median Distance to the Closest
Record (DCR) score. To calculate the DCR of a synthetic sample, we find the
nearest neighbor in the real training set by Euclidean distance. We report the
median of this distance across our synthetic samples. Data with higher DCR
scores will be more private and more novel. There is an inherent trade off
between privacy and quality. Higher quality samples will tend to be closer to
points in the training set and vice versa.

While models such as CTabGAN+[37] and TabDDPM[17] have a fixed trade-off
between privacy and quality after training. TabMT can trade-off between
quality and privacy using temperature scaling. By walking along the Pareto
curve of our model, using temperature scaling, we can controllably tune the
privacy and novelty of our generated data per application. By increasing a
field’s temperature, its generated values become more novel and private, but
they are also less faithful to the underlying data distribution. The trade off
between the quality and privacy here form a Pareto front for TabMT on each
dataset.

We use a separate temperature for each field and perform a search to estimate
the Pareto front. Each search was conducted using a single A10 GPU each.
Search details are available in the Appendix. In Table 2, we compare TabMT’s
DCR and corresponding MLE scores to that of TabDDPM. We are always able to
attain a higher DCR score, and in most cases a higher MLE score as well. This
falls in line with recent results in other domains showing diffusion models
are less private than other generative models[4]. A comparison with CTabGAN+
is available in the Appendix, compared to CTabGAN+ we obtain both higher
privacy and MLE scores in all tested cases. Figure 4 shows the Pareto fronts
of TabMT across several datasets.

Table 2: DCR score comparison between TabDDPM and TabMT. Corresponding MLE
scores are in parentheses.

DS | TabDDPM | TabMT
---|---|---
AB | 0.050(0.550) |  0.249(0.533)
AD | 0.104(0.795) |  1.01(0.811)
BU | 0.143(0.906) |  0.165(0.908)
CA | 0.041(0.836) |  0.117(0.832)
CAR | 0.012(0.737) |  0.041(0.737)
CH | 0.157(0.755) |  0.281(0.758)
DI | 0.204(0.740) |  0.243(0.740)
FB | 0.112(0.713) |  0.252(0.787)

DS | TabDDPM | TabMT
---|---|---
GE | 0.059(0.597) |  0.234(0.599)
HI | 0.449(0.722) |  0.483(0.727)
HO | 0.086(0.677) |  0.151(0.607)
IN | 0.041(0.809) |  0.061(0.816)
KI | 0.189(0.833) |  0.335(0.868)
MI | 0.022(0.936) |  0.026(0.936)
WI | 0.016(0.904) |  0.063(0.881)

Figure 4: Pareto Fronts of TabMT balancing the tradeoff between privacy (DCR)
and data quality (Validation Score). While some datasets have a smooth
transition the temperature changes, others have a sharp drop-off.

###  4.3 Missing Data

Real world data is often missing many values that make training difficult.
When a row has a missing value we must either drop the row, or find a method
to impute the missing value. Other techniques such as the RealTabformer[30] or
TabDDPM[17] cannot natively handle real world missing data, and must either
use a different imputation technique or drop the corresponding rows. Our
masking procedure allows TabMT to natively handle arbitrary missing data. To
demonstrate this, we randomly drop 25% of values from the dataset, ensuring
nearly every row is permanently missing data. Nevertheless, our model is still
able to train, producing synthetic rows with no missing values in them. This
facilitates training on real world data.Table 3 shows our accuracy when
training with missing data.

Table 3: MLE of TabMT when training with 25% of values missing. Delta represents the difference in MLE from training with no missing values. DS | MLE | Delta
---|---|---
AD | 0.813 | -0.001
KI | 0.868 | -0.008

Additionally, our model can be arbitrarily conditioned to produce any subset
of the data distribution at no additional cost, allowing us to more
effectively augment underrepresented portions of data. Prior art is largely
incapable of conditioning when producing outputs.

###  4.4 Scaling

In this section, we examine the performance of our model when scaling to very
large datasets. We use the CIDDS-001 dataset[27] as our benchmark dataset. We
do not use anomalous traffic from the dataset, and randomly select 5% of the
dataset as the validation set for reporting results. The results in this
section, together with those in Section 4.1, demonstrate that TabMT is both
sample-efficient enough to learn with just a few hundred samples, while
remaining general enough to scale to over thirty million samples. We train
three model sizes on this dataset and compute metrics on the resulting
samples. The model topologies are outlined in Table 4. Each model was trained
on a single A10 GPU with the exception of TabMT-L which was trained using 4
V100s. We use the AdamW[19] optimizer with a learning rate of 0.002 and weight
decay of 0.01, a batch size of 2048 and a cosine annealing learning rate
schedule for 350,000 training steps and 10000 warm-up steps.

Table 4: Model topologies used in scaling experiments. The large model sizes here demonstrate we can scale well in terms of model size and dataset size. Model | Width | Depth | Heads
---|---|---|---
TabMT-S | 64 | 12 | 4
TabMT-M | 384 | 12 | 8
TabMT-L | 576 | 24 | 12

Because ML Efficiency and DCR are very costly to compute on a dataset of this
scale, we instead adapt Precision and Recall[18] to the tabular domain. The
original definitions of these metrics[18] rely on vision models to produce the
embeddings used. We find our masking procedure produces strong embeddings for
each sample, so we use the embeddings produced by TabMT-S. Specifically, we
average the embeddings across the fields to produce 64 dimensional for each
flow. We used a fixed neighborhood size of $k=3$. We include an additional
diversity metric defined as the average set coverage across all properties of
the generated data. Results in Table 5 demonstrate strong scaling and
performance across model sizes. We can see Precision and Recall are both very
close to that of the validation set. Sample diversity and quality both scale
as model size increases.

Table 5: Precision, Recall, and Diversity metrics for all tested models. TabMT outperforms NFGAN[28] on all metrics Source | Validation Set | NFGAN | TabMT-S | TabMT-M | TabMT-L
---|---|---|---|---|---
Precision(%) | 90.42 | 77.64 | 82.58 | 85.77 | 88.10
Recall(%) | 90.58 | 63.32 | 91.88 | 91.82 | 91.12
Diversity(%) | 100.0 | 46.97 | 89.81 | 90.56 | 99.43

We compare against the prior state-of-the-art NetflowGAN, or NFGAN[28]. This
GAN was tuned specifically for this dataset. It is trained in two phases.
First IP2Vec[26] is trained to produce Netflow embeddings. These embeddings
are then used as targets for the generator during GAN training. Results from
NFGAN are shown in Table 5. We can see that NFGAN obtains reasonably high
precision, but poor recall and diversity. This is because the model suffers
from mode collapse, producing samples in only a small portion of the full
distribution.

Netflow has both correlations between the fields and complex invariants
between fields. We can measure the violation rate of these invariants to
understand how well our model is detecting patterns within the data. We
measure against seven invariants proposed in [28]. As shown in Table 6 TabMT
produces substantially more diverse data, while achieving a median 20x
improvement in violation probability over NFGAN.

Table 6: Error rates on netflow invariant tests. *: because we construct embeddings per field, our model cannot violate check 5. These tests check structural rules reflected in Netflow, such as the fact that two public ip addresses cannot communicate to each other. | NFGAN | TabMT-S | TabMT-M | TabMT-L
---|---|---|---|---
TCP Flags | 2.33e-03 | 4.63e-04 | 2.16e-04 | 6.54e-06
Private IPs | 2.00e-04 | 7.25e-05 | 2.67e-05 | 1.14e-05
TCP Port | 3.00e-04 | 9.32e-05 | 3.94e-05 | 1.47e-05
DNS | 1.60e-03 | 2.34e-03 | 1.56e-03 | 2.05e-04
Valid Values* | 2.00e-03 | 0.00e-00 | 0.00e-00 | 0.00e-00
NetBios | 7.43e-01 | 4.73e-03 | 4.64e-03 | 1.27e-03
Packet Ratios | 5.10e-03 | 2.28e-03 | 1.79e-03 | 1.16e-03

##  5 Limitations and Future Work

TabMT presents strong results even on large datasets, but our Transformer
backbone means TabMT is slower than more lightweight methods built around
small MLPs, or GANs which can produce a row in a single inference. Searching
temperatures also adds time if optimal privacy is needed. Additionally, we
must quantize continuous fields, while we outperform methods which do not
quantize fields, this could pose issues in some applications. Future work
might examine learning across tabular datasets, alternative masking procedures
and networks to improve speed, or integration with diffusion models to better
tackle continuous fields.

#### Broader Impact

Synthetic data generation allows for privacy preservation, protecting
sensitive data while still enabling data analysis. High Quality synthetic data
may ease the pressure to resort to unethical methods of collection such as
relying on underpaid labor. With this in mind, trading off data for additional
compute does mean that the additional compute will contribute to increased CO2
emissions. Additionally, synthetic data carries the risk of misuse, such as
the potential for manipulating results or research findings with fabricated
data. All experiments were conducted using cloud A10 or V100 GPUs. For
algorithm design and experiment result generation roughly 410 GPU days of
compute were used.

##  6 Conclusion

In this paper, we outlined a novel Masked Transformer design and training
procedure, TabMT, for generating synthetic tabular data. Through a
comprehensive series of benchmarks we demonstrate that our model achieves
state-of-the-art generation quality. This quality is verified at scales that
are orders of magnitude larger than prior work and with missing data present.
Our model achieves superior privacy and is able to easily trade off between
privacy and quality. Our model is a substantial advancement compared to
previous work, due to its scalability, missing data robustness, privacy-
preserving generation, and superior data quality.

## References

  * Akiba et al. [2019] Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama.  Optuna: A next-generation hyperparameter optimization framework.  In _Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining_, pages 2623–2631, 2019.
  * Baldi et al. [2014] Pierre Baldi, Peter Sadowski, and Daniel Whiteson.  Searching for exotic particles in high-energy physics with deep learning.  _Nature communications_ , 5(1):4308, 2014.
  * Borisov et al. [2022] Vadim Borisov, Kathrin Seßler, Tobias Leemann, Martin Pawelczyk, and Gjergji Kasneci.  Language models are realistic tabular data generators.  _arXiv preprint arXiv:2210.06280_ , 2022.
  * Carlini et al. [2023] Nicholas Carlini, Jamie Hayes, Milad Nasr, Matthew Jagielski, Vikash Sehwag, Florian Tramer, Borja Balle, Daphne Ippolito, and Eric Wallace.  Extracting training data from diffusion models.  _arXiv preprint arXiv:2301.13188_ , 2023.
  * Chang et al. [2022] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman.  Maskgit: Masked generative image transformer.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 11315–11325, 2022.
  * Chen et al. [2020] Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever.  Generative pretraining from pixels.  In _International conference on machine learning_ , pages 1691–1703. PMLR, 2020.
  * Deb et al. [2002] K. Deb, A. Pratap, S. Agarwal, and T. Meyarivan.  A fast and elitist multiobjective genetic algorithm: Nsga-ii.  _IEEE Transactions on Evolutionary Computation_ , 6(2):182–197, 2002.  doi: 10.1109/4235.996017.
  * Devlin et al. [2018] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Dorogush et al. [2018] Anna Veronika Dorogush, Vasily Ershov, and Andrey Gulin.  Catboost: gradient boosting with categorical features support.  _arXiv preprint arXiv:1810.11363_ , 2018.
  * Engelmann and Lessmann [2021] Justin Engelmann and Stefan Lessmann.  Conditional wasserstein gan-based oversampling of tabular data for imbalanced learning.  _Expert Systems with Applications_ , 174:114582, 2021.
  * Floridi and Chiriatti [2020] Luciano Floridi and Massimo Chiriatti.  Gpt-3: Its nature, scope, limits, and consequences.  _Minds and Machines_ , 30:681–694, 2020.
  * He et al. [2022] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.  Masked autoencoders are scalable vision learners.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 16000–16009, 2022.
  * Huang et al. [2022] Po-Yao Huang, Hu Xu, Juncheng Li, Alexei Baevski, Michael Auli, Wojciech Galuba, Florian Metze, and Christoph Feichtenhofer.  Masked autoencoders that listen.  _Advances in Neural Information Processing Systems_ , 35:28708–28720, 2022.
  * Jiang et al. [2021] Yifan Jiang, Shiyu Chang, and Zhangyang Wang.  Transgan: Two pure transformers can make one strong gan, and that can scale up.  _Advances in Neural Information Processing Systems_ , 34:14745–14758, 2021.
  * Kim et al. [2022] Jayoung Kim, Chaejeong Lee, and Noseong Park.  Stasy: Score-based tabular data synthesis.  _arXiv preprint arXiv:2210.04018_ , 2022.
  * Kohavi et al. [1996] Ron Kohavi et al.  Scaling up the accuracy of naive-bayes classifiers: A decision-tree hybrid.  In _Kdd_ , volume 96, pages 202–207, 1996.
  * Kotelnikov et al. [2022] Akim Kotelnikov, Dmitry Baranchuk, Ivan Rubachev, and Artem Babenko.  Tabddpm: Modelling tabular data with diffusion models.  _arXiv preprint arXiv:2209.15421_ , 2022.
  * Kynkäänniemi et al. [2019] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila.  Improved precision and recall metric for assessing generative models.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Loshchilov and Hutter [2019] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization, 2019.
  * Madeo et al. [2016] Renata Cristina Barros Madeo, Sarajane Marques Peres, and Clodoaldo Aparecido de Moraes Lima.  Gesture phase segmentation using support vector machines.  _Expert Systems with Applications_ , 56:100–115, 2016.
  * Manocchio et al. [2021] Liam Daly Manocchio, Siamak Layeghy, and Marius Portmann.  Flowgan-synthetic network flow generation using generative adversarial networks.  In _2021 IEEE 24th International Conference on Computational Science and Engineering (CSE)_ , pages 168–176. IEEE, 2021.
  * Pace and Barry [1997] R Kelley Pace and Ronald Barry.  Sparse spatial autoregressions.  _Statistics & Probability Letters_, 33(3):291–297, 1997.
  * Padhi et al. [2021] Inkit Padhi, Yair Schiff, Igor Melnyk, Mattia Rigotti, Youssef Mroueh, Pierre Dognin, Jerret Ross, Ravi Nair, and Erik Altman.  Tabular transformers for modeling multivariate time series.  In _ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 3565–3569. IEEE, 2021.
  * Radford et al. [2022] Alec Radford, Jong Wook Kim, Tao Xu, Greg Brockman, Christine McLeavey, and Ilya Sutskever.  Robust speech recognition via large-scale weak supervision.  _arXiv preprint arXiv:2212.04356_ , 2022.
  * Ramesh et al. [2021] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever.  Zero-shot text-to-image generation.  In _International Conference on Machine Learning_ , pages 8821–8831. PMLR, 2021.
  * Ring et al. [2017a] Markus Ring, Alexander Dallmann, Dieter Landes, and Andreas Hotho.  Ip2vec: Learning similarities between ip addresses.  In _2017 IEEE International Conference on Data Mining Workshops (ICDMW)_ , pages 657–666, 2017a.  doi: 10.1109/ICDMW.2017.93.
  * Ring et al. [2017b] Markus Ring, Sarah Wunderlich, Dominik Grüdl, Dieter Landes, and Andreas Hotho.  Flow-based benchmark data sets for intrusion detection.  In _Proceedings of the 16th European conference on cyber warfare and security. ACPI_ , pages 361–369, 2017b.
  * Ring et al. [2019] Markus Ring, Daniel Schlör, Dieter Landes, and Andreas Hotho.  Flow-based network traffic generation using generative adversarial networks.  _Computers & Security_, 82:156–172, 2019.
  * Singh [2015] Kamaljot Singh.  Facebook comment volume prediction.  _International Journal of Simulation: Systems, Science and Technologies_ , 16(5):16–1, 2015.
  * Solatorio and Dupriez [2023] Aivin V Solatorio and Olivier Dupriez.  Realtabformer: Generating realistic relational and tabular data using transformers.  _arXiv preprint arXiv:2302.02041_ , 2023.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Wang et al. [2021] Yingfan Wang, Haiyang Huang, Cynthia Rudin, and Yaron Shaposhnik.  Understanding how dimension reduction tools work: An empirical approach to deciphering t-sne, umap, trimap, and pacmap for data visualization.  _Journal of Machine Learning Research_ , 22(201):1–73, 2021.  URL http://jmlr.org/papers/v22/20-1061.html.
  * Wolf et al. [2021] Maximilian Wolf, Markus Ring, and Dieter Landes.  Impact of generative adversarial networks on netflow-based traffic classification.  In _13th International Conference on Computational Intelligence in Security for Information Systems (CISIS 2020) 12_ , pages 393–404. Springer, 2021.
  * Xu et al. [2019] Lei Xu, Maria Skoularidou, Alfredo Cuesta-Infante, and Kalyan Veeramachaneni.  Modeling tabular data using conditional gan.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Yu et al. [2022] Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, et al.  Scaling autoregressive models for content-rich text-to-image generation.  _arXiv preprint arXiv:2206.10789_ , 2022.
  * Zhao et al. [2021] Zilong Zhao, Aditya Kunar, Hiek Van der Scheer, Robert Birke, and Lydia Y. Chen.  CTAB-GAN: effective table data synthesizing.  _CoRR_ , abs/2102.08369, 2021.  URL https://arxiv.org/abs/2102.08369.
  * Zhao et al. [2022] Zilong Zhao, Aditya Kunar, Robert Birke, and Lydia Y Chen.  Ctab-gan+: Enhancing tabular data synthesis.  _arXiv preprint arXiv:2204.00401_ , 2022.

##  Appendix A Additional Figures

Figure 5: Additional Pareto Fronts Figure 6: Fake and Real Data Manifolds on
the CIDDS-001[27] dataset for TabMT using PaCMAP[32]

##  Appendix B Additional Tables

Hyperparameter | Search Space
---|---
Network Dim. | Range(16, 512, 16)
Num. of Heads | Choice(1, 2, 4, 8, 16)
Network Depth | Range(2, 12, 2)
Drop Path Rate | Range(0.0, 0.25, 0.05)
Dropout Rate | Range(0.0, 0.5, 0.05)
Learning Rate | LogInterval(5e-5, 5e-3)
Weight Decay | LogInterval(0.001, 0.03)
Max Steps | Range(10000, 40000, 5000)
Batch Size | Range(256, 2048, 256)
Max Bins | Range(10, 200, 10)
Data Multiplier | 8
Gen. Batch Size | 512
Optimizer |  AdamW[19]
Annealing | Cosine
Quantizer | KMeans
Num. Trials | 50
Table 7: The hyperparameter search space for training TabMT. Max bins
represents the maximum number of bins used during quantization of continuous
attributes. We find many hyperparameters in this search space to be
unimportant to final model quality, as such the space can likely be pruned for
better tuning results. We use the Optuna[1] framework with default settings
for optimization

Field | Type
---|---
Date | timestamp
Source IP | categorical
Destination IP | categorical
Protocol | categorical
Source Port | categorical
Destination Port | categorical
Duration | continuous
Bytes | ordinal
Packets | ordinal
Flags | categorical
Type of Service | categorical
Table 8: Original Netflow Structure

Field | Cardinality
---|---
Weekday | 7
Hour | 24
Minute | 60
Second | 60
Millisecond | 1000
Source IP | 51092
Destination IP | 50923
Protocol | 5
Source Port | 61700
Destination Port | 61448
Duration | 37787
Bytes | 181536
Packets | 10489
Flags | 34
Type of Service | 5
Table 9: Post-processed Netflow Structure. All fields are processed to be
categorical variables for training.

Hyperparameter | Search Space
---|---
$\tau_{i}$ | Interval(0.5, 5.0)
sampler |  NSGAII[7]
Table 10: Hyperparamters used when tuning temperatures, we search a separate temperature for each field, optimizing privacy and ML Efficiency simultaneously. This can certainly be done more efficiently, as NSGAII is sample inefficient for this use case. We also use the full 8x data multiplier when searching here, using a smaller multiplier or number of trial repeats during search would speed this up significantly. We didn’t use an explicit trial budget and qualitatively evaluated convergence. This took 0.25-5 GPU days except for MI and FB which each took roughly 15. DS | CTabGAN+ | TabMT
---|---|---
AB | 0.075(0.467) |  0.249(0.533)
AD | 0.119(0.772) |  1.01(0.811)
BU | 0.164(0.864) |  0.165(0.908)
CA | 0.056(0.525) |  0.117(0.832)
CAR | 0.012(0.733) |  0.041(0.737)
CH | 0.212(0.702) |  0.281(0.758)
DI | 0.196(0.734) |  0.243(0.740)
FB | 0.427(0.509) |  0.429(0.566)
Table 11: DCR scores for CTabGAN+ and TabMT. MLE scores are in parentheses. We win on both privacy and MLE for all tested datasets. Our increased privacy on FB vs. the main privacy table demonstrates our models unique ability to control the trade-off between MLE and DCR DS | STaSY | TabMT
---|---|---
AB | 0.482 | 0.535
AD | 0.790 | 0.814
BU | 0.881 | 0.908
CA | 0.762 | 0.838
CAR | 0.725 | 0.738
CH | 0.738 | 0.741
DI | 0.727 | 0.769
Table 12: MLE scores for a recent work, STaSY[15], compared to TabMT. TabMT
obtains a higher score on all tested datasets. MLP Width, depth, and learning
rate were included in the search.

##  Appendix C Pseudocode

Snippet 1: Ordered Embedding

⬇

class OrderedEmbedding:

def __init__(self, occ: Tensor, dim: int):

# occ: ordered cluster centers

self.E = zeros(len(occ), dim)

self.l = randn(dim) * 0.05

self.h = randn(dim) * 0.05

self.r = (occ - occ[0]) / (occ[-1] - occ[0])

@property

def weight(self):

return self.r * self.l + (1 - self.r) * self.h + self.E

def forward(self, idx: Tensor):

return self.weight[idx]

Snippet 2: Dynamic Linear layer

⬇

class DynamicLinear:

def __init__(self, embedding: OrderedEmbedding | Embedding):

self.E = embedding

self.temp = ones(1)

self.bias = zeros(len(embedding))

def forward(self, x: Tensor):

raw_logits = x @ self.E.weight.T + self.bias

return raw_logits / sigmoid(self.temp)

Snippet 3: Training step

⬇

def training_step(mdl: MaskedTransformer, batch: Tensor):

mask = rand_like(batch) > rand(len(batch), 1)

preds = mdl(batch, mask)

batch[~mask] = -1 #ignore unmasked entries in loss

loss = cross_entropy(preds, batch)

return loss

Snippet 4: Batched generation

⬇

def gen_batch(mdl: MaskedTransformer, n: int, l: int, temps: Tensor):

batch = zeros(n, l)

mask = ones_like(batch)

for i in randperm(l):

preds = mdl(batch, mask)

batch[:, i] = Categorical(logits=preds[i] / temps[i]).sample()

mask[:, i] = False

return batch

MaskedTransformer: We use a standard Transformer Encoder architecture. We use
learned positional embeddings initialized with normal distribution with
$\sigma=0.01$. We use a LayerNorm before the DynamicLinear layers. Our mask
token is intialized with a normal distribution as well with $\sigma=0.05$, to
match the other embeddings.

Embedding: We use a standard embedding initialized with a normal distribution
and $\sigma=0.05$, and no magnitude clipping or rescaling during training.

##  Appendix D Dataset Info

For comparison, we use the same datasets and splits as Kotelnikov et al. [17]
in our main comparison. We list them below along with their sources.

  * •

AB: Abalone: OpenML, CC Licensce

  * •

AD: Adult ROC[16]: UCI Data Repo, CC License

  * •

BU: Adopt a Buddy: Kaggle, Other

  * •

CA: California Housing[22]: Kaggle, CC License

  * •

CAR: Cardiovascular Disease: Kaggle, CC License

  * •

CH: Churn Modeling: Kaggle, CC License

  * •

DI: Diabetes: OpenML, CC Licensce

  * •

FB: Facebook Comment Volume[29]: UCI Data Repo, CC License

  * •

GE: Gesture Phase Prediction[20]: UCI Data Repo, CC License

  * •

HI: Higgs (98K)[2]: OpenML, CC License

  * •

HO: House 16H: OpenML, CC Licensce

  * •

IN: Medical Costs, Kaggle, ODbL

  * •

KI: King Count Housing Prices: OpenML, CC License

  * •

MI: MiniBOONE Particle Prediction: OpenML, CC License

  * •

WI: Wilt Remote Sensing: OpenML, CC License

We additionally use the CIDDS-001[27] dataset for scaling experiments, which
has a CC license.