04_improve_minipile.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Improve the MiniPile Pipeline\n",
    "\n",
    "**Objectives:**\n",
    "- [x] Improve the dataset creation process, create new SuperMiniPile dataset (ideally smaller and more information-retaining)\n",
    "- [x] Evaluate Pythia $160\\text{M}$ on SuperMiniPile on MMLU and ARC-Challenge\n",
    "- [x] Evaluate and compare performances of Pythia $1.4\\text{B}$ Pile Deduplicated vs. trained on \"SuperMiniPile\" on the MMLU and ARC benchmarks\n",
    "\n",
    "All produced models and datasets have been published on HuggingFace.\n",
    "\n",
    "Each trained model also holds a model card comparing its performance to the model trained on the original MiniPile.<br>\n",
    "Comprehensive overviews on benchmark results can be found [here](./img/benchmark_results.pdf)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "The results and implementation specifics from the reproduction directly affect the investigation for improvements.<br>\n",
    "**The aim is to investigate different ideas for creating a new SuperMiniPile dataset that is ideally smaller and/or more information-retaining.**\n",
    "\n",
    "**We apply a specific restriction to the ideas below:**<br>\n",
    "Ideas aim to hit the mark of $1,010,500$ examples total.<br>\n",
    "Different sampling techniques for MiniPile may result in datasets with varying memory footprints, even when the number of examples remains the same or very close to that of the original MiniPile.<br>\n",
    "However, for our research purposes, we've decided to maintain the original step count as derived and scaled for MiniPile for as long as the dataset example count is identical or very close to that of said original MiniPile.<br>\n",
    "By keeping this training parameter consistent, we can more directly compare the step-wise effectiveness of our new ideas.<br>\n",
    "In that sense, we expect that a better sampling method will enable the model to produce improved benchmark results under these identical, be it maybe slightly stricter or slightly more relaxed settings.\n",
    "\n",
    "### Idea 1 - Cluster-Proportionate Sampling\n",
    "\n",
    "The original MiniPile dataset was created by sampling *equal* amounts of documents from each of the non-excluded clusters. This results in a MiniPile that cannot represent the original dataset's cluster distribution anymore, but rather imposes a uniform distribution across the clusters, no matter their size or importance/'weight'.\n",
    "\n",
    "For a first improvement attempt, named 'Proportionate', we keep as close to the reproduction code as possible. But, instead of sampling equal amounts of documents from each remaining cluster, we sample a proportionate amount of documents based on cluster sizes (by document count). This requires to make the amount of data points an upper bound rather than a fixed requirement, as we may not be able to sample the exact amount of documents from each cluster by their size. We just don't want to go over the MiniPile-document count.\n",
    "\n",
    "As a side note, the original script was observed as being relatively memory-demanding and, moreover, cache-guzzling.<br>\n",
    "The parquet data-retrieval process was to blame for that, and got thoroughly improved, ditching efforts to wrestle with Pandas (the library), replacing it with Numpy and using explicit cache management, which in turn lifted needs for deep copying.\n",
    "\n",
    "This idea's distillation script is implemented in `03_distill_pile_embed_idea_1_proportionate.py`.<br>\n",
    "Based on thus assembled $1,010,409$ documents, a Pythia $160\\text{M}$ was trained for evaluation.\n",
    "\n",
    "This idea's dataset can be found here: [https://huggingface.co/datasets/Marcus2112/minipile_cluster-proportioned](https://huggingface.co/datasets/Marcus2112/minipile_cluster-proportioned)<br>\n",
    "This idea's trained model can be found here: [https://huggingface.co/Marcus2112/pythia-160m-minipile_cluster-proportioned](https://huggingface.co/Marcus2112/pythia-160m-minipile_cluster-proportioned)\n",
    "\n",
    "The benchmark results are as follows, compared to \"$160\\text{M}$ Pile\", \"$160\\text{M}$ MiniPile\" and \"$160\\text{M}$ MiniPile Reproduction\":\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M Pile Deduplicated | 160M MiniPile               | 160M Reproduction           | 160M Cluster-Proportionate   |\n",
    "| ---------------- | ---------- | --- | ---------------------- | --------------------------- | --------------------------- | ---------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | 0.1997 ± 0.0117        | **0.2125 ± 0.0120**         | 0.1894 ± 0.0115             | 0.1928 ± 0.0115              |\n",
    "| MMLU             | acc        | ↑   | 0.2299 ± 0.0035        | **0.2699 ± 0.0037**         | 0.2295 ± 0.0035             | 0.2295 ± 0.0035              |\n",
    "| HellaSwag        | acc        | ↑   | **0.2903 ± 0.0045**    | 0.2560 ± 0.0044             | 0.2604 ± 0.0044             | 0.2613 ± 0.0044              |\n",
    "| WinoGrande       | acc        | ↑   | 0.4964 ± 0.0141        | 0.4720 ± 0.0140             | **0.5122 ± 0.0140**         | 0.5051 ± 0.0141              |\n",
    "| Lambada (OpenAI) | acc        | ↑   | **0.3689 ± 0.0067**    | 0.0000 ± 0.0000             | 0.0000 ± 0.0000             | 0.0000 ± 0.0000              |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | **31.2589 ± 1.1594**   | 3033175.2693 ± 288926.5827  | 1854408.3999 ± 148101.5978  | 2214257.4651 ± 184064.6008   |\n",
    "| Lambada (Std)    | acc        | ↑   | **0.2335 ± 0.0059**    | 0.0000 ± 0.0000             | 0.0000 ± 0.0000             | 0.0000 ± 0.0000              |\n",
    "| Lambada (Std)    | perplexity | ↓   | **172.7619 ± 7.7265**  | 27067951.3460 ± 2710040.191 | 11927123.2514 ± 1063672.928 | 15143084.5983 ± 1387627.8650 |\n",
    "| BLiMP            | acc        | ↑   | **0.7294 ± 0.0015**    | 0.5194 ± 0.0018             | 0.5481 ± 0.0017             | 0.5452 ± 0.0017              |\n",
    "\n",
    "The proportional sampling approach yields a dataset that is $\\sim 300$ MB smaller than the Reproduction dataset, while also containing $91$ entries less.<br>\n",
    "At the same time, proportional sampling yields close to equal results in all but the perplexity benchmarks, where it underperforms compared to the Reproduction dataset, yet still beats the original MiniPile by a large margin.\n",
    "\n",
    "> We seem to have lost some highly informative examples from smaller clusters when scaling down their representation to match cluster proportions.\n",
    "\n",
    "*The proportional sampling approach is not a clear improvement over the reproduction*, but rather a compromise between the original MiniPile and the Reproduction dataset (tending towards the latter in results though). This is interesting, as it implies the uniform sampling approach may have been more effective in capturing the dataset's most speaking examples than the proportional sampling approach. I do not think that is a fault of the proportional sampling approach, but rather a symptom arising from the original dataset's composition. \n",
    "\n",
    "> Even though we incorporate the original dataset topology more directly, this doesn't necessarily come along with a better focus on the most informative examples.\n",
    "\n",
    "Proportional sampling is neither a failure nor a universal solution, and neither is uniform sampling, but the dismissal of either approach would be premature.<br>\n",
    "Instead, sampling styles should be context-dependent.<br>\n",
    "I assume that when performance on specific downstream tasks, such as language modeling, is stated as the primary goal, incorporating adaptive weighting mechanisms that emphasize the contribution of critical clusters might yield superior results.\n",
    "\n",
    "I conclude that, while itself not a clear improvement, the proportional sampling idea is valid for e.g. approaches where it is important to capture the original dataset's cluster distribution more closely. If the main goal is retaining performance, however, future ideas should investigate how to find a content-based weighting factor per cluster.\n",
    "\n",
    "### Idea 2 - Hybrid Loss-based Sampling\n",
    "\n",
    "In order to sample documents considered most representative and informative, the original MiniPile uses a one-shot proxy-based geometric sampling strategy. After all, the pipeline doesn't really 'select' documents by their content, but by their embedding's position in the cluster space, relative to other documents by proxy of the centroid.<br>\n",
    "Beyond that, once clusters have been determined and selected, the pipeline samples randomly across each cluster.<br>\n",
    "The comparison of different subset assembly techniques performed by [(Guo et al. 2022)](https://arxiv.org/abs/2204.08499) concludes that random sampling can be considered a very robust baseline for custom subset selection efforts, adaptable to various tasks.\n",
    "\n",
    "I deduct that employing cluster-wise random sampling, as performed already, while effective, fast and versatile, can not explicitly consider the point-by-point actual degree of informativeness. To an extent, we rely on luck and sampling spread for finding a best MiniPile distillate.\n",
    "\n",
    "Instead, what if we could utilize the heavy lifting by the embedding and clustering steps, while adding into the process an instance for information-based guiding of cluster sample factor weighting and specific document selection, which could thus address the findings of Idea $1$ (Cluster-Proportionate Sampling)?\n",
    "\n",
    "**Idea $2$, called Lossi (Loss-informed Sampling)**, utilizes much of the existing pipeline for time and efficiency sake.\n",
    "\n",
    "Lossi as a whole consists of several adaptations.<br>\n",
    "The main idea is to use a small proxy model to determine the informativeness of each cluster and then sample documents from each cluster proportionate to their informativeness. We can do this at two points during dataset assembly. \n",
    "\n",
    "What do I mean by that?\n",
    "\n",
    "- Per cluster: Uniformly sample $n$ (e.g. $1,000$) documents and determine their loss with a small Pythia $70\\text{M}$ proxy model\n",
    "- Use the mean loss as a heuristic for the cluster's informativeness and weight the cluster's representation in the final dataset by this value\n",
    "\n",
    "Note that we select the smallest Pythia model trained half-way through the Pile, which assumes that the Pile is shuffled with regards to cluster assignments. I distinctly chose to use a small proxy model for this. MiniPile was intended for use in constrained academic settings, so we have to make do with small models that are most universally trainable, while not being too small as not to allow for their proxy use. Same goes for why the $70\\text{M}$ model is only trained to half of the full Pile dataset.\n",
    "\n",
    "The distillation script is implemented in `03_distill_pile_embed_idea_2_lossi_1.py` and `03_distill_pile_embed_idea_2_lossi_2.py`.<br>\n",
    "Based on the resulting dataset, a Pythia $160\\text{M}$ was trained for evaluation.\n",
    "\n",
    "The dataset can be found here: [https://huggingface.co/datasets/Marcus2112/minipile_loss-sampled](https://huggingface.co/datasets/Marcus2112/minipile_loss-sampled)<br>\n",
    "This idea's trained model can be found here: [https://huggingface.co/Marcus2112/pythia-160m-minipile_loss-sampled](https://huggingface.co/Marcus2112/pythia-160m-minipile_loss-sampled)\n",
    "\n",
    "The benchmark results are as follows, compared to the original MiniPile and the Reproduction:\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M MiniPile                | 160M Lossi                       | Percentage Difference of Means | 95% Confidence Interval         | Interpretation             |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | -------------------------------- | ------------------------------ | ------------------------------- | -------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.2125 ± 0.0120**          | 0.1980 ± 0.0116                  | -6.8235                        | (0.0182; -0.0472)               | Different not significant  |\n",
    "| MMLU             | acc        | ↑   | **0.2699 ± 0.0037**          | 0.2295 ± 0.0035                  | -14.9685                       | (-0.0304; -0.0504)              | MiniPile-trained better    |\n",
    "| HellaSwag        | acc        | ↑   | 0.2560 ± 0.0044              | **0.2599 ± 0.0044**              | 1.5234                         | (0.0161; -0.0083)               | Different not significant  |\n",
    "| WinoGrande       | acc        | ↑   | 0.4720 ± 0.0140              | **0.5107 ± 0.0140**              | 8.1992                         | (0.0775; -0.0001)               | Difference not significant |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                               | -                          |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 3033175.2693 ± 288926.5827   | **2116445.1732 ± 175403.0579**   | -30.2234                       | (-254247.7681; -1579212.4241)   | Lossi 1 severely better    |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                               | -                          |\n",
    "| Lambada (Std)    | perplexity | ↓   | 27067951.3461 ± 2710040.1910 | **14896599.9251 ± 1366937.5470** | -44.9659                       | (-6222231.2223; -18120471.6197) | Lossi 1 severely better    |\n",
    "| BLiMP            | acc        | ↑   | 0.5194 ± 0.0018              | **0.5492 ± 0.0017**              | 5.7374                         | (0.0347; 0.0249)                | Lossi 1 better             |\n",
    "\n",
    "<br>\n",
    "<br>\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M Reproduction             | 160M Lossi                   | Percentage Difference of Means | 95% Confidence Interval      | Interpretation             |\n",
    "| ---------------- | ---------- | --- | -------------------------------- | ---------------------------- | ------------------------------ | ---------------------------- | -------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | 0.1894 ± 0.0115                  | **0.1980 ± 0.0116**          | 4.5407                         | (0.0406; -0.0234)            | Difference not significant |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035                  | 0.2295 ± 0.0035              | 0.0000                         | (0.0097; -0.0097)            | Difference not significant |\n",
    "| HellaSwag        | acc        | ↑   | **0.2604 ± 0.0044**              | 0.2599 ± 0.0044              | -0.1920                        | (0.0117; -0.0127)            | Difference not significant |\n",
    "| WinoGrande       | acc        | ↑   | **0.5122 ± 0.0140**              | 0.5107 ± 0.0140              | -0.2929                        | (0.0373; -0.0403)            | Difference not significant |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000                  | 0.0000 ± 0.0000              | -                              | -                            | -                          |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | **1854408.3999 ± 148101.5978**   | 2116445.1732 ± 175403.0579   | 14.1305                        | (711985.1414; -187911.5948)  | Difference not significant |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000                  | 0.0000 ± 0.0000              | -                              | -                            | -                          |\n",
    "| Lambada (Std)    | perplexity | ↓   | **11927123.2514 ± 1063672.9280** | 14896599.9251 ± 1366937.5470 | 24.8968                        | (6364250.0618; -425296.7144) | Difference not significant |\n",
    "| BLiMP            | acc        | ↑   | 0.5481 ± 0.0017                  | **0.5492 ± 0.0017**          | 0.2007                         | (0.0058; -0.0036)            | Difference not significant |\n",
    "\n",
    "A clear improvement over the original MiniPile can be observed. However, this can't be attributed to the loss-informed sampling approach, as the Reproduction dataset is remarkably equal in performance. This implies that the measures taken unwittingly during the reproduction process pose a more effective set of improvements than the loss-informed sampling approach itself additionally provides, if it does at all. These measures include:\n",
    "\n",
    "- The embedding process, which during the Reproduction uses a slightly worse model than the original MiniPile, but processes larger context sizes for the embeddings ($1024$ instead of $512$)\n",
    "- The clustering process, which may induce different cluster formations based on the different embeddings\n",
    "- The interpretability inherent in the cluster exclusion process, which may have led to a more effective dataset\n",
    "\n",
    "This can imply that changes in the embedding and clustering process may be more effective than the loss-informed sampling approach itself.<br>\n",
    "The proxy model may not be able to capture the full complexity of the dataset, and thus the loss-informed sampling approach may not be able to capture the most informative examples. This in turn would put the loss-informed sampling approach at a disadvantage compared to the uniform sampling approach due to decreased efficiency.\n",
    "\n",
    "### Idea 3 - Density-Proportionate Cluster Sampling\n",
    "\n",
    "An inherent danger with loss-based informativeness approximation is the potential for skewing the dataset towards harder examples, which could make the dataset less representative of general tasks. While loss can be considered a related indicator, it is not a direct measure of informativeness.\n",
    "\n",
    "Taking a step back from this insight, we should re-focus that MiniPile's goal is to capture the most representative subset of the 'insights' gainable from the original Pile dataset. To achieve that, we can focus diversity: Given a dataset represented by embeddings, the sum of all embeddings is scattered across some subspace of that embedding space. Focusing diversity in this plastic example would meaan to implement a way of sampling that aims to recover a wide-spread subset from within the dataset to still cover a large part of the original dataset's embedding subspace. If the embeddings are really representative, which we'll assume for the sake of this argument, sampling with an enforced embeddings spread enforces a selection of more diverse, distinct, characteristically informative data points. This could improve generalization, especially when requiring such broad informational coverage.\n",
    "\n",
    "Specifically, sampling with a factor that is inversely proportional to the cluster density will prioritize sparse regions of the dataset and thus the gaining of a broad coverage of the dataset's information. However, as lunch can't be just free like this, over-sampling/representing sparse clusters could introduce noise, especially when cluster sparsity may correlate with low-quality examples. One could argue that we already excluded the most uninteresting, thus noise-inducing clusters, but the potential for noise itself is in no way mitigated (/mitigatable at this point?). Density-based sampling therefore has to be understood as a technique incorporating a trade-off between spread for representation, avoidance of accidental capturing of noise and capturing of dense regions, as their contents may still be important.\n",
    "\n",
    "And still, even with this notion of having weight density against a notion of informational richness potentially findable within dense areas, density-based sampling could be helpful in capturing a most diverse set of examples, which in turn could be beneficial for generalization.\n",
    "\n",
    "We have to find a feasible way to mitigate noise and over-representation at least to some extent.<br>\n",
    "I therefore propose a density-based sampling approach that calculates cluster contribution proportions like so:\n",
    "$$\\text{Cluster Proportion}_i = \\frac{|C_i|}{|\\bigcup_{j} C_j|} \\cdot (1 - \\omega \\cdot \\rho(C_i))$$\n",
    "\n",
    "where $|C_i|$ is the number of documents in cluster $i$, $|\\bigcup_{j} C_j|$ is the total number of documents in all non-excluded clusters, and $\\rho(C_i)$ is the density of cluster $i$. The impact of the density is scaled by the hyperparameter $\\omega$, making the factor of over-representation of thoroughly sparse clusters controllable.<br>\n",
    "I set $\\omega = 0.5$ with the intention of having density not overly govern the proportion calculation and as to still allow for size-proportionate sampling of clusters and to not overly ignore dense clusters, which could still contain important and therefore often repeated information. But, effectively, cluster-wise embedding density is now a factor in the cluster contribution proportion calculation.\n",
    "\n",
    "The dataset can be found here: [https://huggingface.co/datasets/Marcus2112/minipile_density-proportioned](https://huggingface.co/datasets/Marcus2112/minipile_density-proportioned)<br>\n",
    "The trained model can be found here: [https://huggingface.co/Marcus2112/pythia-160m-minipile_density-proportioned](https://huggingface.co/Marcus2112/pythia-160m-minipile_density-proportioned)\n",
    "\n",
    "These results emerge when comparing benchmarks against Pythia \"$160\\text{M}$ MiniPile\" and Pythia \"$160\\text{M}$ MiniPile Reproduction\":\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M MiniPile                | 160M Density                     | Percentage Difference of Means | 95% Confidence Interval         | Interpretation              |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | -------------------------------- | ------------------------------ | ------------------------------- | --------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.2125 ± 0.0120**          | 0.1920 ± 0.0115                  | -9.6471                        | (0.0121; -0.531)                | Difference not significant  |\n",
    "| MMLU             | acc        | ↑   | **0.2699 ± 0.0037**          | 0.2295 ± 0.0035                  | -14.9685                       | (-0.0304; -0.0504)              | MiniPile better             |\n",
    "| HellaSwag        | acc        | ↑   | 0.2560 ± 0.0044              | **0.2604 ± 0.0044**              | 1.7188                         | (0.0166; -0.0078)               | Difference not significant  |\n",
    "| WinoGrande       | acc        | ↑   | 0.4720 ± 0.0140              | **0.5201 ± 0.0140**              | 10.1907                        | (0.0869; 0.0093)                | **Density better**          |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                               | -                           |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 3033175.2693 ± 288926.5827   | **2099002.0912 ± 170652.6222**   | -30.7985                       | (-276474.4857; -1591871.8705)   | **Density severely better** |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                               | -                           |\n",
    "| Lambada (Std)    | perplexity | ↓   | 27067951.3461 ± 2710040.1910 | **13347273.6076 ± 1997894.6360** | -50.6898                       | (-7121587.1522; -20319768.3248) | **Density severely better** |\n",
    "| BLiMP            | acc        | ↑   | 0.5194 ± 0.0018              | **0.5501 ± 0.0017**              | 5.9107                         | (0.0356; 0.0258)                | **Density better**          |\n",
    "\n",
    "<br>\n",
    "<br>\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M Reproduction                | 160M Density                 | Percentage Difference of Means | 95% Confidence Interval       | Interpretation             |\n",
    "| ---------------- | ---------- | --- | -------------------------------- | ---------------------------- | ------------------------------ | ----------------------------- | -------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | 0.1894 ± 0.0115                  | **0.1920 ± 0.0115**          | 1.3728                         | (0.0345; -0.0293)             | Difference not significant |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035                  | 0.2295 ± 0.0035              | 0.0000                         | (0.0097; -0.0097)             | Difference not significant |\n",
    "| HellaSwag        | acc        | ↑   | 0.2604 ± 0.0044                  | 0.2604 ± 0.0044              | 0.0000                         | (0.0122; -0.0122)             | Difference not significant |\n",
    "| WinoGrande       | acc        | ↑   | 0.5122 ± 0.0140                  | **0.5201 ± 0.0140**          | 1.5424                         | (0.0467; -0.0309)             | Difference not significant |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000                  | 0.0000 ± 0.0000              | -                              | -                             | -                          |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | **1854408.3999 ± 148101.5978**   | 2099002.0912 ± 170652.6222   | 13.1899                        | (687468.6952; -198281.3126)   | Difference not significant |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000                  | 0.0000 ± 0.0000              | -                              | -                             | -                          |\n",
    "| Lambada (Std)    | perplexity | ↓   | **11927123.2514 ± 1063672.9280** | 13347273.6076 ± 1997894.6360 | 11.9069                        | (5856415.8784; -3016115.1660) | Difference not significant |\n",
    "| BLiMP            | acc        | ↑   | 0.5481 ± 0.0017                  | **0.5501 ± 0.0017**          | 0.3649                         | (0.0067; -0.0027)             | Difference not significant |\n",
    "\n",
    "Other than the loss-based approach and except for the perplexity benchmarks, the density-based sampling approach is on par or marginally better than the Reproduction and thus noticably better than the original MiniPile. This does not imply that density-based sampling is an improvement, but in turn, it would signal that the Reproduction pipeline itself, as was the case with the loss-based sampling approach, is the largest contributor to the improvements seen in the benchmarks. All be it statistically insignificant, the density-based sampling approach brings some slight improvement over the Reproduction in some areas.\n",
    "\n",
    "**However**, contrary to these slightest improvements or equal results seen in the benchmarks, I conclude that this approach is a real improvement.<br>\n",
    "This is because while benchmarks show near equal results between the Reproduction and the Density-based approach, the datasets themselves are not of the same size:\n",
    "- MiniPile Reproduction: $1,010,500$ documents at $3.76$ GB,\n",
    "- MiniPile Density-Proportioned: $946,465$ documents at $3.25$ GB.\n",
    "\n",
    "(I actually calculated whether this deviation in example counts would imply a need for change of the step count. I found this to not be the case, with the calculated step count for the Density-Proportioned dataset being $\\sim 1,024$ still.)\n",
    "\n",
    "The dataset contains $64,035$ documents less ($\\sim 6.34\\%$), which amount to $0.51$ GB ($\\sim 13.56\\%$) less of a footprint.<br>\n",
    "(For the training split specifically, this means a reduction from $1,000,000$ to $936,630$, meaning $63,370$ documents less, or $\\sim 6.34\\%$)\n",
    "Still, it produces ever so slightly better or equal results (except for perplexity) compared to the Reproduction and much better results (except on MMLU) than the original MiniPile. Therefore, I conclude that applying weighted density-based sampling to The Pile to indeed be a valid improvement over the Reproduction and the original distilled subset.\n",
    "\n",
    "While we see improvements when adding the context of the size reduction, these observations are not necessarily conclusive. Therefore, further ablations and ideas were explored.\n",
    "\n",
    "### Idea 3.1 - Changing $\\omega$\n",
    "\n",
    "The above described density-based sampling approach uses a hyperparameter $\\omega$ to scale the impact of the density on the cluster proportion calculation like so:\n",
    "$$\\text{Cluster Proportion}_i = \\frac{|C_i|}{|\\bigcup_{j} C_j|} \\cdot (1 - \\omega \\cdot \\rho(C_i))$$\n",
    "\n",
    "In the above attempt, $\\omega$ was set to $0.5$ with an intent of reducing the weight/impact of the density. However, this value was chosen arbitrarily and may not be optimal. Thus, the impact of choosing a different value for $\\omega$ should be investigated for potential improvements and further insights.\n",
    "\n",
    "The higher we set $\\omega$, the more we increase the impact of the density on the cluster proportion calculation. In theory, this could lead to a more diverse dataset, but also to a higher risk of over-representing sparse clusters and thus introducing noise.\n",
    "\n",
    "We've seen the results of cluster-proportionate sampling, i.e. $\\omega = 0$, and density-based sampling, i.e. $\\omega = 0.5$. Now, we will investigate the impact of setting $\\omega = 0.75$ on the dataset creation process. By increasing $\\omega$, the above formula emphasizes sampling from clusters in lower-density regions more strongly. With a higher $\\omega$, the term $1 - \\omega \\cdot \\rho(C_i)$ will reach 0 more easily, thus only leaving room effectively for clusters with lower $\\rho(C_i)$ to be sampled from.\n",
    "\n",
    "The dataset can be found here: [https://huggingface.co/datasets/Marcus2112/minipile_low-density-proportioned](https://huggingface.co/datasets/Marcus2112/minipile_low-density-proportioned)<br>\n",
    "The trained model can be found here: [https://huggingface.co/Marcus2112/pythia-160m-minipile_low-density](https://huggingface.co/Marcus2112/pythia-160m-minipile_low-density)\n",
    "\n",
    "The benchmark results are as follows, where $\\omega = 0.5$ is called 'Density' and $\\omega = 0.75$ is called 'Low Density':\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M Density                     | 160M Low Density             | Percentage Difference of Means | 95% Confidence Interval       | Interpretation             |\n",
    "| ---------------- | ---------- | --- | -------------------------------- | ---------------------------- | ------------------------------ | ----------------------------- | -------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.1920 ± 0.0115**              | 0.1886 ± 0.0114              | -1.7708                        | (0.0283; -0.0351)             | Difference not significant |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035                  | 0.2295 ± 0.0035              | 0.0000                         | -                             | -                          |\n",
    "| HellaSwag        | acc        | ↑   | **0.2604 ± 0.0044**              | 0.2508 ± 0.0044              | -3.6866                        | (0.0026; -0.0218)             | Difference not significant |\n",
    "| WinoGrande       | acc        | ↑   | **0.5201 ± 0.0140**              | 0.5067 ± 0.0141              | -2.5764                        | (0.0255; -0.0523)             | Difference not significant |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000                  | 0.0000 ± 0.0000              | 0.0000                         | -                             | -                          |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | **2099002.0912 ± 170652.6222**   | 2287598.5548 ± 192724.6151   | 8.8951                         | (693139.8095; -315946.8823)   | Difference not significant |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000                  | 0.0000 ± 0.0000              | 0.0000                         | -                             | -                          |\n",
    "| Lambada (Std)    | perplexity | ↓   | **13347273.6076 ± 1997894.6360** | 16223747.0588 ± 1503858.3054 | 21.5510                        | (7777717.0232; -2024770.1208) | Difference not significant |\n",
    "| BLiMP            | acc        | ↑   | 0.5501 ± 0.0017                  | **0.5504 ± 0.0170**          | 0.0545                         | (0.0338; -0.0332)             | Difference not significant |\n",
    "\n",
    "Apart from BLiMP results, all benchmark report insignificantly lower scores for the 'Low Density' dataset compared to the 'Density' dataset.<br>\n",
    "This implies that overly enforcing sampling from low-density clusters may hurt overall dataset representation and thus generalization abilities. The 'Density' dataset, with $\\omega = 0.5$, is therefore the better choice of these two approaches for the dataset creation process.\n",
    "\n",
    "### Idea 4 - Increasing $k$ for Clustering\n",
    "\n",
    "Given the results of ideas 1 to 3, we can see that the reproduction in itself already is a strong improvement over the original MiniPile. It is such a strong improvement in fact, that the sampling and arrangement ideas seemingly only marginally improve the dataset over the reproduction, if at all. This implies that the clustering process itself should be investigated for potential improvements.\n",
    "\n",
    "Revisiting the K-Means clustering process, it became evident that the number of clusters $k$ chosen by the paper was chosen somewhat arbitrarily. While the deduplicated Pile consists of $220$ data subsets, this doesn't mean that $k = 220$ is the optimal choice for the clustering process, neglecting e.g. the potential for capturing of finer group structures within subsets, which would in turn enable more detailed document selection.\n",
    "\n",
    "To investigate if a higher number of clusters could lead to a more representative dataset, $k$ was increased from $220$ to $440$.\n",
    "We can't however couple the doubling of $k$ with the expectation of doubling the number of clusters to be excluded. Clusters may be excluded only if they fit to the paper's original criteria. By these measures, a total of $70$ clusters were excluded from the dataset. Other parameters were kept unchanged.\n",
    "\n",
    "For dataset distillation, the reproduction code with an increased $k$ value was used.\n",
    "\n",
    "The dataset can be found here: [https://huggingface.co/datasets/Marcus2112/minipile_k440](https://huggingface.co/datasets/Marcus2112/minipile_k440)<br>\n",
    "The trained model can be found here: [https://huggingface.co/Marcus2112/pythia-160m-minipile_k440](https://huggingface.co/Marcus2112/pythia-160m-minipile_k440)\n",
    "\n",
    "The benchmark results are as follows, where $k = 220$ is called 'Reproduction' and $k = 440$ is called 'Increased $k$':\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M Density                 | 160M k440                        | Percentage Difference of Means | 95% Confidence Interval       | Interpretation             |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | -------------------------------- | ------------------------------ | ----------------------------- | -------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | 0.1920 ± 0.0115              | **0.1971 ± 0.0116**              | 2.6563                         | (0.0371; -0.0269)             | Difference not significant |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035              | 0.2295 ± 0.0035                  | -                              | -                             | -                          |\n",
    "| HellaSwag        | acc        | ↑   | 0.2604 ± 0.0044              | **0.2615 ± 0.0044**              | 0.4224                         | (0.0133; -0.0111)             | Difference not significant |\n",
    "| WinoGrande       | acc        | ↑   | 0.5201 ± 0.0140              | **0.5107 ± 0.0140**              | -1.8073                        | (0.0294; -0.0482)             | Difference not significant |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                             | -                          |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 2099002.0912 ± 170652.6222   | **1854900.7910 ± 147593.4812**   | -11.6294                       | (198121.5825; -686324.1829)   | Difference not significant |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                             | -                          |\n",
    "| Lambada (Std)    | perplexity | ↓   | 13347273.6076 ± 1997894.6360 | **11658172.4311 ± 1033012.4141** | -12.6550                       | (2719242.3669; -6097444.7199) | Difference not significant |\n",
    "| BLiMP            | acc        | ↑   | **0.5501 ± 0.0017**          | 0.5466 ± 0.0017                  | -0.6362                        | (0.0012; -0.0082)             | Difference not significant |\n",
    "\n",
    "The density-based sampling with $\\omega=0.5$ on $k=220$ clusters produces non-significantly worse results than random sampling across finer $k=440$ clusters. The results are not significantly different, but the increased $k$ value does not seem to hurt the dataset's generalization abilities.<br>\n",
    "I conclude that the clustering process itself can be improved by increasing the number of clusters, but only ever so slightly. Given this insight and the fact that cluster selection is based on human judgement, therefore more error-/bias-prone and time consuming, the increased $k$ alone should not be considered a preferable, reliable improvement over the original MiniPile.\n",
    "\n",
    "It has to be noted that \"MiniPile Density\", compared to \"MiniPile $\\text{k}440$\", contains $6.33\\%$ less documents at $~7.41\\%$ less disk space.<br>\n",
    "Given that \"MiniPile Density\" requires less human judgement, while providing a smaller dataset at comparable performance to \"$160$M MiniPile $\\text{k}440$\", I conclude that overall \"MiniPile Density\" is the better choice of the two.\n",
    "\n",
    "### Idea 5 - Increasing $k$ for Clustering with Density-based Sampling\n",
    "\n",
    "Density-based sampling and the increased $k$ value for clustering have been combined to investigate if the two approaches could lead to a more representative dataset when combined. The increased $k$ value is set to $440$ and the density-based sampling approach uses $\\omega = 0.5$.\n",
    "\n",
    "The dataset can be found here: [https://huggingface.co/datasets/Marcus2112/minipile_k440_density-proportioned](https://huggingface.co/datasets/Marcus2112/minipile_k440_density-proportioned)<br>\n",
    "The trained model is available here: [https://huggingface.co/Marcus2112/pythia-160m-minipile_k440_density-proportioned](https://huggingface.co/Marcus2112/pythia-160m-minipile_k440_density-proportioned)\n",
    "\n",
    "The benchmark results are as follows, where $k = 220,\\ \\omega=0.5$ is called 'Density' and $k = 440,\\ \\omega = 0.5$ is called 'k440 Density':\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M Density                 | 160M k440 Density                | Percentage Difference of Means | 95% Confidence Interval       | Interpretation             |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | -------------------------------- | ------------------------------ | ----------------------------- | -------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | 0.1920 ± 0.0115              | **0.1928 ± 0.0115**              | 0.4167                         | (0.0327; -0.0311)             | Difference not significant |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035              | 0.2295 ± 0.0035                  | 0.0000                         | -                             | -                          |\n",
    "| HellaSwag        | acc        | ↑   | **0.2604 ± 0.0044**          | 0.2603 ± 0.0044                  | -0.0384                        | (0.0121; -0.0123)             | Difference not significant |\n",
    "| WinoGrande       | acc        | ↑   | **0.5201 ± 0.0140**          | 0.4941 ± 0.0141                  | -4.9990                        | (0.0129; -0.0649)             | Difference not significant |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | 0.0000                         | -                             | -                          |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 2099002.0912 ± 170652.6222   | **2025523.7766 ± 164221.8893**   | -3.5006                        | (390719.6474; -537676.2766)   | Difference not significant |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | 0.0000                         | -                             | -                          |\n",
    "| Lambada (Std)    | perplexity | ↓   | 13347273.6076 ± 1997894.6360 | **12959844.9407 ± 1160155.0647** | -2.9027                        | (4140783.3681; -4915640.7019) | Difference not significant |\n",
    "| BLiMP            | acc        | ↑   | 0.5501 ± 0.0017              | **0.5520 ± 0.0017**              | 0.3454                         | (0.0066; -0.0028)             | Difference not significant |\n",
    "\n",
    "Differences across benchmarks are not significant. The results are about equal to the 'Density' dataset, which implies that the increased $k$ value does not significantly improve the dataset when combined with density-based sampling. Given this result and the fact that higher $k$ implies more human judgement, the increased $k$ value should not be considered a reliable/preferable improvement over the density-based lower $k$ MiniPile, even though this newer dataset is smaller in size (by $\\sim3.46\\%$, equal in memory footprint to the original MiniPile).\n",
    "\n",
    "- if an approach with a minimal human involvement is preferred, choose the 'Density' dataset\n",
    "- if an approach with a more detailed cluster structure is preferred, reducing a little more in size while generally reaching the same performance, choose the 'k440 Density' dataset\n",
    "\n",
    "### Idea 6 - Inter-Intra Cluster Sampling with $k=440$\n",
    "\n",
    "Most of the above approaches consider clusters and their inner workings and topology. This showed to be fruitful, e.g. with the density-based sampling approach. We could extend this idea to also consider (at least primitive) inter-cluster relationships. Specifically, we could involve the following measures:\n",
    "- Cluster density\n",
    "- Cluster size\n",
    "- Inter-cluster diversity (based on cosine distance of centroids)\n",
    "\n",
    "Increasing $k$ to $k=440$ didn't prove to be too effective as such, yet it provides a more detailed cluster landscape, which could be leveraged more effectively through the inter-cluster diversity measure than would be possible for $k=220$ clusters.\n",
    "\n",
    "Adapting the formula from Idea 3, we have to realize cluster sample count calculation like so:\n",
    "$$\\text{Cluster Score}_i = \\omega_s \\cdot \\frac{|C_i|}{|\\bigcup_{j} C_j|} + \\omega_d \\cdot (1 - \\rho(C_i)) + \\omega_v \\cdot \\frac{\\delta(C_i)}{\\max_{j} \\delta(C_j)}$$\n",
    "\n",
    "Here, the additional $\\delta(C_i)$ is the inter-cluster diversity score, and the single $\\omega$ had to evolve into more distinct $\\omega_s$, $\\omega_d$, $\\omega_v$.\n",
    "As for their values, I elected to set all $\\omega$ to $0.33$ to equally consider all three factors.<br>\n",
    "This makes comparing against the density-based sampling approach more clear, because while the original approach gave equal weight ($0.5$) to size and density-adjusted sampling, the new approach maintains this balance by allocating equal thirds ($0.33$) across all factors, effectively preserving the relative influence. My idea is that if we maintain this balance of factors, we could more clearly see the impact of the presence of the inter-cluster diversity measure.\n",
    "\n",
    "The dataset can be found here: [Marcus2112/minipile_k440_inter-density-proportioned](https://huggingface.co/datasets/Marcus2112/minipile_k440_inter-density-proportioned)<br>\n",
    "The trained model can be found here: [Marcus2112/pythia-160m-minipile_k440_inter-density-proportioned](https://huggingface.co/Marcus2112/pythia-160m-minipile_k440_inter-density-proportioned)\n",
    "\n",
    "The benchmark results compared to \"$160\\text{M}$ Density\" are as follows:\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M Density                 | 160M k440 Inter                  | Percentage Difference of Means | 95% Confidence Interval     | Interpretation             |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | -------------------------------- | ------------------------------ | --------------------------- | -------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | 0.1920 ± 0.0115              | **0.1937 ± 0.0115**              | 0.8854                         | 0.0336; -0.0302             | Difference not significant |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035              | 0.2295 ± 0.0035                  | 0.0000                         | 0.0097; -0.0097             | Difference not significant |\n",
    "| HellaSwag        | acc        | ↑   | 0.2604 ± 0.0044              | **0.2608 ± 0.0044**              | 0.1536                         | 0.0126; -0.0118             | Difference not significant |\n",
    "| WinoGrande       | acc        | ↑   | **0.5201 ± 0.0140**          | 0.4996 ± 0.0141                  | -3.9415                        | 0.0184; -0.0594             | Difference not significant |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                           | -                          |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 2099002.0912 ± 170652.6222   | **1858348.2053 ± 147853.1415**   | -11.4652                       | 201902.0871; -683209.8589   | Difference not significant |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                           | -                          |\n",
    "| Lambada (Std)    | perplexity | ↓   | 13347273.6076 ± 1997894.6360 | **11655568.3141 ± 1032438.4286** | -12.6745                       | 2716121.6575; -6099532.2445 | Difference not significant |\n",
    "| BLiMP            | acc        | ↑   | 0.5501 ± 0.0017              | **0.5510 ± 0.0017**              | 0.1636                         | 0.0056; -0.0038             | Difference not significant |\n",
    "\n",
    "The results indiciate that the inter-cluster diversity measure does not significantly improve the dataset over the density-based sampling approach, especially when considering that the density-based sampling dataset is smaller. However, neither could the impact of inter-cluster diversity be conclusively be determined, as the benchmark results only deviate sightly from both \"$160\\text{M}$ Density\" and \"$160\\text{M}$ $\\text{k}440$ Density\". For example, at this point, one cannot rule out that the improvement in the perplexity measure is due to $k=440$ or the inter-cluster diversity measure, or both, as results for \"$160\\text{M}$ $\\text{k}440$ Density\" suggest that the increased $k$ value alone could have a positive impact on the perplexity measure.\n",
    "\n",
    "We can not conclusively determine positive effects of considering the inter-cluster diversity measure, but we can conclude that the dataset is not significantly worse than the density-based sampling approach.\n",
    "\n",
    "### Idea 6.1 - Increasing $\\omega_v$\n",
    "\n",
    "To investigate the importance of inter-cluster diversity further, the impact of increasing $\\omega_v$ should be investigated. We set $\\omega_s = 0.25$, $\\omega_d = 0.25$, and $\\omega_v = 0.5$.\n",
    "\n",
    "The dataset can be found here: [Marcus2112/minipile_k440_high-inter_density](https://huggingface.co/datasets/Marcus2112/minipile_k440_high-inter_density)<br>\n",
    "The trained model can be found here: [Marcus2112/pythia-160m-minipile_k440_high-inter_density-proportioned](https://huggingface.co/Marcus2112/pythia-160m-minipile_k440_high-inter_density-proportioned)\n",
    "\n",
    "Comparing the benchmark results to \"$160\\text{M}$ Density\" and \"$160\\text{M}$ $\\text{k}440$ Inter\", we get:\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M Density                 | 160M k440 Inter High             | Percentage Difference of Means | 95% Confidence Interval     | Interpretation             |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | -------------------------------- | ------------------------------ | --------------------------- | -------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.1920 ± 0.0115**          | 0.1911 ± 00115                   | -0.4688                        | 0.0310; -0.0328             | Difference not significant |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035              | 0.2295 ± 0.0035                  | 0.0000                         | 0.0097; -0.0097             | Difference not significant |\n",
    "| HellaSwag        | acc        | ↑   | 0.2604 ± 0.0044              | **0.2610 ± 0.0044**              | 0.2304                         | 0.0128; -0.0116             | Difference not significant |\n",
    "| WinoGrande       | acc        | ↑   | **0.5201 ± 0.0140**          | 0.5185 ± 0.0140                  | -0.3076                        | 0.0372; -0.0404             | Difference not significant |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                           | -                          |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 2099002.0912 ± 170652.6222   | **1976271.1659 ± 158805.4234**   | -5.8471                        | 334169.7526; -579631.6032   | Difference not significant |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                           | -                          |\n",
    "| Lambada (Std)    | perplexity | ↓   | 13347273.6076 ± 1997894.6360 | **12395759.9269 ± 1104763.2927** | -7.1289                        | 3523164.5637; -5426191.9251 | Difference not significant |\n",
    "| BLiMP            | acc        | ↑   | **0.5501 ± 0.0017**          | 0.5440 ± 0.0017                  | -1.1089                        | -0.0014; -0.0108            | MiniPile Density better    |\n",
    "\n",
    "We see no significant differences, except for BLiMP, where the \"$160\\text{M}$ Density\" dataset is slightly better. This implies that the inter-cluster diversity measure does not significantly improve the dataset over the density-based sampling approach. Further evidence is provided by comparing \"$160\\text{M}$ $\\text{k}440$ Inter\" to \"160M $\\text{k}440$ Inter High\":\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M k440 Inter                  | 160M k440 Inter High         | Percentage Difference of Means | 95% Confidence Interval     | Interpretation             |\n",
    "| ---------------- | ---------- | --- | -------------------------------- | ---------------------------- | ------------------------------ | --------------------------- | -------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.1937 ± 0.0115**              | 0.1911 ± 00115               | -1.3423                        | 0.0293; -0.0345             | Difference not significant |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035                  | 0.2295 ± 0.0035              | 0.0000                         | 0.0097; -0.0097             | Difference not significant |\n",
    "| HellaSwag        | acc        | ↑   | 0.2608 ± 0.0044                  | **0.2610 ± 0.0044**          | 0.0767                         | 0.0124; -0.0120             | Difference not significant |\n",
    "| WinoGrande       | acc        | ↑   | 0.4996 ± 0.0141                  | **0.5185 ± 0.0140**          | 3.7830                         | 0.0578; -0.0200             | Difference not significant |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000                  | 0.0000 ± 0.0000              | -                              | -                           | -                          |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | **1858348.2053 ± 147853.1415**   | 1976271.1659 ± 158805.4234   | 6.3456                         | 543201.0215; -307355.1003   | Difference not significant |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000                  | 0.0000 ± 0.0000              | -                              | -                           | -                          |\n",
    "| Lambada (Std)    | perplexity | ↓   | **11655568.3141 ± 1032438.4286** | 12395759.9269 ± 1104763.2927 | 6.3505                         | 3703897.6525; -2223514.4269 | Difference not significant |\n",
    "| BLiMP            | acc        | ↑   | **0.5510 ± 0.0017**              | 0.5440 ± 0.0017              | -1.2704                        | -0.0023; -0.0117            | 160M k440 Inter better     |\n",
    "\n",
    "The latter, i.e. \"$160\\text{M}$ $\\text{k}440$ Inter High\" is not significantly better. This implies that the inter-cluster diversity measure does not significantly improve the dataset over the density-based sampling approach, even and especially when increasing its impact."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Downsizing\n",
    "\n",
    "At this point, several approaches have been undertaken with the aim of improving performance of Pythia $160\\text{M}$ when trained on MiniPile.<br>\n",
    "The most successful approach was the density-based sampling approach, which showed to be a valid improvement over the original MiniPile and the Reproduction. The increased $k$ value for clustering did not significantly improve the dataset, but it did not hurt either. The inter-cluster diversity measure did not significantly improve the dataset over the density-based sampling approach, but it did not hurt it either.\n",
    "\n",
    "To summarize, key contributing factors for subset creation for MiniPile include:\n",
    "- Selection of excluded clusters based on human judgement\n",
    "- Clustering kept at $k=220$\n",
    "- Cluster-wise density- and size-weighted sampling with always the same $\\omega=0.5$\n",
    "\n",
    "How low can we now go with the dataset example count while still maintaining scores with a majority least on a level with the original MiniPile?<br>\n",
    "In other words, to which extreme can data point sampling be reduced?\n",
    "\n",
    "**Critically, note this:**<br>\n",
    "I now actively change example count and dataset size significantly. Therefore, I will have to adapt the step count for respective training runs accordingly. This is not as trivial as it was in chapter 2, where we calculated the step count as an approximation proportional to the uncompressed memory footprint of the original MiniPile. Arguably, even that was an approximation, but I say this doesn't work at all anymore now, because the produced mini-datasets always have an additional (debugging/sanity-check-) column `pile_idx`, watering the expressiveness of the memory footprint further. Still, to prove that I considered this memory-based step count scaling, I attached `04_step-count_calculator.ipynb`. *But I didn't use that. I considered, but didn't apply. It seemed too watered now.* Instead, I removed the `max_steps` and `save_steps` arguments from the training args alltogether, and let the training run for $1.5$ epochs, which is then automatically proportional to the dataset size. I would start this script, would get informed about the number of steps, would stop the script, adapt the learning rate scheduling with this step count, and would finally start the script again for good.\n",
    "\n",
    "### Density Tiny\n",
    "\n",
    "Once again, comparing Pythia $160\\text{M}$ trained on the original MiniPile againsted a version trained on [minipile_density-proportioned](https://huggingface.co/datasets/Marcus2112/minipile_density-proportioned):\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M MiniPile                | 160M Density                     | Percentage Difference of Means | 95% Confidence Interval         | Interpretation              |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | -------------------------------- | ------------------------------ | ------------------------------- | --------------------------- |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.2125 ± 0.0120**          | 0.1920 ± 0.0115                  | -9.6471                        | (0.0121; -0.531)                | Difference not significant  |\n",
    "| MMLU             | acc        | ↑   | **0.2699 ± 0.0037**          | 0.2295 ± 0.0035                  | -14.9685                       | (-0.0304; -0.0504)              | MiniPile better             |\n",
    "| HellaSwag        | acc        | ↑   | 0.2560 ± 0.0044              | **0.2604 ± 0.0044**              | 1.7188                         | (0.0166; -0.0078)               | Difference not significant  |\n",
    "| WinoGrande       | acc        | ↑   | 0.4720 ± 0.0140              | **0.5201 ± 0.0140**              | 10.1907                        | (0.0869; 0.0093)                | **Density better**          |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                               | -                           |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 3033175.2693 ± 288926.5827   | **2099002.0912 ± 170652.6222**   | -30.7985                       | (-276474.4857; -1591871.8705)   | **Density severely better** |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                  | -                              | -                               | -                           |\n",
    "| Lambada (Std)    | perplexity | ↓   | 27067951.3461 ± 2710040.1910 | **13347273.6076 ± 1997894.6360** | -50.6898                       | (-7121587.1522; -20319768.3248) | **Density severely better** |\n",
    "| BLiMP            | acc        | ↑   | 0.5194 ± 0.0018              | **0.5501 ± 0.0017**              | 5.9107                         | (0.0356; 0.0258)                | **Density better**          |\n",
    "\n",
    "Size of MiniPile: $1,000,000$ training<br>\n",
    "Size of minipile_density-proportioned: $936,630$ training documents\n",
    "\n",
    "In a first attempt, taking inspiration from the improvement on Winogrande, we reduce minipile_density-proportioned by $10\\%$ to $842,967$ documents for training.\n",
    "\n",
    "The dataset can be found here: [Marcus2112/minipile_density-proportioned_tiny](https://huggingface.co/datasets/Marcus2112/minipile_density-proportioned_tiny)<br>\n",
    "The trained model can be found here: [Marcus2112/pythia-160m-minipile_tiny_density-proportioned](https://huggingface.co/Marcus2112/pythia-160m-minipile_tiny_density-proportioned)\n",
    "\n",
    "Comparing the benchmarks against the original MiniPile:\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M MiniPile                | 160M Density Tiny               | Percentage Difference in Means |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | ------------------------------- | ------------------------------ |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.2125 ± 0.0120**          | 0.1843 ± 0.0113                 | -13.2706                       |\n",
    "| MMLU             | acc        | ↑   | **0.2699 ± 0.0037**          | 0.2295 ± 0.0035                 | -14.9685                       |\n",
    "| HellaSwag        | acc        | ↑   | 0.2560 ± 0.0044              | **0.2598 ± 0.0044**             | 1.4844                         |\n",
    "| WinoGrande       | acc        | ↑   | 0.4720 ± 0.0140              | **0.4980 ± 0.0141**             | 5.5085                         |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                 | -                              |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 3033175.2693 ± 288926.5827   | **1934160.4021 ± 153855.8657**  | -36.2331                       |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                 | -                              |\n",
    "| Lambada (Std)    | perplexity | ↓   | 27067951.3461 ± 2710040.1910 | **10354382.8435 ± 900493.0078** | -61.7467                       |\n",
    "| BLiMP            | acc        | ↑   | 0.5194 ± 0.0018              | **0.5363 ± 0.0016**             | 3.2538                         |\n",
    "\n",
    "The downsized dataset is not significantly worse than the original MiniPile, except in the MMLU benchmark, but this effect has been observed with the MiniPile Density dataset as well and to the exact same extent.<br>\n",
    "Stricingly speaking, when comparing the \"$160\\text{M}$ Density Tiny\" to the \"$160\\text{M}$ Density\" performance results, the downsized dataset is not worse at all, showing a noticable drop in performance only on the ARC-Challenge.<br>\n",
    "\n",
    "I conclude that from the $3.25$ GB of \"MiniPile Density\", we were able to shave off another $350$ MB total for a $2.9$ GB dataset almost accomplishing the same.<br>\n",
    "I particularly find the drop in perplexity for the downsized dataset to be quite remarkable, because apparently, the downsized dataset is able to provide a more diverse and informative training signal to the model, even though it is literally the same sampling approach as \"MiniPile Density\".\n",
    "\n",
    "Mind you, we downsized the already downsized \"MiniPile Density\". The $842,967$ elements strong \"MiniPile Density Tiny\" is $15.303\\%$ smaller than the original MiniPile by example count.\n",
    "\n",
    "### Density Nano\n",
    "\n",
    "Let's test this further.<br>\n",
    "We now reduce the dataset size by $25\\%$ compared to the original MiniPile, making for $750,000$ training examples.\n",
    "\n",
    "The dataset can be found here: [Marcus2112/minipile_density-proportioned_nano](https://huggingface.co/datasets/Marcus2112/minipile_density-proportioned_nano)<br>\n",
    "The trained model can be found here: [Marcus2112/pythia-160m-minipile_nano_density-proportioned](https://huggingface.co/Marcus2112/pythia-160m-minipile_nano_density-proportioned)\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M MiniPile                | 160M Density Nano              | Percentage Difference in Means |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | ------------------------------ | ------------------------------ |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.2125 ± 0.0120**          | 0.1928 ± 0.0115                | -9.2706                        |\n",
    "| MMLU             | acc        | ↑   | **0.2699 ± 0.0037**          | 0.2295 ± 0.0035                | -14.9685                       |\n",
    "| HellaSwag        | acc        | ↑   | 0.2560 ± 0.0044              | **0.2596 ± 0.0044**            | 1.4063                         |\n",
    "| WinoGrande       | acc        | ↑   | 0.4720 ± 0.0140              | **0.5036 ± 0.0141**            | 6.6949                         |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                | -                              |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 3033175.2693 ± 288926.5827   | **1871303.2176 ± 150515.6406** | -38.3055                       |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                | -                              |\n",
    "| Lambada (Std)    | perplexity | ↓   | 27067951.3461 ± 2710040.1910 | **10513877.8584 ± 92624.3394** | -61.1575                       |\n",
    "| BLiMP            | acc        | ↑   | 0.5194 ± 0.0018              | **0.5364 ± 0.0016**            | 3.2730                         |\n",
    "\n",
    "The density-based sampling technique affords us to reduce the dataset to $75\\%$ of the original MiniPile size, while nearly still maintaining the same performance level, if not better. The perplexity measure has not significantly worsened, and the WinoGrande benchmark has even improved slightly. ARC-Challenge and MMLU have worsened, with MMLU significantly so. But, I interpret that the overall performance is still largely retained.\n",
    "\n",
    "### Density Pico\n",
    "\n",
    "At this point, \"$160\\text{M}$ Density\", \"$160\\text{M}$ Density Tiny\" and \"$160\\text{M}$ Density Nano\" indicate a trend: A subset smaller or equal to $75\\%$ of the volume of the original MiniPile, is responsible of a majority of the retained performance.<br>\n",
    "This subset seems to be sufficiently captured by the size-density-weighted sampling approach, yet we don't really know which exact data points make up this performance-contributing subset of the distillation. After all, Size-Density-Weighted sampling does just that: It samples data points based on cluster size and density, it does not precisely tell us which data points are the most contributing.<br>\n",
    "However, being able to reduce the dataset size by $25\\%$ of MiniPile, while maintaining largely the same performance level is a strong indicator that size-density-weighted sampling is a useful, and importantly, efficient, approach.\n",
    "\n",
    "To be honest, I didn't expect this trend to occur.<br>\n",
    "I expected benchmark results to crumble hilariously as soon as we'd deviate from the $1\\text{M}$ MiniPile training example count by any more than $10\\%$.<br>\n",
    "But it seems we can afford to go fruther down in example count.<br>\n",
    "Frankly, MiniPile Density Pico with its $250,000$ training examples was built to try and break the observed performance retention trend and to indicate a definite lower bound for the dataset size.\n",
    "\n",
    "The dataset can be found here: [Marcus2112/minipile_density-proportioned_pico](https://huggingface.co/datasets/Marcus2112/minipile_density-proportioned_pico)<br>\n",
    "It is $891$ MB in size, which is $28.38\\%$ of the original MiniPile.\n",
    "The trained model can be found here: [Marcus2112/pythia-160m-minipile_pico_density-proportioned](https://huggingface.co/Marcus2112/pythia-160m-minipile_pico_density-proportioned)\n",
    "\n",
    "Comparing the benchmarks against the original MiniPile, we retain most of the performance except on ARC-Challenge and MMLU and gain significantly on perplexity benchmarks, BLiMP and WinoGrande:\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M MiniPile                | 160M Density Pico               | Percentage Difference in Means |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | ------------------------------- | ------------------------------ |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.2125 ± 0.0120**          | 0.1903 ± 0.0115                 | -10.4471                       |\n",
    "| MMLU             | acc        | ↑   | **0.2699 ± 0.0037**          | 0.2295 ± 0.0035                 | -14.9685                       |\n",
    "| HellaSwag        | acc        | ↑   | 0.2560 ± 0.0044              | **0.2583 ± 0.0044**             | 0.8984                         |\n",
    "| WinoGrande       | acc        | ↑   | 0.4720 ± 0.0140              | **0.4957 ± 0.0141**             | 5.0212                         |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                 | -                              |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 3033175.2693 ± 288926.5827   | **1964196.9255 ± 153419.7846**  | -35.2429                       |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                 | -                              |\n",
    "| Lambada (Std)    | perplexity | ↓   | 27067951.3461 ± 2710040.1910 | **10720344.5522 ± 925236.7038** | -60.3947                       |\n",
    "| BLiMP            | acc        | ↑   | 0.5194 ± 0.0018              | **0.5383 ± 0.0017**             | 3.6388                         |\n",
    "\n",
    "The quarter of a MiniPile dataset is not significantly worse, and in several cases even better.<br>\n",
    "We see that every benchmark that \"160M MiniPile Density\" had improved upon when compared to the 160M original MiniPile\n",
    "is still improved when retained in \"MiniPile Density Pico\".\n",
    "\n",
    "If I hadn't uploaded and viewed the datasets and models, I wouldn't believe this.<br>\n",
    "But it seems, at least for the 160M model size, that the size-density-weighted sampling approach is a highly effective way to distill a dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Increasing the step count\n",
    "\n",
    "We have to rule out that the retaining effects recorded for the downsized datasets are due to the model architecture overfitting on the larger Density dataset. To do this, we increase the step count for the training runs of \"160M Density\" and \"160M Density Pico\" from $1.5$ epochs to $2$ epochs.\n",
    "\n",
    "If the performance of \"160M Density Pico 2 Epochs\" is improved over \"160M Density Pico\", we can conclude that the model architecture is not overfitting on the smaller dataset. If instead the performance is worse, we can conclude that the model architecture starts to overfit on the dataset. If that were the case, it would be the model architecture instead of the dataset structure limiting the performance of the downsized datasets and therefore it would water the expressiveness of the size-density-weighted sampling approach.\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M Density Pico               | 160M Density Pico 2 Epochs  | Percentage Difference in Means |\n",
    "| ---------------- | ---------- | --- | ------------------------------- | --------------------------- | ------------------------------ |\n",
    "| ARC-Challenge    | acc        | ↑   | 0.1903 ± 0.0115                 | **0.1928 ± 0.0115**         | 1.3137                         |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035                 | 0.2295 ± 0.0035             | 0.0000                         |\n",
    "| HellaSwag        | acc        | ↑   | **0.2583 ± 0.0044**             | 0.2573 ± 0.0044             | -0.3871                        |\n",
    "| WinoGrande       | acc        | ↑   | **0.4957 ± 0.0141**             | 0.4933 ± 0.0141             | -0.4842                        |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000                 | 0.0000 ± 0.0000             | -                              |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | **1964196.9255 ± 153419.7846**  | 2017680.7049 ± 159090.0606  | 2.7229                         |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000                 | 0.0000 ± 0.0000             | -                              |\n",
    "| Lambada (Std)    | perplexity | ↓   | 10720344.5522 ± 925236.7038     | **10465698.6884 ± 903166.5195** | -2.3754                    |\n",
    "| BLiMP            | acc        | ↑   | **0.5383 ± 0.0017**             | **0.5405 ± 0.0017**         | 0.4087                         |\n",
    "\n",
    "Most changes in benchmark performance are below $1\\%$, we can observe improvement in the Lambada (Std) perplexity score, yet there is also an equal percentage deterioration on the Lambada (OpenAI) perplexity score.<br>\n",
    "This suggests the model saturation-wise has reached a plateau: The step count changes neither make the model clearly underfit nor overfit. The changes are small enough for me to have to consider them to be noise margins, and more importantly, no systematic change can be observed. But, the stability across most metrics suggests the model has reached its capacity limit for what it can learn from this dataset, and we only are at $488$ steps, compared to the $1024$ steps for a 1.5 epoch traversal through the original MiniPile and MiniPile Density datasets.\n",
    "\n",
    "Our hypothesis at this point could therefore be that the model reached an early saturation at around $488$ steps. Any more than that would not significantly affect the model's performance, or worse, it could lead to overfitting. But is this the case?\n",
    "\n",
    "With MiniPile Density, we now went from $1024$ steps for 1.5 epochs to $1366$ steps:\n",
    "\n",
    "| Benchmark        | Measure    |     | 160M Density                 | 160M Density 2 Epochs          | Percentage Difference in Means |\n",
    "| ---------------- | ---------- | --- | ---------------------------- | ------------------------------ | ------------------------------ |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.1920 ± 0.0115**          | 0.1894 ± 0.0115                | -1.3542                        |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035              | 0.2295 ± 0.0035                | 0.0000                         |\n",
    "| HellaSwag        | acc        | ↑   | **0.2604 ± 0.0044**          | 0.2568 ± 0.0044                | -1.3825                        |\n",
    "| WinoGrande       | acc        | ↑   | **0.5201 ± 0.0140**          | 0.5012 ± 0.0141                | -3.6339                        |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                | -                              |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 2099002.0912 ± 170652.6222   | **1587737.3755 ± 121555.3148** | -24.3575                       |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000              | 0.0000 ± 0.0000                | -                              |\n",
    "| Lambada (Std)    | perplexity | ↓   | 13347273.6076 ± 1997894.6360 | **8366924.7603 ± 713077.3579** | -37.3136                       |\n",
    "| BLiMP            | acc        | ↑   | **0.5501 ± 0.0017**          | 0.5378 ± 0.0017                | -2.2360                        |\n",
    "\n",
    "We do not see the nuanced behavior from \"160M Density Pico 2 Epochs\" reflect here. Instead, we encounter sharp improvements on perplexity scores, but also noticable deteriorations across all accuracy benchmarks except MMLU. As we each worse accuracy with more iterations, we can conclude that the model starts to overfit on the dataset, except for overall exposure, reflected in perplexity scores. The observations further support that the \"160M Density Pico 2 Epochs\" model has reached closer to the model' capacity limit, and that to a noticable extent, we now encounter the effects of stepping over a point of \"early saturation\".\n",
    "\n",
    "However, we can see from the perplexity scores that in some aspects at least, model capacity has not been reached. This makes the perplexity results of \"160M Density Pico\" stand out, as they are lower than the perplexity results of \"160M Density\", suggesting that the downsized dataset is able to provide to some extent a more diverse and informative training signal to the model, even though it is literally the same sampling approach as \"MiniPile Density\". Therefore, the \"160M Density 2 Epochs\" is overfitting but in a particular way: It is still getting better at pure language modeling while losing its sense for reasoning tasks.\n",
    "\n",
    "With these results, we can say that:\n",
    "\n",
    "- The smaller dataset might actually be providing a more \"concentrated\" learning signal, regarding at least the perplexity scores with certainty\n",
    "- The reduction in size might actually be helping prevent the model from overfitting in reasoning tasks\n",
    "- The sampling approach scales differently at different dataset sizes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Evaluate Pythia $1.4\\text{B}$ Pile Deduplicated vs. Pythia $1.4\\text{B}$ \"SuperMiniPile\"\n",
    "\n",
    "In the prior chapter, training was conducted with Pythia $1.4\\text{B}$ on the reproduced MiniPile.<br>\n",
    "This was compared to the Pythia $1.4\\text{B}$ trained on the full Pile Deduplicated dataset.<br>\n",
    "\n",
    "- Pythia $1.4\\text{B}$ MiniPile can be found here: [Marcus2112/pythia-1.4b-minipile](https://huggingface.co/Marcus2112/pythia-1.4b-minipile)\n",
    "- Pythia $1.4\\text{B}$ MiniPile Reproduction can be found here: [Marcus2112/pythia-1.4b-minipile_reproduction](https://huggingface.co/Marcus2112/pythia-1.4b-minipile_reproduction)\n",
    "- Pythia $1.4\\text{B}$ MiniPile Density can be found here: [Marcus2112/pythia-1.4b-minipile_density-proportioned](https://huggingface.co/Marcus2112/pythia-1.4b-minipile_density-proportioned)\n",
    "\n",
    "\n",
    "Further training was conducted to establish an additional baseline with training Pythia $1.4\\text{B}$ on the original MiniPile.<br>\n",
    "This was first compared to the Pythia $1.4\\text{B}$ trained on the full Pile Deduplicated dataset:\n",
    "\n",
    "| Benchmark        | Measure    |     | 1.4B Pile Deduplicated | 1.4B MiniPile              | Percentage Difference in Means |\n",
    "| ---------------- | ---------- | --- | ---------------------- | -------------------------- | ------------------------------ |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.2600 ± 0.0130**    | 0.1903 ± 0.0115            | -26.8077                       |\n",
    "| MMLU             | acc        | ↑   | **0.2388 ± 0.0036**    | 0.2295 ± 0.0035            | -3.8945                        |\n",
    "| HellaSwag        | acc        | ↑   | **0.4177 ± 0.0049**    | 0.2579 ± 0.0044            | -38.2571                       |\n",
    "| WinoGrande       | acc        | ↑   | **0.5730 ± 0.0140**    | 0.5185 ± 0.0140            | -9.5133                        |\n",
    "| Lambada (OpenAI) | acc        | ↑   | **0.6202 ± 0.0068**    | 0.0000 ± 0.0000            | -100.0000                      |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | **6.1041 ± 0.1531**    | 1564928.5258 ± 118691.4565 | 25637234.3458                  |\n",
    "| Lambada (Std)    | acc        | ↑   | **0.4898 ± 0.0070**    | 0.0000 ± 0.0000            | -100.0000                      |\n",
    "| Lambada (Std)    | perplexity | ↓   | **11.2448 ± 0.3305**   | 8848600.9409 ± 745031.8900 | 78690503.1312                  |\n",
    "| BLiMP            | acc        | ↑   | **0.8154 ± 0.0013**    | 0.5483 ± 0.0017            | -32.7569                       |\n",
    "| ARC-Easy         | acc        | ↑   | **0.6174 ± 0.0100**    | 0.2715 ± 0.0091            | -56.0253                       |\n",
    "\n",
    "The Pythia $1.4\\text{B}$ trained on the original MiniPile was then compared to the Pythia $1.4\\text{B}$ trained on the reproduced MiniPile:\n",
    "\n",
    "| Benchmark        | Measure    |     | 1.4B MiniPile              | 1.4B Reproduction              | Percentage Difference in Means |\n",
    "| ---------------- | ---------- | --- | -------------------------- | ------------------------------ | ------------------------------ |\n",
    "| ARC-Challenge    | acc        | ↑   | 0.1903 ± 0.0115            | **0.1928 ± 0.0115**            | 1.3137                         |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035            | 0.2295 ± 0.0035                | 0.0000                         |\n",
    "| HellaSwag        | acc        | ↑   | 0.2579 ± 0.0044            | **0.2584 ± 0.0044**            | 0.1939                         |\n",
    "| WinoGrande       | acc        | ↑   | **0.5185 ± 0.0140**        | 0.5091 ± 0.0141                | -1.8129                        |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000            | 0.0000 ± 0.0000                | -                              |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 1564928.5258 ± 118691.4565 | **1520707.8702 ± 115261.3664** | -2.8257                        |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000            | 0.0000 ± 0.0000                | -                              |\n",
    "| Lambada (Std)    | perplexity | ↓   | 8848600.9409 ± 745031.8900 | **8651201.8876 ± 735161.5236** | -1.5685                        |\n",
    "| BLiMP            | acc        | ↑   | **0.5483 ± 0.0017**        | 0.5397 ± 0.0016                | -2.2309                        |\n",
    "| ARC-Easy         | acc        | ↑   | **0.2715 ± 0.0091**        | 0.2673 ± 0.0091                | -1.5470                        |\n",
    "\n",
    "Now, the \"SuperMiniPile\", i.e. the density-sampled MiniPile (Idea 3) was compared to the Pythia $1.4\\text{B}$ trained on the original MiniPile.<br>\n",
    "Note that \"MiniPile Density\" contains $64,035$ documents less ($\\sim 6.34\\%$):\n",
    "\n",
    "| Benchmark        | Measure    |     | 1.4B MiniPile              | 1.4B Density                   | Percentage Difference in Means |\n",
    "| ---------------- | ---------- | --- | -------------------------- | ------------------------------ | ------------------------------ |\n",
    "| ARC-Challenge    | acc        | ↑   | **0.1903 ± 0.0115**        | 0.1852 ± 0.0114                | -2.6800                        |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035            | 0.2295 ± 0.0035                | 0.0000                         |\n",
    "| HellaSwag        | acc        | ↑   | 0.2579 ± 0.0044            | **0.2589 ± 0.0044**            | 0.3877                         |\n",
    "| WinoGrande       | acc        | ↑   | **0.5185 ± 0.0140**        | 0.5043 ± 0.0141                | -2.7387                        |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000            | 0.0000 ± 0.0000                | -                              |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | 1564928.5258 ± 118691.4565 | **1420846.8323 ± 106563.1327** | -9.2069                        |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000            | 0.0000 ± 0.0000                | -                              |\n",
    "| Lambada (Std)    | perplexity | ↓   | 8848600.9409 ± 745031.8900 | **7916035.3527 ± 664805.9178** | -10.5391                       |\n",
    "| BLiMP            | acc        | ↑   | **0.5483 ± 0.0017**        | 0.5422 ± 0.0017                | -1.1125                        |\n",
    "| ARC-Easy         | acc        | ↑   | **0.2715 ± 0.0091**        | 0.2698 ± 0.0091                | -0.6262                        |\n",
    "\n",
    "Finally, the impact of dataset size reduction on the larger model's retention of performance was evaluated by training on the \"MiniPile Density Pico\" dataset.\n",
    "\n",
    "| Benchmark        | Measure    |     | 1.4B Density                   | 1.4B Density Pico          | Percentage Difference in Means |\n",
    "| ---------------- | ---------- | --- | ------------------------------ | -------------------------- | ------------------------------ |\n",
    "| ARC-Challenge    | acc        | ↑   | 0.1852 ± 0.0114                | **0.1928 ± 0.0115**        | 4.1037                         |\n",
    "| MMLU             | acc        | ↑   | 0.2295 ± 0.0035                | 0.2295 ± 0.0035            | 0.0000                         |\n",
    "| HellaSwag        | acc        | ↑   | 0.2589 ± 0.0044                | **0.2600 ± 0.0044**        | 0.4249                         |\n",
    "| WinoGrande       | acc        | ↑   | 0.5043 ± 0.0141                | **0.5122 ± 0.0140**        | 1.5665                         |\n",
    "| Lambada (OpenAI) | acc        | ↑   | 0.0000 ± 0.0000                | 0.0000 ± 0.0000            | -                              |\n",
    "| Lambada (OpenAI) | perplexity | ↓   | **1420846.8323 ± 106563.1327** | 1662608.9444 ± 128444.3607 | 17.0154                        |\n",
    "| Lambada (Std)    | acc        | ↑   | 0.0000 ± 0.0000                | 0.0000 ± 0.0000            | -                              |\n",
    "| Lambada (Std)    | perplexity | ↓   | **7916035.3527 ± 664805.9178** | 8543578.1832 ± 737889.9436 | 7.9654                         |\n",
    "| BLiMP            | acc        | ↑   | 0.5422 ± 0.0017                | **0.5445 ± 0.0017**        | 7.9654                         |\n",
    "| ARC-Easy         | acc        | ↑   | 0.2698 ± 0.0091                | **0.2761 ± 0.0091**        | 2.3351                         |\n",
    "\n",
    "I can't say I expected this. The downsized dataset is not significantly worse, and in several cases even better.<br>\n",
    "This seems to underline that density-proportioned sampling in combination with size-reduction of the target distillate dataset has an actual positive effect on dataset quality, as improvements are evident across both 160M and 1.4B models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Cancelled Ideas\n",
    "\n",
    "The project's constraints for time and resource availability needed to be considered during the pursuit of different improvement ideas.\n",
    "As MiniPile itself is designed purposefully to be applied in academic settings, ideally, its creation should be feasible within this same setting as well.\n",
    "However, techniques like HDBSCAN clustering, which could potentially improve the clustering process, require more time and resources than could be deemed feasible for this project. The following ideas have been cancelled due to these constraints, but have been added in the `cancelled_ideas` folder and may be worth investigating in future projects:\n",
    "\n",
    "- **HDBSCAN clustering:** Replacing $k$-Means with HDBSCAN could help in identifying complex-shaped clusters, as well as outliers\n",
    "- **Lossi 3:** The loss-based sampling approach could have been improved further by using the proxy model during a second loss-based selection strategy. Where the cluster proportions where governed by losses already, the samples themselves could have been, too. For that, we would sample $1.5\\times$ the proportion per cluster and measure each example's loss with the proxy. From that, from a list sorted by loss in descending order, we would pick the top as $0.8\\times$ of the proportion and the lower as $0.2\\times$ of the proportion. This would have been a more fine-grained approach to the loss-based sampling idea, but loss itself is was deemed not a suitable measure, because it tends to favor complexity over information content, even with the $0.2\\times$ lower complexity proportion.\n",
    "- **Density-based sampling with $\\omega = 0.25$:** This approach was discarded to save time. In effect, the cluster size-proportioned sampling approach provided a reference for an $\\omega = 0.00$ setting already. More interest was therefore put into lower-density-based sampling, resulting in the tests with $\\omega=0.75$.\n",
    "    \n",
    "## Theoretical Ideas\n",
    "\n",
    "While the above ideas were cancelled during the implementation phase, others were put aside during brainstorming for similar reasons (feasibility within constraints).\n",
    "\n",
    "### Sparse Sampling of Embeddings with Similarity Hashing\n",
    "\n",
    "This idea arose from the Lossi $3$ cancellation and aims to improve the dataset creation process specifically without the need for a proxy model. This idea applies Locality Sensitive Hashing to the clusters resulting from the $k$-Means clustering step. Per cluster, LSH would hash embeddings into buckets based on similarity, potentially capturing more detailed relationships between documents than would be possible by $k$-Means clusters alone. Note that I would keep $k$-Means in place and would not replace it with LSH, as $k$-Means poses an overall low-cost grouping option allowing for human-guided filtering. Also, using LSH on cluster-level may help with resource usage. (This choice can admittedly be seen as a shortcoming as we don't leverage LSH's full potential of capturing relationships even across clusters, yet clustwe-wise handling would be more likely to execute.)\n",
    "\n",
    "Per cluster, LSH hashes embeddings into similarity-based buckets (e.g. $16$ buckets, arbitrarily). We'd attain $16$ sub-clusters per cluster by similarity and it is across these that we would sample from the cluster. Sampling could thus be more informed of inner-cluster topology. Still, sampling itself would not solely rely on LSH, but treat it as one of three configurably weighted factors:\n",
    "\n",
    "- **Proportional Sampling:** As in Idea 1, we would sample documents proportional to cluster sizes, enabling capture of original dataset distribution.\n",
    "- **Random Exploration:** Introduces 'total' inner-cluster randomness to capture some potentially unique or median examples.\n",
    "- **Semantic Diversity Sampling:** Prioritizes semantically diverse examples using LSH-based similarity groups, aiming at capturing more detailed relationships between documents.\n",
    "\n",
    "Rest of assembly would equal the implemented ideas.<br>\n",
    "I've read that the python framework `faiss` is praised for being efficient in LSH-based similarity hashing. Maybe this would be a good starting point.<br>\n",
    "This is what I stumbled across while researching LSH and its applications and what might be interesting:\n",
    "- [Locality Sensitive Hashing (LSH): The Illustrated Guide (pinecone.io)](https://www.pinecone.io/learn/series/faiss/locality-sensitive-hashing/)\n",
    "- [Embedding Compression with Hashing for Efficient Representation Learning in Large-Scale Graph (Yeh, et al. 2022)](https://arxiv.org/pdf/2208.05648)\n",
    "- [Hashing for Similarity Search: A Survey (Wang, et al. 2014)](https://arxiv.org/pdf/1408.2927)\n",
    "- [Efficient Measures for Processing and Exploring Large-Scale Data Domains (Reinbold, Christian. 2022)](https://mediatum.ub.tum.de/doc/1655492/ptz317jfxiatpjlwxrpcuimu1.2022-11-01_einreichung_mediatum_titelblatt_neu.pdf) (although not as directly related I fear)\n",
    "\n",
    "### Double-Proxied Cross-Entropy-based Sampling\n",
    "\n",
    "This idea would be closely aligned with the findings of [Extracting representative subset from extensive text data for\n",
    "training pre-trained language models (Suzuki, et al. 2023)](https://www.sciencedirect.com/science/article/pii/S0306457322003508). The authors claim that \"[...] the representative subset obtained using [a] likelihood difference score can achieve the $90\\%$ performance level even when the size of the dataset is reduced to approximately two to three orders of magnitude smaller than the original dataset.\"\n",
    "\n",
    "The approach scores and ranks data samples based on the likelihood difference between two differently pretrained language models (PreLMs):\n",
    "1. An in-domain PreLM, trained on domain-specific data, i.e. a conceptually interrelated subset of the larger dataset\n",
    "2. A non-domain PreLM, trained on a general domain, i.e. the full dataset\n",
    "\n",
    "For each document $X$, the likelihood difference score $S(X)$ is calculated as:\n",
    "$$S(X)=\\overline{H}^{(I)}(X) - \\overline{H}^{(N)}(X)$$\n",
    "$\\overline{H}^{(I)}$ and $\\overline{H}^{(N)}$ are the normalized per-word cross-entropy losses from in-domain PreLM and non-domain PreLM, respectively.\n",
    "\n",
    "Essentially, we would:\n",
    "1. Calculate $S(X)$ for all documents in the dataset.\n",
    "2. Rank documents in ascending order of $S(X)$.\n",
    "3. Deduplicate redundant samples (I'd argue this is not necessary for The Pile Deduplicated).\n",
    "4. Selecting the top-$K$ ranked samples as representative dataset (RepSet).\n",
    "\n",
    "It would be interesting to follow up on the paper's claims, but I'd fear this is a resource-intense and in the end still loss-based approach, overrepresenting compliated and not representing/informative documents in the final assembly. Further, this approach poses a catch $22$, because in order to train more cost-effective, we have to train the non-domain PreLM on the full dataset. For a constrained academic setting, this might not be possible, even if the proxy model size is downscaled.\n",
    "\n",
    "### Semantic Deduplication as Post-Processing for the distilled dataset\n",
    "\n",
    "An idea for a post-processing step would be to analyze the now smaller subset for semantic duplicates. I suggest the work of [SemDeDup: Data-efficient learning at web-scale through semantic deduplication (Abbas, et al. 2023)](https://arxiv.org/pdf/2303.09540) for this semantic deduplication as basis. It is claimed that for the LAION dataset, \"[...] SemDeDup can remove 50% of the data with minimal performance loss\". This could be a valuable step to further reduce the dataset size without losing performance.\n",
    "\n",
    "While this idea is interesting, I was hesitant to allocate resources to it, mostly because reducing the distillate solely by the perspective of semantic similarity might remove syntactically informative, yet semantically similar examples. This could lead to a loss of diversity in the dataset, which is not necessarily desirable. However, testing could be done relatively flexibly, as this can be put into action post-assembly.\n",
    "\n",
    "## Reinforcement Learning-based sampling\n",
    "\n",
    "Reinforcement learning (RL) techniques like DQN could be used to iteratively revise a sampling strategy that optimizes model performance for specific task sets. An agent could be rewarded for selecting examples that would lead to improved model performance in those optimized for tasks. From the above, we see that density-based sampling is based on human judgment. RL, on the other hand, would allow for learning a data-driven sampling strategy that adapts over time, autonomously. Staying with the terminology of RL, we could formulate the problem as follows:\n",
    "- **Environment:** The embedding-enriched Pile dataset, a current state of MiniPile, and a language model (like Pythia 160M), trained on the current MiniPile state.\n",
    "- **Agent:** An RL agent whose task it is to pick out examples from the Pile in order to have MiniPile emerge.\n",
    "- **Actions:** The agent can perform actions such as adding or removing from MiniPile on an individual example level.\n",
    "- **States:** A state includes all current examples in MiniPile and performance metrics of the language model trained thereon, e.g. perplexity scores and accuracies.\n",
    "- **Rewards:** The agent receives a reward based on the performance of the language model trained on MiniPile, measured on a validation set.\n",
    "- Better performance on the validation set would lead to a higher reward.\n",
    "\n",
    "I didn't follow up on this. I view it as possible to have RL employed without the need to be downstream-task-specific in its choosing, i.e. just produce a dataset for which examples are picked based on learned correlations with metric outcomes (benchmark results) to be best attainable, but I fear that the complexity (room for error) would be too high for a feasible implementation within the constraints of this project.<br>\n",
    "However, this approach could lift a need for clustering alltogether.<br>\n",
    "Training the agent and then repeatedly training the language model, though, even if just incremental, would be a resource-intensive and time-consuming process.\n",
    "\n",
    "Hypothetically, if implementation complexity were to be deemed neglegible, we could both increase interpretability and lift the need for clustering by employing time-wise incremental training, benchmarking, sampling and resampling more interpretable results by e.g. employing a solution akin to [Mastering Diverse Domains through World Models (Hafner, et al. 2021)](https://arxiv.org/abs/2301.04104). This algorithm in particular was proven to be executable in rather resource-constrained settings, i.e. robotics. Also, it contains components for simulation of world/model behavior, ideally, potentially, lifting the need for extreme step by step benchmarking otherwise needed as environment response. It would be interesting to see how this would perform on the dataset distillation problem, but I deem it architectually and resource-wise still to demanding for constrained academic settings."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}