AI-modeller og platforme

BlackMamba: En introduktion til MoE for State Space Models

mm
Føj Unite.AI til dine foretrukne kilder på Google

Udviklingen af Large Language Models (LLM’er) bygget på decoder-only transformer-modeller har spillet en afgørende rolle i transformationen af Natural Language Processing (NLP)-domænet samt fremme af diverse dybe læringssammenhænge, herunder reinforcement learning, tidsrækkeanalyse, billedbehandling og meget mere. Dog står LLM’er bygget på decoder-only transformer-modeller stadig over for betydelige udfordringer. Selvom de er udtryksfulde, kræver attention-mekanismen i transformer-afledte LLM’er omfattende beregningsressourcer under både inferens og træning, hvilket kræver en betydelig mængde hukommelse til sekvenslængden og kvadratiske FLOPs. Dette høje beregningskrav begrænser modellernes kontekstlængde, gør autoregressiv genereringstasks proportionelt dyre med skalaen, og hindrer modellernes evne til at lære fra kontinuerte datastrømme og proces sekvenser af ubegrænset længde effektivt.

I de seneste år har State Space Models (SSM’er) demonstreret bemærkelsesværdige evner og præstationer, konkurrerende med transformer-arkitekturmodeller i store skalamodelleringstests, mens de opnår lineær kompleksitet i forhold til sekvenslængden. Desuden har Mamba, en nyligt udgivet State Space Model, vist fremragende præstationer i en række sprogmodellering- og lange-sekvensbehandlingsopgaver. Samtidig har Mixture of Expert-modeller (MoE) også vist imponerende præstationer, mens de betydeligt reducerer latencyen og beregningsomkostningerne ved inferens, dog til en højere pris for en større hukommelsesaftryk. Byggende på Mamba og MoE-modeller, vil denne artikel diskutere BlackMamba, en ny arkitektur, der kombinerer Mamba State Space Model med MoE-modeller for at udnytte fordelene fra begge rammer. Eksperimenter med BlackMamba har demonstreret dets evne til at overgå den eksisterende Mamba-ramme og transformer-baselines i både trænings-FLOPs og inferens.

… (resten af artiklen følger)

BlackMamba: En introduktion til MoE for State Space Models

Fremkomsten af Large Language Models (LLM’er), især de baseret på decoder-only transformer-arkitekturer, har haft en betydelig indvirkning på Natural Language Processing (NLP)-feltet og udvidet til diverse dybe læringssammenhænge, herunder reinforcement learning, tidsrækkeanalyse, billedbehandling og mere. Dog står disse decoder-only transformer-baserede LLM’er over for betydelige udfordringer. Attention-mekanismen, en central del af transformer-baserede LLM’er, kræver omfattende beregningsressourcer under både inferens og træning.

Betydelige bestræbelser er gjort i de seneste år for at overvinde disse begrænsninger, og fokus er skiftet mod at udvikle alternative arkitekturer til de kanoniske tætte attention-transformer-modeller med SSM’er og MoE-modeller som de mest lovende kandidatarkitekturer. Det centrale fordel ved at foretrække State Space Models over transformer-arkitekturmodeller er den lineære beregningskompleksitet i forhold til sekvenslængden, som SSM’er tilbyder i modsætning til den kvadratiske kompleksitet, transformer-modeller tilbyder.

… (resten af artiklen følger)

BlackMamba: Arkitektur og Metodologi

State Space Models

State Space Models hører til gruppen af sekvensmodeller med lineær kompleksitet i forhold til sekvenslængden. Arkitekturen af State Space Models ligner mere rekursive neurale netværk og convolutionelle neurale netværk end attention-baseret arkitektur og er inspireret af et kontinuert dynamisk system, der mapper en 1-dimensionel funktion gennem en implicit latent rum.

Den indbyggede kvadratiske kompleksitet i forhold til sekvenslængde i transformer-modeller er en velkendt begrænsning, der hindrer forståelse og forståelse over meget lange kontekster. Nylige innovationer har introduceret idéen om at udvide kontekstlængden, så transformer-modeller kan trænes på en realistisk skala, før de anvendes på langt længere kontekster under inferens.

… (resten af artiklen følger)

Mixture of Expert Models

Mixture of Expert-modeller (MoE) opnår en adskillelse mellem inferensomkostningerne og det totale antal parametre ved at aktivere parametre selektivt under en enkelt forward-pass. I stedet for at bruge alle parametre, dirigerer disse modeller tokens til bestemte Multilayer Perceptron (MLP)-eksperter.

Arkitektur

BlackMamba anvender en standard transformer-model, der består af interpolerede MLP-blokke og attention-blokke, der føjes sekventielt langs en residual-strøm. Nu erstatter de fleste Mixture of Expert-modeller blot de multilagrede perceptron-blokke med et router-ekspertlag.

Træning og Dataset

BlackMamba-modellen er trænet på over 300 milliarder tokens på et brugerdefineret dataset og anvender SwiGLU-aktiveringsfunktionen for ekspert-multilagrede perceptron. Rammen træner med 8 eksperter, et antal, som udviklerne fandt var den rette balance og kompromis mellem hukommelsesaftryk og inferensomkostninger for modellen.

BlackMamba: Resultater

For at sikre en fair sammenligning mellem Mamba og BlackMamba, har udviklerne trænet begge modellerne med de samme træningsparametre på det samme træningsdata. BlackMamba-rammen kan overgå både Mamba og transformer-modeller for identisk forward-pass modelstørrelse ved både inferens- og træningstidspunkt.

Desuden kombinerer BlackMamba-rammen latencyfordelene fra både Mixture of Expert- og Mamba-modeller, resulterende i betydeligt hurtigere inferenstider i forhold til transformer-modeller, rene Mamba-modeller og MoE-modeller.

Endelige tanker

I denne artikel har vi talt om BlackMamba, en ny arkitektur, der kombinerer Mamba State Space Model med MoE-modeller for at udnytte fordelene fra begge rammer. Eksperimenter med BlackMamba har demonstreret dets evne til at overgå den eksisterende Mamba-ramme og transformer-baselines i både trænings-FLOPs og inferens.

En ingeniør af profession, en forfatter af hjerte. Kunal er en teknisk forfatter med en dyb kærlighed og forståelse af AI og ML, dedikeret til at forenkle komplekse koncepter inden for disse felter gennem sin engagerende og informative dokumentation.