Technische Samenvatting: HERMES: Een Multi-Granulariteit Labelingssubstraat voor Pre-training Datamengsels

1. Probleemstelling

Huidige datamengstrategieën voor LLM pre-training kampen met een bottleneck in het labelsysteem in plaats van in de mengalgoritmen zelf. Bestaande benaderingen verdelen corpora doorgaans in groepen met behulp van vaste taxonomieën (herkomst, onderwerpen, formaten) of vlakke embedding-clusters (bijv. KMeans). Deze methoden lijden onder drie primaire beperkingen:

1. Enkele Granulariteit/As: Ze committeren zich aan een enkele semantische as op een vaste resolutie. Het wijzigen van de granulariteit (bijv. van brede onderwerpen naar fijnmazige clusters) vereist het volledig opnieuw opbouen van de volledige labelset.
2. Herberekeningskosten: Vlakke clusteringmethoden (zoals KMeans) vereisen herberekening om het aantal clusters ($K$) te wijzigen, waardoor systematische sweeps van granulariteit computationeel duur en onsamenhangend zijn.
3. Onflexibiliteit: De keuze voor groepen bepaalt wat een mixer kan uitdrukken, wat de mogelijkheid beperkt om interacties tussen de resolutie van data en bemonsteringsstrategieën te testen.

De auteurs stellen dat de bottleneck de labelinfrastructuur is, niet de mixer-machinerie. Er is behoefte aan een herbruikbaar, uit data afgeleid substraat dat het voor practitioners mogelijk maakt om door een hiërarchie van granulariteiten te navigeren zonder opnieuw te clusteren.

2. Methodologie: HERMES

HERMES (Hierarchical Embedding-based RVQ for Mixture Selection) is een one-shot, hiërarchisch RVQ-substraat ontworpen om documenten te annoteren met grove-naar-fijne codes. Het bestaat uit twee hoofdstadia:

A. Learned Semantic Transform (LST)

Om embeddings voor kwantisatie voor te bereiden, gebruikt HERMES een Learned Semantic Transform.

• Input: Bevroren document-embeddings ($e_i \in \mathbb{R}^{1024}$) van een pre-trained encoder.
• Mechanisme: Een enkele lineaire laag met normalisatie ($h_i = \text{normalize}(W e_i + b)$) roteert de embeddingruimte.
• Trainingsdoel: De transformatie wordt gezamenlijk getraind met de kwantisatie-codeboeken om drie verliezen te optimaliseren:
  1. Behoud van Paarstructuur: Handhaaft de cosinus-geometrie tussen ruwe en getransformeerde embeddings.
  2. Kwantisatie-bewuste Reconstructie: Minimaliseert de reconstructiefout van de vectorkwantizer.
  3. Orthogonaliteit: Voorkomt representatie-instorting door de bijna-orthogonaliteit van de transformatiematrix $W$ af te dwingen.

B. 3-Stage Residual Vector Quantization (RVQ)

De getransformeerde embeddings ($h_i$) worden gecodeerd door een gekascadeerde 3-staps RVQ.

• Architectuur: Drie opeenvolgende vectorkwantizers, elk met een codeboekgrootte $K=256$.
• Proces:
  1. Stadium 1 (Grof): Codeert $h_i$ om de dichtstbijzijnde centroid $c_1$ te vinden, wat een residu $r_2 = h_i - q_1$ produceert.
  2. Stadium 2 (Intermediair): Codeert $r_2$ om $c_2$ te vinden, wat een residu $r_3$ produceert.
  3. Stadium 3 (Fijn): Codeert $r_3$ om $c_3$ te vinden.
• Output: Elk document ontvangt een hiërarchische code $(c_1, c_2, c_3)$.
• Granulariteitscontrole: De granulariteit wordt gedefinieerd door de prefixlengte van de code:
  • L1: Code $(c_1)$ $\rightarrow$ 256 buckets.
  • L12: Code $(c_1, c_2)$ $\rightarrow$ $\sim$65k buckets.
  • L123: Code $(c_1, c_2, c_3)$ $\rightarrow$ $\sim$130k geobserveerde buckets (van een nominale $1.67 \times 10^7$).

C. Sampling Taxonomie

Het paper definieert een sampler als een compositie van twee stadia die opereren op het HERMES-substraat:

1. Stadium 1 (Outer Weight): Wint gewichten toe aan L1-voorouders. Het paper gebruikt Uniform (max-entropie) of DoReMi (group-DRO op een proxy-model).
2. Stadium 2 (Sub-bucket Sampler): Beslist over geschiktheid en selectie binnen L1-voorouders.
   • Massaverdeling: Gelijk (max-entropie) versus grootte-proportioneel.
   • Per-document Geschiktheid: Alle documenten versus de top-30% op basis van kwaliteitsscore (FineWeb-Edu) binnen de sub-bucket.

3. Belangrijkste Bijdragen

