AI-mallit ja alustat

Jamba: AI21 Labsin Uusi Hybridirakenteinen Transformer-Mamba-kielikoodi

mm
Lisää Unite.AI suosikkilähteisiisi Google-palvelussa

Kielenmallit ovat kokeneet nopeaa kehitystä, ja Transformer-pohjaiset arkkitehtuurit ovat johtaneet kehitystä luonnollisen kielen prosessoinnissa. Kuitenkin, kun malleja skaalataan, haasteet pitkien kontekstien käsittelyssä, muistin tehokkuudessa ja läpipääsyssä ovat tulleet yhä selkeämmiksi.

AI21 Labs on esittänyt uuden ratkaisun Jamban kanssa, joka on valmis, suurikokoinen kielenmalli (LLM), joka yhdistää Transformer- ja Mamba-arkkitehtuureiden vahvuudet hybridirakenteisessa kehyksessä. Tämä artikkeli käy Jamban arkkitehtuuriin, suorituskykyyn ja mahdollisiin sovelluksiin.

Jamban yleiskatsaus

Jamba on hybridisuuri kielenmalli, jota on kehittänyt AI21 Labs, hyödyntäen Transformer-kerrosten ja Mamba-kerrosten yhdistelmää, johon on integroitu Mixture-of-Experts (MoE) -moduuli. Tämä arkkitehtuuri mahdollistaa Jamballe tasapainon muistin käytön, läpipääsyn ja suorituskyvyn välillä, mikä tekee siitä voimakkaan työkalun laajalle valikoimalle NLP-tehtäville. Malli on suunniteltu mahtumaan yhteen 80 GB:n GPU:hen, tarjoten korkean läpipääsyn ja pienen muistijäljen samalla, kun ylläpidetään huippuluokan suorituskykyä erilaisilla benchmark-kojeilla.

Jamban arkkitehtuuri

Jamban arkkitehtuuri on avain sen kykyihin. Se perustuu uuteen hybridisuunnitteluun, jossa Transformer-kerrokset ja Mamba-kerrokset vuorottelevat, ja siihen on integroitu MoE-moduuli mallin kapasiteetin lisäämiseksi ilman merkittävää laskentakustannusten kasvua.

1. Transformer-kerrokset

Transformer-arkkitehtuuri on modernien LLM-mallien standardi sen kyvyn vuoksi käsitellä rinnakkaisprosessointia tehokkaasti ja sieventää pitkänmatkan riippuvuuksia tekstissä. Kuitenkin sen suorituskyky on usein rajoitettu korkeiden muisti- ja laskentavaatimusten vuoksi, erityisesti pitkien kontekstien prosessoinnissa. Jamba ratkaisee nämä rajoitukset integroimalla Mamba-kerrokset, joista keskustellaan seuraavaksi.

2. Mamba-kerrokset

Mamba on viimeaikainen tila-avaruusmalli (SSM), joka on suunniteltu käsittelemään pitkänmatkan suhteita sekvensseissä tehokkaammin kuin perinteiset RNN:t tai jopa Transformerit. Mamba-kerrokset ovat erityisen tehokkaita vähentämään Transformerien muistijälkeä, joka liittyy avain-arvo (KV) -välimuistiin. Vuorottelemalla Mamba-kerroksia Transformer-kerrosten kanssa Jamba vähentää kokonaismuistin käyttöä ylläpitäen samalla korkeaa suorituskykyä, erityisesti tehtävissä, jotka vaativat pitkän kontekstin käsittelyä.

3. Mixture-of-Experts (MoE) -moduulit

MoE-moduuli Jambassa esittää joustavan lähestymistavan mallin kapasiteetin skaalauttamiseen. MoE sallii mallin lisätä käytettävissä olevien parametrejensa määrää ilman, että aktiiviset parametrit kasvavat suhteettomasti johtuen. Jambassa MoE on sovellettu joissakin MLP-kerroksissa, ja reititysmekanismi valitsee aktivoitavat asiantuntijat kunkin tokenin kohdalla. Tämä valikoiva aktivoituminen mahdollistaa Jamballe ylläpitää korkeaa tehokkuutta käsitellessään monimutkaisia tehtäviä.

