Ayn: A Tiny yet Competitive Indian Legal Language Model Pretrained from Scratch

Dit paper introduceert Ayn, een compact Indiaas juridisch taalmodel met slechts 88 miljoen parameters dat, ondanks zijn kleine formaat, presteert die vergelijkbaar zijn met of zelfs beter zijn dan veel grotere taalmodellen op specifieke juridische taken.

De kern

🇳🇱 AYN: De Slimme, Kleine Jurist in je Broekzak

Stel je voor dat je een gigantische bibliotheek hebt met alle boeken ter wereld. De huidige "super-intelligente" kunstmatige intelligenties (zoals de grote LLM's) zijn als reuzen die die hele bibliotheek hebben ingeslikt. Ze kunnen over alles praten, maar ze zijn enorm zwaar, duur om te voeden (trainen) en soms vergeten ze de specifieke regels van jouw eigen dorp.

De auteurs van dit paper vragen zich af: "Hebben we echt die enorme reus nodig om een specifieke juridische vraag in India te beantwoorden, of kan een slimme, kleine dwerg het misschien zelfs beter doen?"

Het antwoord is AYN (een naam die klinkt als 'Ogen', maar hier staat voor een Tiny Language Model).

1. Het Probleem: De Reus vs. De Specialist

De grote AI-modellen zijn als all-inclusive resorts. Ze hebben een zwembad, een golfbaan en een restaurant, maar als je een heel specifiek juridisch probleem hebt in India, kunnen ze soms de lokale dialecten of oude wetten niet goed begrijpen. Ze zijn ook duur om te bouwen en te gebruiken.

In de Indiase wetgeving is het nog ingewikkelder:

• Taal: Het is een mix van Engels, Hindi en andere talen (code-switching).
• Woorden: Er zijn oude, moeilijke woorden (zoals "notwithstanding" of "hereinafter") en specifieke verwijzingen naar wetten (bijv. "Sectie 3(1)(b)").
• Grootte: De documenten zijn vaak enorm lang en volgestopt met details.

De grote AI's zien deze zinnen vaak als een rommelige brij van losse letters, omdat hun "woordenboek" (tokenizer) niet is gemaakt voor juridische taal.

2. De Oplossing: AYN, de Slimme Dwerg

De onderzoekers bouwden AYN, een model met slechts 88 miljoen parameters. Ter vergelijking: de grote modellen hebben 1 tot 8 miljard parameters. AYN is dus 80 tot 90 keer kleiner dan de grootste concurrenten.

Hoe hebben ze dit gedaan?

• Specifiek Voedsel: In plaats van alles te eten, kregen ze alleen maar Indiase Hooggerechtshof-uitspraken, de Indiase Grondwet en het Wetboek van Strafrecht.
• Speciaal Woordenboek: Ze bouwden een eigen "woordenboek" (tokenizer) dat specifiek is gemaakt voor juridische termen.
  • Analogie: Stel je voor dat een gewone AI het woord "statutory" (wettelijk) opbreekt in rare stukjes als "stat", "ut", "ory". AYN ziet het als één heel woord. Dat maakt het veel sneller en slimmer in het begrijpen van de zin.
• Efficiëntie: Ze trainden dit model op één enkele computerchip (een A100 GPU) gedurende 185 uur. Dat is alsof je een auto bouwt in je garage in plaats van in een fabriek. Het kostte minder dan $500 en was heel milieuvriendelijk (weinig CO2-uitstoot).

3. De Wedstrijd: Klein vs. Groot

De onderzoekers lieten AYN heten tegen de grote reuzen (zoals Llama-3 en OLMo) in twee belangrijke juridische spellen:

Spel 1: Het Voorspellen van Uitspraken

• De taak: Lees een rechtszaak en voorspel of de klager wint of verliest.
• Het resultaat: AYN won! Zelfs zonder dat ze het model eerst extra hadden getraind (zero-shot), was AYN beter dan de modellen die 30 tot 80 keer groter waren.
• Waarom? Omdat AYN zich volledig heeft ingelezen in de Indiase wetgeving, terwijl de grote modellen te veel "afleiding" hebben van andere landen en onderwerpen.

Spel 2: Het Samenvatten van Rechtszaken

• De taak: Lees een 28.000 woorden lang document en schrijf een samenvatting van 5.000 woorden.
• Het resultaat: AYN deed het net zo goed als modellen die 30 keer groter waren. Het kon de kern van de zaak goed vastleggen.
• Let op: Als je het model vraagt om een extreem lange samenvatting (bijna de hele tekst), raakt het de draad kwijt. Maar voor normale samenvattingen is het een kampioen.

Spel 3: Algemene Vragen (MMLU, Logica, etc.)

• De taak: Vragen beantwoorden over algemene kennis, logica en taal.
• Het resultaat: AYN deed het verrassend goed! Het was beter dan sommige modellen die 10 keer groter waren, en deed het ongeveer even goed als een model dat 3 keer groter was.
• Conclusie: Zelfs als je alleen maar juridische boeken leest, leer je toch genoeg over de wereld om algemene vragen te beantwoorden.

4. Waarom is dit belangrijk?

Stel je voor dat je een dokter nodig hebt voor een zeldzame ziekte in een afgelegen dorp.

• De grote AI is als een beroemd internationaal ziekenhuis: geweldig, maar je moet er uren voor reizen, het kost een fortuin, en de artsen kennen misschien de lokale gewoonten niet.
• AYN is als een lokale specialist die precies weet hoe de ziekte in dat dorp zich voordoet. Hij is goedkoop, snel, en werkt perfect in de lokale taal.

De voordelen van AYN:

1. Kosten: Je kunt het draaien op een gewone laptop of server, niet op dure supercomputers.
2. Snelheid: Het is veel sneller in het geven van antwoorden.
3. Eerlijkheid: Grote AI's zijn vaak "Westerse" modellen die de Indiase cultuur en wetten niet goed begrijpen. AYN is puur voor India gemaakt.
4. Milieu: Het verbruikt een fractie van de energie van de grote modellen.

5. De Kijk in de Kristallen Bol (Beperkingen)

De auteurs zijn eerlijk: AYN is niet perfect.

• Het is alleen getraind op Engelse documenten van het Hooggerechtshof. Het kan nog geen Hindi of andere Indiase talen.
• Het is nog niet getest op alle soorten rechtbanken (zoals lagere rechtbanken).
• Omdat het een generatief model is, kan het soms "hallucineren" (verzonnen feiten vertellen). Daarom zeggen ze: Gebruik dit niet zomaar voor echte juridische zaken zonder dat een menselijke expert het nakijkt.

Samenvatting in één zin

AYN bewijst dat je niet altijd de grootste en duurste AI nodig hebt; een klein, slim en specifiek getraind model kan op zijn eigen terrein (Indiase wetgeving) de reuzen verslaan, terwijl het goedkoper, sneller en milieuvriendelijker is.

Technische samenvatting

Probleemstelling

Grote Taalmodellen (LLMs) zijn momenteel de standaard voor Natural Language Processing (NLP), maar ze zijn duur om te trainen en te gebruiken. Voor domeinspecifieke taken, zoals juridische analyse in India, zijn er specifieke uitdagingen:

1. Kosten en Toegankelijkheid: Het trainen van grote modellen (1B-8B parameters) vereist enorme rekenkracht.
2. Data Schaarste: Geannoteerde juridische data is schaars, duur en tijdrovend om te produceren.
3. Taalkundige Complexiteit: Indiaas juridisch taalgebruik kenmerkt zich door archaïsch vocabulaire, lange geneste zinnen, complexe citaten en code-switching (meertaligheid). Algemene tokenizers (zoals die van LLaMA) fragmenteren deze termen inefficiënt, wat de modelprestaties schaadt.
4. Westers Bias: Bestaande LLMs zijn voornamelijk getraind op westerse juridische systemen en missen de nuances van het Indiase rechtssysteem.

De centrale vraag is of een Tiny Language Model (TLM) met minder dan 100 miljoen parameters, dat vanaf nul is voorgeïntroduceerd op specifieke domeindata, concurrerend kan zijn met veel grotere modellen.

Methodologie

De auteurs hebben AYN ontwikkeld, een decoder-only taalmodel met 88 miljoen parameters, specifiek voor het Indiase rechtssysteem.

1. Data en Corpus:

• Het trainingscorpus is een uitgebreide versie van de Indian Legal Documents Corpus (ILDC).
• Het bevat 38.222 Supreme Court-uitspraken (1947-2023), de Grondwet van India en het Indiase Wetboek van Strafrecht.
• Totaal: ongeveer 142,6 miljoen woorden.
• De data is schoongemaakt met reguliere expressies om meta-informatie (zoals namen van rechters, data) te verwijderen, terwijl de kern van de uitspraak behouden bleef.

2. Tokenisatie:

• In plaats van een bestaande tokenizer te gebruiken, hebben de auteurs een domeinspecifieke BPE-tokenizer (Byte-Pair Encoding) vanaf nul getraind.
• Deze tokenizer is geoptimaliseerd voor juridisch jargon (bijv. "hereinafter", "statutory bar") en voorkomt inefficiënte fragmentatie die bij algemene tokenizers optreedt.
• Het vocabulaire is compact (3.500 tokens) maar zeer effectief voor juridische tekst.

3. Model Architectuur en Training:

• Architectuur: Gebaseerd op de Transformer-architectuur (causal decoder) met RMSNorm, SwiGLU-activatie en weight tying.
• Grootte: 768 hidden dimension, 12 lagen, feedforward hidden size van 2048.
• Contextlengte: Trainen met een contextlengte van 8192 tokens op slechts één NVIDIA A100 GPU. Dit werd mogelijk gemaakt door een interpolatiemethode voor positionele embeddings (RoPE) met een "shrinking factor".
• Training: Het model is 185 uur getraind op één A100 GPU. Er werd gebruikgemaakt van µP-transfer (zero-shot hyperparameter transfer) van een kleiner model om de optimale leercurve te bepalen.
• Efficiëntie: Het totale energieverbruik was 50,875 kW, wat resulteert in een zeer lage CO2-voetafdruk (0,0196 tCO2eq).

Belangrijkste Bijdragen

1. Nieuw Corpus: Een uitgebreid corpus van Indiase Supreme Court-uitspraken, de Grondwet en het Wetboek van Strafrecht.
2. Domeinspecifieke Tokenizer: Een nieuwe BPE-tokenizer die specifiek is ontworpen voor de complexiteit van Indiase juridische taal.
3. AYN Model: Een 88M-parameter TLM dat volledig vanaf nul is getraind op dit corpus, met een contextlengte van 8192 tokens, getraind op beperkte hardware.
4. Empirisch Bewijs: Een uitgebreide vergelijking die aantoont dat een klein, domeinspecifiek model grotere, algemene modellen kan verslaan op specifieke taken.

Resultaten

Het model werd geëvalueerd op twee juridische taken en algemene NLP-benchmarks, vergeleken met modellen van 1B tot 8B parameters (zoals Llama-3, Llama-2, OLMo).

1. Juridische Uitspraakvoorspelling (Classification):

• Zero-shot: AYN (88M) presteerde beter dan alle vergeleken LLMs (tot 8B) en zelfs beter dan een specifiek voorgetrainde Llama-2 7B (CPTLlama-2).
• Discriminative Classifier: Wanneer een eenvoudige classificatielaag bovenop het model werd getraind, bereikte AYN 69,69% nauwkeurigheid, wat aanzienlijk hoger was dan de grote modellen (die rond de 60% bleven). Dit bevestigt dat AYN dichte domeinrepresentaties leert.

2. Juridische Samenvatting (Abstractive Summarization):

• AYN kon samenvattingen genereren die qua semantische coherentie en overlap met referenties (ROUGE, METEOR, BERTScore) concurreren met modellen tot 30 keer groter (tot 3B parameters).
• Het presteerde echter minder goed dan modellen groter dan 7B op deze specifieke generatietaken, maar deed dit met een veel lager rekenkosten.

3. Algemene NLP-taken (Benchmarks):

• Op benchmarks zoals MMLU, LogiQA en QNLI presteerde AYN verrassend goed.
• Het versloeg gemiddeld zes grotere modellen (zoals Pythia 6.9B, Falcon 7B) en deed het op gelijke voet met OLMo 1B, ondanks dat het uitsluitend op juridische data is getraind. Dit toont aan dat domeinspecifieke TLMs bepaalde generalisatievermogens kunnen ontwikkelen.

Betekenis en Conclusie

Het paper demonstreert dat Tiny Language Models (TLMs), wanneer ze vanaf nul worden getraind op een zorgvuldig samengesteld, domeinspecifiek corpus met een aangepaste tokenizer, een zeer krachtig alternatief kunnen zijn voor grote, dure LLMs in gespecialiseerde domeinen.

• Kostenefficiëntie: Het trainen kostte minder dan $500 en 185 uur op één GPU, vergeleken met miljoenen dollars voor grote modellen.
• Duurzaamheid: De CO2-voetafdruk is meerdere ordes van grootte lager dan die van standaard LLMs.
• Domeinprestaties: Voor specifieke taken zoals juridische classificatie in India, overtreft een 88M-parameter model zelfs 8B-parameter modellen.
• Toekomstperspectief: Dit benadrukt de waarde van "small but specialized" modellen voor regio's en domeinen die ondervertegenwoordigd zijn in de huidige AI-landschap, waarbij het vermijden van "one-size-fits-all" modellen leidt tot betere en eerlijkere resultaten.

Beperkingen: Het model is momenteel beperkt tot Engels en Supreme Court-uitspraken (niet districtshoven). Er is geen menselijke juridische expert-evaluatie uitgevoerd voor hallucinaties, en het model wordt niet aanbevolen voor gebruik zonder menselijke supervisie.