1. HERMES Substraat: Een one-shot hiërarchisch RVQ-systeem dat $\sim$50 miljoen documenten annoteert in maximaal $\sim$130k cellen. Het maakt granulariteit bloot als een prefixlengte op een enkele getrainde codeboekstap, waardoor het niet nodig is om opnieuw te clusteren wanneer de resolutie verandera.
2. Granulariteit-Sampler Interactie: De auteurs demonstreren dat de effectiviteit van een bemonsteringsregel (bijv. kwaliteitsgebaseerde selectie) afhankelijk is van het granulariteitsniveau.
3. Candidate Competition Diagnostic: Het paper identificeert een structurele limiet waarbij het vergroten van de granulariteit de omvang van de kandidaat-pool verkleint tot een punt waar ranking binnen de bucket instabiel wordt.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op een 1B-parameter LLaMA-stijl model getraind voor 25B tokens, geëvalueerd op een 16-taak bekwaamheids macro-gemiddelde (Avg.).

A. Het Clustering Plateau (L1=256)

• Bevinding: Op de grofste granulariteit (L1=256) presteert HERMES op gelijke voet met standaard clusteringmethoden (KMeans, MiniBatchKMeans, Plain RVQ).
• Metrieken: Op intrinsieke metrieken (compactheid, massabalans, entropie) bevindt HERMES zich op een plateau met andere 256-weg clusteraars.
• Downstream: Wanneer gekoppeld aan een Uniform + L1 max-entropie sampler, verschillen HERMES en KMeans slechts 0.0002 in Avg. nauwkeurigheid.
• Implicatie: De prestatiewinst van HERMES komt niet voort uit een superieur clusteringalgoritme, maar uit het substraat zelf.

B. Bevinding A: Stadium-2 Regel bij Intermediaire Granulariteit (L12)

• Setup: Vaste DoReMi L1 buitenste gewichten, vaste HERMES codeboek, granulariteit op L12 ($\sim$65k buckets).
• Vergelijking: Het wisselen van de Stadium-2 sampler van max-entropie dekking (gelijke sub-bucket massa, alle docs) naar kwaliteit top-30% (grootte-proportionele massa, top-30% volgens FineWeb-Edu score).
• Resultaat: Deze switch verhoogde de 16-taak Avg. met +0.0253 (z-score +2.09).
• Interpretatie: Bij L12 zijn sub-buckets groot genoeg (mediaan pool $\sim$2.271 documenten) zodat kwaliteitsranking een stabiele proxy is voor globale kwaliteit.

C. Bevinding B: Instorting bij Fijne Granulariteit (L123)

• Setup: Dezelfde regels toegepast op L123 ($\sim$130k buckets).
• Observatie: De mediane kandidaat-pool krimpt 5.3x van L12 (2.271 docs) naar L123 (429 docs).
• Resultaat: Het voordeel van de kwaliteit top-30% regel stort in. Het prestatieverschil tussen de kwaliteit top-30% en lokale random dekking krimpt naar +0.0002 (0.3988 vs. 0.3986), wat statistisch ononderscheidbaar is onder single-seed evaluatie.
• Implicatie: Dit ondersteunt de candidate competition hypothese: wanneer sub-buckets te klein worden, is ranking binnen de bucket niet langer een betrouwbare proxy voor globale kwaliteit, wat het voordeel van kwaliteitsgebaseerde selectie tenietdoet.

D. Beste Configuratie

De hoogst presterende configuratie in de sweep was DoReMi-L1 buitenste gewichten + L12 kwaliteit top-30% (corrected reader), die een Avg. van 0.4222 behaalde.

5. Betekenis en Claims

Het paper claimt bescheiden dat HERMES de vormgeving van datamengsels herdefinieert. De primaire betekenis ligt in het verschuiven van de focus van "kiezen tussen vaste labelsets" naar "navigeren door een herbruikbare, uit data afgeleide granulariteitshierarchie."

• Substraat boven Clusterer: De bijdrage is de infrastructuur die het mogelijk maakt om granulariteit een vrije variabele (een draaiknop) te maken, niet de clusteringprestatie zelf.
• Gezamenlijke Bepaling: Granulariteit en de Stadium-2 sampler zijn gezamenlijk bepaalde assen. Een regel die werkt op één granulariteit (L12) kan falen op een andere (L123) vanwege structurele beperkingen zoals de omvang van de kandidaat-pool.
• Beperkingen: De auteurs erkennen dat conclusies beperkt zijn tot een 1B/25B regime op een interne corpus. De optimale granulariteitsdiepte wordt begrensd door "candidate competition," wat kan verschuiven met grotere modellen en corpora. Ze merken ook op dat de keuze van de buitenste gewichten (DoReMi vs. Uniform) bij de optimale Stadium-2 regel (L12 kwaliteit top-30%) niet volledig is gecross-validated (de Uniform + L12 kwaliteit top-30% run is niet uitgevoerd).

Samenvattend biedt HERMES een mechanisme om datagranulariteit te behandelen als een continue, instelbare parameter, wat onthult dat de effectiviteit van dataselectiestrategieën inherent verbonden is met de resolutie van het onderliggende labelsysteem.