Alempana oleva kuva havainnollistaa induktio-pään toimintaa hybridisessä Attention-Mamba-mallissa, joka on Jamban avainominaisuus. Tässä esimerkissä huomio-pää on vastuussa tunnisteen, kuten “Positiivinen” tai “Negatiivinen”, ennustamisesta mielipidetutkimustehtävissä. Korostetut sanat osoittavat, miten mallin huomio on vahvasti keskittynyt tunnistelabel-tokeneihin harvojen esimerkkien kohdalla, erityisesti ennen lopputunnisteen ennustamista. Tämä huomio-mekanismi on tärkeä osa mallin kyvyssä suorittaa kontekstissä oppimista, jossa malli on pääteltävä sopiva tunniste annetun kontekstin ja harvojen esimerkkien perusteella.

Suorituskyvyn parantaminen, jota MoE:n integrointi Attention-Mamba-hybridirakenteeseen tarjoaa, korostuu taulukossa. Käyttämällä MoE:ta Jamba lisää kapasiteettiaan ilman, että laskentakustannukset kasvavat suhteettomasti. Tämä on erityisen näkyvää suorituskyvyn merkittävässä parantumisessa erilaisilla benchmark-kojeilla, kuten HellaSwag, WinoGrande ja Natural Questions (NQ). MoE-malli saavuttaa korkeamman tarkin (esim. 66,0 % WinoGrandessa verrattuna 62,5 %:iin ilman MoE:ta) ja osoittaa parannettuja log-probabiliteettejä eri aloilla (esim. -0,534 C4:llä).

Avainarkkitehtuurin ominaisuudet

  • Kerrosten koostumus: Jamban arkkitehtuuri koostuu blokeista, jotka yhdistävät Mamba– ja Transformer-kerrokset tiettyssä suhteessa (esim. 1:7, tarkoittaen yhtä Transformer-kerrosta jokaista seitsemää Mamba-kerrosta). Tämä suhde on säätelty optimaalista suorituskykyä ja tehokkuutta varten.
  • MoE-integrointi: MoE-kerrokset sovelletaan joka muutamalla kerroksella, ja niissä on 16 asiantuntijaa, joista kaksi asiantuntijaa aktivoituu kunkin tokenin kohdalla. Tämä konfiguraatio mahdollistaa Jamballe skaalautumisen tehokkaasti hallitsemalla muistin käytön ja laskentatehokkuuden välistä tasapainoa.
  • Normalisointi ja vakaus: Varmistamaan vakauden aikana koulutusta varten Jamba sisältää RMSNormin Mamba-kerroksissa, mikä auttaa lieventämään ongelmia, kuten suuria aktivaatiopiikkejä, jotka voivat esiintyä skaalautumisen aikana.

Jamban suorituskyky ja benchmarking

Jamba on testattu laajasti eri benchmark-kojeilla, osoittaen kilpailukykyisen suorituskyvyn yleisesti. Seuraavat kohdat korostavat joitakin avainbenchmark-kojeita, joissa Jamba on menestynyt, osoittaen sen vahvuudet sekä yleisissä NLP-tehtävissä että pitkän kontekstin tilanteissa.

1. Yleiset NLP-benchmarkit

Jamba on arvioitu useilla akateemisilla benchmark-kojeilla, mukaan lukien:

  • HellaSwag (10-shot): Yleinen järjen käyttöön perustuva tehtävä, jossa Jamba saavutti 87,1 %:n suorituskyvyn, jättäen useat kilpailevat mallit taakse.
  • WinoGrande (5-shot): Toinen päättelytehtävä, jossa Jamba saavutti 82,5 %:n suorituskyvyn, osoittaen jälleen kykynsä käsitellä monimutkaisia kielellisiä päättelyjä.
  • ARC-Challenge (25-shot): Jamba osoitti vahvan suorituskyvyn 64,4 %:n tuloksella, heijastaen kykyä hallita haastavia monivalintakysymyksiä.

Yhteenvetona MMLU-benchmarkissa (5-shot) Jamba saavutti 67,4 %:n tuloksen, osoittaen vahvuuden moninaisten tehtävien hallinnassa.

2. Pitkän kontekstin arviointi

Yksi Jamban erityisominaisuuksista on kyky käsitellä erittäin pitkiä konteksteja. Malli tukee jopa 256K tokenin kontekstipituutta, mikä on pisin julkaistuista malleista. Tämä kyky testattiin Needle-in-a-Haystack-benchmarkilla, jossa Jamba osoitti poikkeuksellista hakutarkin tarkinuttaa eri kontekstipituuksilla, mukaan lukien jopa 256K tokenia.

3. Läpipääsy ja tehokkuus

Jamban hybridirakenteinen arkkitehtuuri parantaa merkittävästi läpipääsyä, erityisesti pitkissä sekvensseissä.

