Retrieval-Based Transformer for Table Augmentation (2023)

Abstract: Data preparation, also called data wrangling, is considered one of the most expensive and time-consuming steps when performing analytics or building machine learning models. Preparing data typically involves collecting and merging data from complex heterogeneous, and often large-scale data sources, such as data lakes. In this paper, we introduce a novel approach toward automatic data wrangling in an attempt to alleviate the effort of end-users, e.g. data analysts, in structuring dynamic views from data lakes in the form of tabular data. Given a corpus of tables, we propose a retrieval augmented transformer model that is self-trained for the table augmentation tasks of row/column population and data imputation. Our self-learning strategy consists in randomly ablating tables from the corpus and training the retrieval-based model with the objective of reconstructing the partial tables given as input with the original values or headers. We adopt this strategy to first train the dense neural retrieval model encoding portions of tables to vectors, and then the end-to-end model trained to perform table augmentation tasks. We test on EntiTables, the standard benchmark for table augmentation, as well as introduce a new benchmark to advance further research: WebTables. Our model consistently and substantially outperforms both supervised statistical methods and the current state-of-the-art transformer-based models.

## introduction the way organizations store and manage data is rapidly evolving from using strict transactional databases to data lakes that typically consist of large collections of heterogeneous data formats, such as tabular data, spreadsheets, and nosql databases. the absence of a unified schema in data lakes does not allow the usage of declarative query languages, e.g. sql, making the process of data preparation foot_0 dramatically expensive @xcite . data preparation involves several phases, such as data discovery, structuring, cleansing, enrichment and validation, with the purpose of producing views commonly organized in a tabular format used to create reports @xcite or to gather feature sets to build machine learning models @xcite . the schemaless nature of data lakes makes data discovery and structuring even more challenging since the tasks of joinability and unionability among tables become non-deterministic @xcite @xcite . in this work, we propose a novel end-to-end solution based on a retrieval augmented transformer architecture with the aim to support end-users, such as data analysts, in the process of constructing dynamic views from data lakes. to this end, we address three table augmentation tasks @xcite : automatic row and column population and cell filling (or data imputation). figure 1 illustrates the three core tasks in table augmentation. all tasks proceed from a query or seed table. in the case of self-supervised training, this seed table is formed by ablating rows, columns or cell values from an existing table in the data lake. the task of column header population, also simply called column population, is to extend the table with additional possible column names or headers. this is a way of suggesting additional data that could be joined into this table. in the task of cell filling there is a specific unknown cell, for which the model predicts a specific value. the task of row population is only populating the key column for a row. this is the column that contains the primary entity that the remainder of the row contains data for, sometimes referred to as a row header. typically this is the first column in a table. approaches to table augmentation can be purely parametric @xcite , in which case the data lake is used to train the param-figure 1 : given a partially completed table as a query (i.e. a few album releases from the pink floyd discography), the three table augmentation tasks consist of retrieving from the data lake: 1) a list of possible next column headers, such as the "label" or "format", 2) the missing value "1979" for the release date of the row "the wall", 3) a list of other album releases as possible next rows, such as "atom heart mother" and "the division bell". eters of the model, but not used during inference. in this setting, the table augmentation model must draw the possible augmentations for rows, columns and cells from its trained parameters. alternatively, with retrieval-based models @xcite @xcite , the data lake can also be used at inference to provide evidence for proposed augmentations. this has two key advantages: 1) the model need not memorize the data lake -or even a significant fraction of it, and 2) the model can provide justification for its predicted augmentations in the form of a provenance table or tables. the key contributions of this paper are: (1) we introduce the first end-to-end, retrieval-based model for table augmentation. our retrieval augmented table augmentation (rata) model uses a biencoder retrieval model for neural indexing and searching tables from data lake, and a reader transformer to identify augmentations from retrieved tables. (2) our model establishes a new state-ofthe-art across all three tasks in table augmentation, while also providing additional value with its provenance information. (3) we create and release a new dataset for table augmentation, expanding the scope of evaluation beyond wikipedia. this dataset, based on @xcite , is also larger and more diverse than the standard wikipedia-based dataset @xcite . ## related work table augmentation can be divided into three sub-tasks: row population, column population, and cell filling. for row and column population, @xcite identifies and ranks candidate val-ues from both the table corpus and knowledge base. table2vec @xcite trains header and entity embeddings from a table corpus in a skipgram manner and uses the embeddings for the task. although tabert @xcite was developed as a foundational model primarily for question answering, its embeddings have also been applied for row and column population. recent work formulates the task as multi-label classification and fine-tunes large-scale pre-trained models such as tabbie @xcite and turl @xcite . tabbie consists of three transformers for converting cells, columns and rows to vector representations. a corrupt cell detection task is the pretraining task used to learn these embeddings on the table corpus. to fine-tune a trained tabbie model for the column header population task, a concatenated [clscol] embedding of the columns is passed through a single linear and softmax layer and trained with a multi-label classification objective. similarly, for the row population task a multi-class classification is carried out on the first column's [clscol] representation. for cell filling, infogather @xcite retrieves tables from the table corpus and selects values from retrieved tables. @xcite extends the system to retrieve from both the table corpus and knowledge base. their system that uses only the table corpus as the source is called tmatch, which we compare to in section 6. ahmadov et al. (2015) combines predictions both from table retrieval and from a machine learningbased value imputation system. deng et al. (2022) (a) building the table-part index. query encoder ann index cell or header selection ? query table ? top-k tables parts figure 2 : index building and inference system overviews directly applies pre-trained turl model to the task since cell filling is similar with its pre-training objective. cell filling is also related to the task of value imputation, i.e., to provide an assumed value when the actual value is unknown, usually using machine learning methods @xcite . in addition to augmenting individual entities, column headers or cells, some other work aims to join tables over entire rows or columns with retrieved tables @xcite @xcite . retrieval-augmented models have been successfully applied to many tasks. for open-domain question answering (odqa), dpr learns dense representation to retrieve evidence and trains a separate reader to select answer from retrieved evidence @xcite . rag uses a generator to generate outputs conditioned on retrieved evidence and jointly trains dpr with a generator on the downstream task @xcite . rag is shown to achieve good performance on knowledgeintensive nlp tasks such as odqa, fact verification, slot filling, etc @xcite . re 2 g further introduces a reranker to boost performance @xcite . retrievalaugmented models are also shown to be effective on zero-shot slot filling @xcite , and multilingual keyphrase generation @xcite . similar models have also been applied to tablerelated tasks such as open-domain table question answering @xcite . in our work, we apply the architecture to table augmentation. ## approach while the row, column, and cell predictions of purely parametric table augmentation methods may be useful on their own, they can be much more effective for a human-in-the-loop use case if they are supported by provenance. a user of a data preparation application may be unwilling to simply accept the prediction of a model, but when paired with evidence from the data lake, that prediction can be better assessed. furthermore, the retrieval model itself may be useful for exploration and general search in a data lake. in this view, table augmentation can be seen as self-supervised pretraining for table retrieval. fortunately, there is now considerable work on retrieval augmented transformer models @xcite . these models augment the parametric knowledge of the transformer, with non-parametric knowledge in the form of an indexed corpus. to do so, they use a neural retrieval model based on dpr (dense passage retrieval) @xcite ) that is trained end-to-end to assist in generation. we build on this line of research to introduce a general model for all table augmentation tasks: row population, column header population and cell filling. our model, retrieval augmented table augmentation (rata), comprises of an index of tables, a retrieval component, and a reader or selection component. the table index is built from the tables in the training set, which are first decomposed into table-parts, then transformed into sequences for use with standard retrieval approaches. the retrieval component is a biencoder architecture similar to dpr @xcite , but trained without ground truth on correct provenance. we call this dense table retrieval or dtr. the reader component is an extractive approach. an extractive rather than generative approach ensures that the model's predictions are always grounded in actual data, rather than speculative guesses. the extractive approach is also a more natural fit for row and column population tasks, where there is no required order to the answers. finally, the extractive approach permits an initial training phase for the retrieval component where the answer-bearing tables are considered as a bag of positives. figure 1 illustrates the tasks of table augmentation by example. formally, the input i is a table with r rows and c columns comprising a caption c, headers h, and matrix of cell values, v. one of the columns, usually the first, is indicated as the key column key. the input table is ablated in a task specific way to produce a query table and gold answers, q, g , described as follows: ## webtables dataset prior work on table augmentation has focused on tables derived from wikipedia @xcite @xcite @xcite . in order to better assess the proposed methods and provide the research community with a new benchmark, we introduce a new dataset for table augmentation: webtables. we construct this dataset using the tables crawled and extracted by @xcite . we start from the english relational tables of wdc web table corpus 2015. we further filter the dataset to remove the most common types of noisy tables: calendars formatted as tables, lists of forum posts and torrent links, tables with less than four rows or columns, and tables that format large blocks of text. because previous work on table augmentation focused so heavily on wikipedia tables, we exclude from this dataset any tables crawled from any "wikipedia" domain. we also deduplicate the corpus, ensuring that there are no two tables with the same content in their cells. following filtering and deduplication we sample 10 thousand tables each for the development and test sets and one million tables for training. however, in our experiments we use only 300 thousand training examples to limit the computational cost. to parallel the setting of entitables we use the "key column" identified by @xcite as the target column for row population and we consider entities to be those strings that occur at least three times in the key column for any table in the train set. ## experiments we experiment on two datasets of tables across three tasks. entitables @xcite ) contains 1.6m tables collected from wikipedia where entity mentions are normalized into its name in dbpedia. for row and column population, we use the development and test sets released by @xcite each containing 1,000 randomly sampled queries. for cell filling, we use the test set released by @xcite . the test set contains 1,000 queries uniformly sampled from four main column data types: entity, quantity, string, and datetime. though @xcite use human annotations as gold labels, we notice that the human annotations are of low quality, so we use the original values in the table cells as gold labels. webtables is based on cafarella et al. ( 2008 ) -154m relational tables extracted from html tables in common crawl. we process the corpus as described in section 4. for column population we use the original development and test sets of 10,000 tables each. while for row population we necessarily exclude any tables without any entities in the key column after the first n seed rows. for cell filling, we use heuristic rules to classify cell values into three types: quantity, string and datetime. then, we sample 3,000 queries uniformly from the three types as test set and sample another 3,000 queries as development set. we compare our method with two deep learningbased baselines, tabbie @xcite and bart @xcite . both tabbie and bart have no retrieval component involved. tabbie, described in section 2, uses three transformers: one for cell values, one for rows, and one for columns. it produces vector embeddings for each cell and each row and column of a table. we follow @xcite for the row and column population and base our experiments on the par-tial released code and pretrained model foot_2 . to apply tabbie to cell filling, we formulate it as classification on the concatenation of the row and column embedding vectors, similar to row and column population. the classification vocabulary is collected from the training corpus: all cell values that occur at least ten times. we also report the published results for tabbie on the entitables dataset, although we were unable to reproduce these results for row population. bart is a sequence-to-sequence model that takes the linearized table as the source text and generates the row entities, cell headers, or cell value as the target text. we use a beam search in decoding (beam @xmath0. we use the fairseq toolkit @xcite for these experiments. for rag we use the implementation in hugging face transformers @xcite . for both bart and rag, the sequence representation of the query tables is the same as in rata. on the entitables dataset, we also compare against probabilistic methods that first retrieve tables from the table corpus and next select values for table augmentation. we compare against the published results of @xcite for row and column population, and against tmatch @xcite for cell filling. for evaluation, we report mean reciprocal rank (mrr) and normalized discounted cumulative gain over the top ten outputs (ndcg@10) for the final prediction performance of row population, column population, and cell filling. to evaluate the performance of dtr retrieval, we also report answer-bearing mrr, where a retrieved table-part is considered correct if it contains one of the correct answers. to determine the significance of these results we use a 95% confidence interval on the t-distribution. we also applied a sampling permutation test, but this did not change any conclusions regarding significance. ## results table-parts row population column population cell filling figure 4: mrr gain as number of retrieved table-parts increases on the entitables dataset