Kl Reduction

kl-reduction

Data Selection for Language Models via Importance Resampling — Xie et al. (2023) (arXiv:2302.03169, 2023)

What this evaluates

Probes the statistical alignment between a candidate pretraining dataset and a target reference distribution (e.g., The Pile or Wikipedia/books). It quantifies how well the dataset's hashed n-gram frequencies match the desired language model pretraining distribution, serving as a proxy for downstream pretraining performance.

Datasets

• The Pile — total ?; splits: validation (-1)
• RealNews — total ?; splits: (unstated)
• S2ORC — total ?; splits: (unstated)
• Amazon reviews — total ?; splits: (unstated)

Metrics

• KL reduction (primary) — range: other
  • Estimates the hashed n-gram distribution from the first 100k examples of a dataset by normalizing counts to an MLE distribution. Computes the KL divergence between this dataset distribution and a reference distribution (e.g., The Pile validation set or Wikipedia/books). The metric reports the reduction in divergence, indicating how closely the dataset matches the target pretraining distribution.

Input / output format

Input: Raw text documents from a candidate dataset (specifically the first 100,000 examples).

Output: A single scalar value representing the KL reduction score.

Scoring recipe

def compute_kl_reduction(dataset_samples, reference_dist):
    # Extract first 100k examples
    samples = dataset_samples[:100000]
    # Compute hashed n-gram counts and normalize to MLE distribution
    p_dataset = normalize(compute_hashed_ngrams(samples))
    # Reference distribution (e.g., from Pile val or Wikipedia/books)
    p_ref = reference_dist
    # Compute KL divergence and return reduction metric
    kl_div = kl_divergence(p_ref, p_dataset)
    return 1.0 - kl_div  # or raw KL depending on convention

Common pitfalls

• Using more or fewer than 100k examples for the n-gram distribution estimation, which changes the MLE estimate.
• Failing to normalize raw n-gram counts into a proper probability distribution before computing KL divergence.
• Using the wrong reference distribution (e.g., The Pile instead of Wikipedia/books) when evaluating RoBERTa pretraining.

Evidence (verbatim from paper)

To compute the KL reduction metric for a particular dataset, we took the first 100k examples from the dataset and computed the hashed n-gram counts. Normalizing these counts gives an MLE estimate of the hashed n-gram distribution for the dataset. We use the same procedure to compute the hashed n-gram distribution parameters for The Pile (from the Pile validation set).

Citation

@misc{xie2023data,
  title={Data Selection for Language Models via Importance Resampling},
  author={Xie et al. (2023)},
  year={2023},
  note={arXiv:2302.03169}
}

• arXiv: 2302.03169