Unlocking Noisy Real-World Corpora for Foundation Model Pre-Training via Quality-Aware Tokenization

Current tokenization methods process sequential data without accounting for signal quality, limiting their effectiveness on noisy real-world corpora. We present QA-Token (Quality-Aware Tokenization), which incorporates data reliability directly into vocabulary construction. We make three key contributions: (i) a bilevel optimization formulation that jointly optimizes vocabulary construction and downstream performance, (ii) a reinforcement learning approach that learns merge policies through quality-aware rewards with convergence guarantees, and (iii) an adaptive parameter learning mechanism via Gumbel-Softmax relaxation for end-to-end optimization. Our experimental evaluation demonstrates consistent improvements: genomics (6.7 percentage point F1 gain in variant calling over BPE), finance (30% Sharpe ratio improvement). At foundation scale, we tokenize a pretraining corpus comprising 1.7 trillion base-pairs and achieve state-of-the-art pathogen detection (94.53 MCC) while reducing token count by 15%. We unlock noisy real-world corpora, spanning petabases of genomic sequences and terabytes of financial time series, for foundation model training with zero inference overhead.

💡 Research Summary

The paper introduces Quality‑Aware Tokenization (QA‑Token), a novel framework that integrates data quality signals directly into the vocabulary construction process for sequential data such as genomic sequences and financial time‑series. Traditional tokenizers (e.g., BPE, SentencePiece) rely solely on frequency statistics and treat all input symbols equally, which leads to suboptimal performance when the underlying data contain substantial noise—sequencing errors in genomics or micro‑structure noise in high‑frequency finance.

QA‑Token addresses this gap through three technical pillars. First, it formulates tokenization as a bilevel optimization problem: the upper level maximizes downstream language‑model performance while penalizing vocabulary size and rewarding token quality; the lower level selects up to K merge operations to build the vocabulary. The authors prove the problem is NP‑hard and provide a principled approximation scheme. Second, they cast vocabulary construction as a Markov Decision Process (MDP) and train a merge‑selection policy using Proximal Policy Optimization (PPO). The reward function is multi‑objective, combining a quality‑aware term, an information‑gain term, a complexity penalty, and domain‑specific constraints, all normalized via exponential moving averages to ensure bounded, scale‑invariant signals. Third, they enable end‑to‑end learning of quality‑sensitivity parameters (e.g., α controlling how strongly quality influences merges) using a Gumbel‑Softmax relaxation. The training proceeds in two stages: a fast‑timescale RL phase that discovers promising merge policies, followed by a slow‑timescale adaptive‑parameter phase that fine‑tunes α, domain weights, and other continuous hyper‑parameters. The authors prove convergence guarantees for both stages, including a (1‑1/e) approximation bound derived from adaptive submodularity, bounded gradients for the Gumbel‑Softmax estimator, and stationary‑point convergence for PPO.

Quality metrics are defined per domain. In genomics, Phred scores are decayed by position (learned β_pos) and aggregated via a geometric mean, making a single low‑quality base heavily penalize the token. In finance, four micro‑structure signals—liquidity, signal‑to‑noise ratio, volatility stability, and order‑flow information—are linearly combined with learnable weights. Both metrics satisfy boundedness, Lipschitz continuity, and monotonic degradation under increased noise.

The merge score derived from the first‑order approximation of the bilevel objective combines pointwise mutual information (PMI) with a quality term (average token quality raised to α) and a domain‑specific constraint factor ψ(a,b). This score is provably bounded and Lipschitz, ensuring stable learning.

Empirical evaluation spans three settings. On simulated 150‑bp paired‑end reads (30× coverage, doubled error rates) and the GIAB HG002 truth set, QA‑BPE‑seq achieves an F1 of 0.891 in variant calling, a 6.7‑percentage‑point gain over standard BPE (0.824). It also reduces reconstruction time by ~24 % and improves taxonomic classification. Ablation studies show that removing RL, quality rewards, information rewards, or adaptive parameters each degrades performance by 2–7 pp, confirming the contribution of each component. In a high‑frequency finance benchmark, QA‑Token‑enhanced models improve Sharpe ratios by 30 % relative to baselines, while maintaining lower token‑level volatility.

At foundation‑model scale, the authors re‑tokenize a 1.7 trillion‑base‑pair corpus (METAGEN‑1) using QA‑Token, reducing total token count by 15 % without sacrificing information. A 7‑billion‑parameter transformer trained on this tokenization attains a pathogen‑detection Matthews Correlation Coefficient of 0.9453, surpassing prior state‑of‑the‑art.

The paper’s contributions are significant: it provides a mathematically rigorous, quality‑aware tokenization method that can unlock massive noisy corpora for foundation‑model pre‑training, all while incurring zero inference overhead. Limitations include reliance on accurate quality annotations (e.g., Phred scores) and computational overhead during vocabulary construction, which may require further optimization for real‑time streaming scenarios. Overall, QA‑Token represents a compelling step toward more robust, data‑efficient language models in scientific and financial domains.