Tuning Language Models by Mixture-of-Depths Ensemble

Technische Samenvatting: Tuning van Taalmodellen via Mixture-of-Depths Ensemble (MoDE)

Probleemstelling
Traditionele Transformer-gebaseerde Large Language Models (LLMs) vertrouwen uitsluitend op representaties van de laatste lagen voor zowel voorspelling als verliesberekening tijdens het fine-tuning proces. Hoewel interpreteerbaarheidstools zoals de "logit lens" hebben aangetoond dat de verborgen toestanden van latere lagen al grotendeels gevormde, taakrelevante voorspellingen bevatten, worden deze intermediaire representaties doorgaans behandeld als impliciete bijproducten in plaats van expliciete trainingssignalen. Bijgevolg wordt het voorspellende vermogen van late lagen onderbenut, en wordt standaard fine-tuning alleen geoptimaliseerd voor de laatste laag, waardoor intermediaire representaties slechts indirect worden geoptimaliseerd. Bovendien benutten bestaande methoden om intermediaire lagen te exploiteren (bijv. vroege exit voor efficiëntie of representatie-sturing) de laagspecifieke logits vaak niet als een direct, trainbaar onderdeel van de uiteindelijke voorspellingsdoelstelling.

Methodologie: Mixture-of-Depths Ensemble (MoDE)
De auteurs stellen Mixture-of-Depths Ensemble (MoDE) voor, een tuning-framework dat de outputs van de laatste $k$ lagen van een LLM behandelt als een ensemble van "experts" binnen één enkel model. In tegen tegenstelling tot Mixture-of-Experts (MoE), waarbij verschillende tokens naar verschillende modellen worden gerouteerd, routeert MoDE de uiteindelijke voorspelling naar een gewogen combinatie van logits van verschillende dieptes binnen hetzelfde model.

Het framework bestaat uit drie kerncomponenten:

1. Extractie van Late-Layer Logits: Voor een model met $n$ lagen worden de laatste $k$ lagen ($K = \{L_{n-k+1}, \dots, L_n\}$) geselecteerd. De vocabulaire-head $\phi(\cdot)$ wordt toegepast op de verborgen toestanden van elk van deze $k$ lagen om logits $\ell_i(x)$ te genereren.
2. Learned Routing Network: Een lichtgewicht router $G(x)$, geïmplementeerd als een enkele lineaire laag gevolgd door een softmax, berekent de routing-gewichten op basis van de verborgen toestand vlak vóór de laatste $k$ lagen ($x = H_{n-k}$). De uiteindelijke logits worden berekend als een gewogen som:
   $$\ell(x_t | x_{<t}) = \sum_{i=0}^{k-1} G(x)_i \cdot \ell_i(x)$$
   De router kan spars gemaakt worden (Top-K) om vroege exit tijdens inferentie mogelijk te maken.
3. Adaptatie via Normalisatie en Distillatie: Het direct toepassen van de pre-trained head op intermediaire lagen resulteert vaak in een hoge initiële loss. Om dit aan te pakken zonder zware parameters toe te voegen (in tegenstelling tot per-laag affine matrices), introduceren de auteurs:
   • Per-laag Normalisatie ($N_k$): Trainbare normalisatie-modules (die overeenkomen met de pre-trained architectuur, bijv. LayerNorm) worden toegevoegd aan elke late laag om de representaties aan te passen aan de inputruimte van de head.
   • Distillatieverlies ($L_{distill}$): De laatste laag fungeert als een "leraar". Een KL-divergentieverlies wordt toegepast om de outputdistributies van de eerdere late lagen naar de distributie van de laatste laag te trekken, wat de training stabiliseert en het ensemble afstemt op de taakdoelstelling.

De totale loss is $L_{MoDE} = L_{task} + \lambda L_{distill}$. MoDE is ontworpen als een plug-in mechanisme dat kan worden toegepast op elke bestaande parameter-efficiënte fine-tuning methode (PEFT) (bijv. LoRA, DoRA) of onafhankelijk kan worden gebruikt.

