[Paper Review] GENIUS: Sketch-based Language Model Pre-training via Extreme and Selective Masking for Text Generation and Augmentation

Biyang Guo, Yeyun Gong+5 more|arXiv (Cornell University)|Nov 18, 2022

Topic Modeling4 citations

TL;DR

GENIUS is a sketch-based text generation model pre-trained via extreme and selective masking to reconstruct full texts from minimal key information (sketches), enabling high-quality, diverse generation. It powers GeniusAug, a data augmentation method that extracts task-aware sketches and generates semantically rich, diverse training samples, significantly improving performance across text classification, NER, and MRC tasks without fine-tuning.

ABSTRACT

We introduce GENIUS: a conditional text generation model using sketches as input, which can fill in the missing contexts for a given sketch (key information consisting of textual spans, phrases, or words, concatenated by mask tokens). GENIUS is pre-trained on a large-scale textual corpus with a novel reconstruction from sketch objective using an extreme and selective masking strategy, enabling it to generate diverse and high-quality texts given sketches. Comparison with other competitive conditional language models (CLMs) reveals the superiority of GENIUS's text generation quality. We further show that GENIUS can be used as a strong and ready-to-use data augmentation tool for various natural language processing (NLP) tasks. Most existing textual data augmentation methods are either too conservative, by making small changes to the original text, or too aggressive, by creating entirely new samples. With GENIUS, we propose GeniusAug, which first extracts the target-aware sketches from the original training set and then generates new samples based on the sketches. Empirical experiments on 6 text classification datasets show that GeniusAug significantly improves the models' performance in both in-distribution (ID) and out-of-distribution (OOD) settings. We also demonstrate the effectiveness of GeniusAug on named entity recognition (NER) and machine reading comprehension (MRC) tasks. (Code and models are publicly available at https://github.com/microsoft/SCGLab and https://github.com/beyondguo/genius)

Motivation & Objective

• To develop a conditional language model that generates coherent, diverse text from minimal sketch inputs (key phrases or spans) using mask tokens.
• To overcome limitations of existing text reconstruction models, which are constrained by low masking ratios and cannot handle highly corrupted inputs.
• To design a pre-training objective using extreme and selective masking that enables robust reconstruction from sparse sketches.
• To demonstrate the utility of GENIUS as a strong, ready-to-use data augmentation tool for diverse NLP tasks.
• To balance diversity and quality in data augmentation by preserving core semantics while generating varied contexts.

Proposed method

• Pre-train GENIUS on a large-scale corpus using a novel reconstruction-from-sketch objective with extreme masking (high proportion of tokens masked) and selective masking (targeting salient, informative spans).
• Use mask tokens to represent sketch elements, where only key phrases or words are retained and the rest of the text is masked.
• Apply a selective masking strategy that prioritizes important semantic units (e.g., named entities, sentiment-bearing phrases) for retention in the sketch.
• Leverage the pre-trained GENIUS model to generate full, coherent, and diverse text conditioned on the sketch input.
• Design GeniusAug: extract target-aware sketches from original training data, then use GENIUS to generate new training samples around these sketches.
• Integrate GeniusAug into downstream NLP pipelines without requiring task-specific fine-tuning, enabling plug-and-play data augmentation.

Experimental results

Research questions

• RQ1 Can a language model effectively reconstruct full, coherent, and diverse text from extremely masked inputs consisting only of key phrases or spans?
• RQ2 Does extreme and selective masking during pre-training lead to better generalization and generation quality compared to standard masking strategies?
• RQ3 Can sketch-based generation via GENIUS serve as a powerful and controllable data augmentation method for downstream NLP tasks?
• RQ4 How does GeniusAug compare to conservative (e.g., EDA, MLM-based) and aggressive (e.g., LAMBADA) data augmentation methods in terms of diversity, quality, and downstream performance?
• RQ5 To what extent can GeniusAug improve model generalization in both in-distribution and out-of-distribution settings?

Key findings

• GeniusAug significantly improves model performance on six text classification datasets, with gains of up to 4.2% in F1-score on average across in-distribution and out-of-distribution settings.
• GeniusAug outperforms strong baselines like EDA, back-translation, and MLM-based methods in both diversity and downstream accuracy, especially in low-resource scenarios.
• The method achieves state-of-the-art or competitive results on named entity recognition and machine reading comprehension tasks, demonstrating broad applicability.
• GeniusAug generates high-quality samples with strong semantic fidelity to the original while introducing meaningful structural and lexical variation, avoiding the pitfalls of both conservative and aggressive augmentation.
• The pre-trained GENIUS model enables zero-shot data augmentation without fine-tuning, and further fine-tuning on downstream data yields additional performance gains.
• GeniusAug is more controllable than AR-based methods like LAMBADA, as it preserves core semantics while generating diverse contexts, reducing the risk of low-quality or irrelevant samples.