Vertailukokeissa, joissa verrattiin läpipääsyä (tokenia sekunnissa) eri mallien välillä, Jamba ylitti johdonmukaisesti kilpailijansa, erityisesti suurten batch-kokojen ja pitkien kontekstien tilanteissa. Esimerkiksi 128K tokenin kontekstilla Jamba saavutti kolme kertaa Mixtralin läpipääsyn.

Jamban käyttäminen: Python

Kehittäjille ja tutkijoille, jotka haluavat kokeilla Jambaa, AI21 Labs on tarjonnut mallin Hugging Face -alustalla, mikä mahdollistaa sen käytön laajassa sovellusvalikoimassa. Seuraava koodinpätkä osoittaa, miten Jambaa voidaan ladata ja käyttää tekstien generoimiseen:


<p>from transformers import AutoModelForCausalLM, AutoTokenizer</p>

<p>model = AutoModelForCausalLM.from_pretrained(&quot;ai21labs/Jamba-v0.1&quot;)
tokenizer = AutoTokenizer.from_pretrained(&quot;ai21labs/Jamba-v0.1&quot;)</p>

<p>input_ids = tokenizer(&quot;Viimeisimmässä Super Bowl LVIII:ssa,&quot;, return_tensors=&#039;pt&#039;).to(model.device)[&quot;input_ids&quot;]</p>

<p>outputs = model.generate(input_ids, max_new_tokens=216)</p>

print(tokenizer.batch_decode(outputs))

Tämä yksinkertainen skripti lataa Jamban mallin ja tokenisaattorin, generoi tekstiä annetun syötteen perusteella ja tulostaa generoidun tuloksen.

Jamban hienosäätö

Jamba on suunniteltu perusmalliksi, mikä tarkoittaa, että sitä voidaan hienosäätää tiettyihin tehtäviin tai sovelluksiin. Hienosäätö mahdollistaa mallin mukauttamisen erityisiin aloihin, parantaen suorituskykyä erityisissä tehtävissä. Seuraava esimerkki osoittaa, miten Jambaa voidaan hienosäätää PEFT-kirjaston avulla:

import torch
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments

<p>tokenizer = AutoTokenizer.from_pretrained(&quot;ai21labs/Jamba-v0.1&quot;)
model = AutoModelForCausalLM.from_pretrained(
&quot;ai21labs/Jamba-v0.1&quot;, device_map=&#039;auto&#039;, torch_dtype=torch.bfloat16)</p>

<p>lora_config = LoraConfig(r=8,
target_modules=[
&quot;embed_tokens&quot;,&quot;x_proj&quot;, &quot;in_proj&quot;, &quot;out_proj&quot;, # mamba
&quot;gate_proj&quot;, &quot;up_proj&quot;, &quot;down_proj&quot;, # mlp
&quot;q_proj&quot;, &quot;k_proj&quot;, &quot;v_proj&quot;
# attention],
task_type=&quot;CAUSAL_LM&quot;, bias=&quot;none&quot;)</p>

<p>dataset = load_dataset(&quot;Abirate/english_quotes&quot;, split=&quot;train&quot;)
training_args = SFTConfig(output_dir=&quot;./results&quot;,
num_train_epochs=2,
per_device_train_batch_size=4,
logging_dir=&#039;./logs&#039;,
logging_steps=10, learning_rate=1e-5, dataset_text_field=&quot;quote&quot;)
trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args,
peft_config=lora_config, train_dataset=dataset,
)
trainer.train()

Tämä koodinpätkä hienosäätää Jambaa englanninkielisten sitaattien tietokannassa, sovittamalla mallin parametrejä paremmin tiettyyn tehtävään.

Käyttöönotto ja integrointi

AI21 Labs on tehnyt Jamban laajasti saataville eri alustoilla ja käyttöönottovaihtoehdoilla:

  1. Pilvi-alustat:
    • Saatavilla suurimmilla pilvi-palveluntarjoajilla, kuten Google Cloud Vertex AI, Microsoft Azure ja NVIDIA NIM (NVDA ).
    • Tulossa pian Amazon Bedrockiin, Databricks Marketplaceen ja Snowflake Cortexiin.
  2. AI-kehityskehykset:
    • Integroitu suosittuihin kehyksiin, kuten LangChainiin ja LlamaIndexiin (tulossa).
  3. AI21 Studio:
    • Suora pääsy AI21:n omalta kehitysalustalta.
  4. Hugging Face:
    • Mallit ovat ladattavissa ja kokeiltavissa.
  5. Paikallinen käyttöönotto:
    • Vaihtoehdot yksityisille, paikallisille käyttöönotoille organisaatioille, joilla on tiettyjä turvallisuuden tai määräystenmukaisuuden vaatimuksia.
  6. Räätälöidyt ratkaisut:
    • AI21 tarjoaa räätälöityjä mallin mukauttamis- ja hienosäätöpalveluita yritysasiakkaille.

Kehittäjäystävälliset ominaisuudet

Jamban mallit tulevat useiden sisäänrakennettujen ominaisuuksien kanssa, jotka tekevät niistä erityisen houkuttelevia kehittäjille:

  1. Funktiokutsu: Helppo integrointi ulkoisia työkaluja ja API:ja AI-työvirtoihin.
  2. Rakenteellinen JSON-tuloste: Generoi puhdas, parsittavissa oleva tietorakenteet suoraan luonnollisen kielen syötteistä.
  3. Asiakirjan objekti-imuron: Tehokas prosessointi monimutkaisia asiakirjojen rakenteita.
  4. RAG-optimoit: Sisäänrakennetut ominaisuudet parantamaan hakua ja generointia.

Nämä ominaisuudet, yhdistettynä mallin pitkään konteksti-ikkunaan ja tehokkaaseen prosessointiin, tekevät Jambasta monikäyttöisen työkalun laajalle kehittäjien sovellusalueelle.

Eettiset huomioonotot ja vastuullinen AI

Vaikka Jamban kyvyt ovat vaikuttavia, on tärkeää lähestyä sen käyttöä vastuullisen AI-mielentilalla. AI21 Labs korostaa useita tärkeitä seikkoja:

  1. Perusmallin luonne: Jamba 1.5 -mallit ovat esikoulutettuja perusmalleja ilman tiettyjä suuntautumisia tai ohjausjärjestelyjä.
  2. Viittauksien puute: Malleissa ei ole sisäänrakennettuja valvontamekanismeja.
  3. Huolellinen käyttöönotto: Lisäsovittelu ja valvontamekanismit tulisi toteuttaa ennen mallin käyttöönottoa tuotantoympäristöissä tai loppukäyttäjien kanssa.
  4. Tietosuojelu: Pilvipohjaisissa käyttöönotoissa on oltava tietoinen tietojen käsittelystä ja määräystenmukaisuusvaatimuksista.
  5. Harhasuhteiden tiedostaminen: Kuten kaikki suuret kielenmallit, Jamba voi heijastaa harhasuhteita, jotka ovat läsnä sen koulutusaineistossa. Käyttäjien tulisi olla tietoisia tästä ja toteuttaa asianmukaiset lieventämisstrategiat.

Pitämällä nämä tekijät mielessä, kehittäjät ja organisaatiot voivat hyödyntää Jamban kykyjä vastuullisesti ja eettisesti.

Uusi luku AI-kehityksessä?

Jamban esittely AI21 Labsilta merkitsee merkittävää askelta suurten kielenmallien kehityksessä. Yhdistämällä Transformer- ja tila-avaruusmallien vahvuudet, integroimalla Mixture-of-Experts-tekniikoita ja puskiemalla kontekstin pituuden ja prosessoinnin nopeuden rajoja, Jamba avaa uusia mahdollisuuksia AI-sovelluksille eri aloilla.

Kun AI-yhteisö jatkaa tutkimista ja rakentamista tämän innovatiivisen arkkitehtuurin ympärillä, voimme odottaa edelleen kehitystä mallien tehokkuudessa, pitkän kontekstin ymmärtämisessä ja käytännön AI-käyttöönotossa. Jamban perhe edustaa ei vain uutta mallisarjaa, vaan potentiaalista muutosta siinä, miten lähestymme suurten mittakaavojen AI-järjestelmien suunnittelua ja toteutusta.

Olen viettänyt viimeiset viisi vuotta uppoutumalla kiinnostavaan koneoppimisen ja syvän oppimisen maailmaan. Minun intohimoni ja asiantuntemukseni ovat johtaneet minun osallistumiseen yli 50:een monipuoliseen ohjelmistosuunnitteluhankkeeseen, joissa on erityisesti painottunut AI/ML. Minun jatkuva uteliaisuuteni on myös ohjannut minun luontaisen kielen prosessoinnin pariin, jota haluan tutkia tarkemmin.