Kl Reduction

---
name: kl-reduction
description: 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. Use when the user has predictions and gold and needs to compute KL reduction.
metadata:
  skill_kind: metric
  source_arxiv: 2302.03169
  bibtex_key: xie2023data
  confidence: high
---

# 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

```python
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

```bibtex
@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