Belangrijkste Bijdragen

• Trainbare Late-Layer Signalen: Het artikel demonstreert dat late-layer logits, wanneer ze individueel geoptimaliseerd worden met distillatie, convergeren naar losses die vergelijkbaar zijn met de laatste laag en een complementaire set problemen oplossen.
• MoDE Framework: Een eenvoudig, laag-overhead tuning-mechanisme dat late-layer logits enzembelt via een geleerde router. Het voegt verwaarloosbare trainbare parameters toe (ongeveer +0,04% van de modelgrootte), terwijl het in staat is om substantieel grotere trainbare modules (zoals full-layer LoRA) te vervangen met vergelijkbare prestaties.
• Efficiëntie en Flexibiliteit: Het framework ondersteunt sparse routing (Top-K), wat vroege exit en versnellingen van de generatiesnelheid tot 1,6$\times$ mogelijk maakt met minimaal verlies aan nauwkeurigheid. Het combineert met bestaande PEFT-methoden, waardoor hun prestaties worden verbeterd.

Experimentele Resultaten
Experimenten werden uitgevoerd op LLaMA2-7B en LLaMA3-8B voor taken op het gebied van rekenkundig redeneren, commonsense redeneren en instructieopvolging.

• Rekenkundig Redeneren: MoDE verbeterde consequent de prestaties van basismethoden (LoRA, DoRA). De lichtgewicht variant (LoRA toegepast op alle lagen behalve de laatste $k$, plus MoDE) bereikte prestaties binnen 0,3–0,5 punten van full-layer LoRA, terwijl er aanzienlijk minder parameters werden gebruikt.
• Commonsense Redeneren: MoDE bood additieve winst boven gewicht-gebaseerde methoden (LoRA, DoRA). Terwijl representatie-gebaseerde methoden (ReFT) op zichzelf sterk presteerden, bood MoDE betrouwbare verbeteringen wanneer het werd gecombineerd met gewicht-gebaseerde benaderingen.
• Instructieopvolging: MoDE leverde kleine maar consistente winsten op MT-Bench, hoewel de verbeteringen minder bescheiden waren dan bij redeneertaken, wat suggereert dat instructieopvolging sterk leunt op late-layer formatting die ensemblen slechts gedeeltelijk vastleggen.
• Ablatie en Analyse:
  • De router concentreert het gewicht van nature op de laatste één of twee lagen, maar de distillatie- en normalisatiecomponenten zorgen ervoor dat eerdere late lagen effectief kunnen bijdragen.
  • Een ensemble-omvang van $k \approx 3\text{--}4$ is optimaal; een grotere $k$ helpt bij eenvoudige datasets maar schaadt complexe redeneertaken.
  • Sparse Top-K routing behoudt de meeste nauwkeurigheid terwijl het significante versnellingen mogelijk maakt.

Betekenis en Claims
De auteurs presenteren MoDE als een mechanisme-studie die aantoont dat late-layer logits direct bruikbaar kunnen worden gemaakt voor tuning. Het artikel stelt dat:

1. Late-layer representaties bezitten een aanzienlijk voorspellend potentieel dat kan worden ontsloten via eenvoudige adaptatie (normalisatie en distillatie).
2. Deze intermediaire signalen kunnen substantieel grotere trainbare modules (zoals full-layer LoRA) vervangen met vergelijkbare prestaties, wat een parameter-efficiënte alternatief biedt.
3. MoDE een complementaire richting biedt om intermediaire laag-representaties om te zetten in actieerbare trainingssignalen, verschillend van inference-time probing of gewichts-editing benaderingen.

Het artikel blijft bescheiden over de reikwijdte en merkt beperkingen op, zoals de afhankelijkheid van een vaste $k$ (in plaats van dynamische selectie), kleinere winsten bij instructieopvolging vergeleken met redeneren, en een evaluatie die beperkt is tot decoder-only LLMs. Toekomstig werk wordt gesuggereerd voor dynamische $k$-selectie, bredere modelfamilies en grotere modellen.