Modelli e piattaforme di IA
Jamba: il nuovo modello linguistico ibrido di AI21 Labs che combina Transformer e Mamba
Aggiungi Unite.AI alle tue fonti preferite su Google
I modelli linguistici hanno assistito a rapidi progressi, con architetture basate su Transformer che guidano il processo nel campo dell’elaborazione del linguaggio naturale. Tuttavia, man mano che i modelli crescono, le sfide relative alla gestione di contesti lunghi, all’efficienza della memoria e alla velocità di elaborazione sono diventate più pronunciate.
AI21 Labs ha introdotto una nuova soluzione con Jamba, un modello linguistico di grandi dimensioni (LLM) di ultima generazione che combina le forze di entrambe le architetture Transformer e Mamba in un framework ibrido. Questo articolo descrive l’architettura di Jamba, le sue prestazioni e le sue potenziali applicazioni.
Panoramica di Jamba
Jamba è un modello linguistico di grandi dimensioni ibrido sviluppato da AI21 Labs, che sfrutta una combinazione di layer Transformer e layer Mamba, integrati con un modulo Mixture-of-Experts (MoE). Questa architettura consente a Jamba di bilanciare l’utilizzo della memoria, la velocità di elaborazione e le prestazioni, rendendolo uno strumento potente per una vasta gamma di compiti di elaborazione del linguaggio naturale. Il modello è progettato per adattarsi a una singola GPU da 80 GB, offrendo una velocità di elaborazione elevata e un’impronta di memoria ridotta, mantenendo al contempo prestazioni all’avanguardia in vari benchmark.
L’architettura di Jamba
L’architettura di Jamba è il fondamento delle sue capacità. È costruita su un design ibrido innovativo che alterna layer Transformer con layer Mamba, incorporando moduli MoE per aumentare la capacità del modello senza aumentare significativamente le richieste computazionali.
1. Layer Transformer
L’architettura Transformer è diventata lo standard per i moderni LLM a causa della sua capacità di gestire l’elaborazione parallela in modo efficiente e di catturare le dipendenze a lungo raggio nel testo. Tuttavia, le sue prestazioni sono spesso limitate dalle elevate richieste di memoria e di calcolo, in particolare quando si elaborano contesti lunghi. Jamba affronta queste limitazioni integrando layer Mamba, che esamineremo nel prossimo punto.
2. Layer Mamba
Mamba è un modello di stato recente (SSM) progettato per gestire le relazioni a lunga distanza nelle sequenze in modo più efficiente rispetto ai tradizionali RNN o anche ai Transformer. I layer Mamba sono particolarmente efficaci nel ridurre l’impronta di memoria associata alla memorizzazione delle cache di chiave-valore (KV) nei Transformer. Alternando i layer Mamba con i layer Transformer, Jamba riduce l’utilizzo complessivo della memoria, mantenendo al contempo elevate prestazioni, in particolare nei compiti che richiedono la gestione di contesti lunghi.
3. Mixture-of-Experts (MoE) Modules
Il modulo MoE in Jamba introduce un approccio flessibile per scalare la capacità del modello. MoE consente al modello di aumentare il numero di parametri disponibili senza aumentare proporzionalmente i parametri attivi durante l’inferenza. In Jamba, MoE è applicato ad alcuni dei layer MLP, con il meccanismo di routing che seleziona gli esperti migliori da attivare per ogni token. Questa attivazione selettiva consente a Jamba di mantenere un’elevata efficienza mentre gestisce compiti complessi.
L’immagine seguente dimostra la funzionalità di una testa di induzione in un modello ibrido di attenzione-Mamba, una caratteristica chiave di Jamba. In questo esempio, la testa di attenzione è responsabile della previsione di etichette come “Positivo” o “Negativo” in risposta a compiti di analisi del sentimento. Le parole evidenziate mostrano come l’attenzione del modello sia fortemente focalizzata sui token di etichetta degli esempi di few-shot, in particolare nel momento critico prima della previsione dell’etichetta finale. Questo meccanismo di attenzione svolge un ruolo cruciale nella capacità del modello di eseguire l’apprendimento in contesto, in cui il modello deve inferire l’etichetta appropriata in base al contesto e agli esempi di few-shot forniti.
I miglioramenti delle prestazioni offerti dall’integrazione di Mixture-of-Experts (MoE) con l’architettura ibrida di attenzione-Mamba sono evidenziati nella tabella. Utilizzando MoE, Jamba aumenta la sua capacità senza aumentare proporzionalmente i costi computazionali. Ciò è particolarmente evidente nel significativo aumento delle prestazioni in vari benchmark come HellaSwag, WinoGrande e Natural Questions (NQ). Il modello con MoE non solo raggiunge una maggiore accuratezza (ad esempio, 66,0% su WinoGrande rispetto al 62,5% senza MoE), ma dimostra anche log-probabilità migliorate in diversi domini (ad esempio, -0,534 su C4).
Caratteristiche architettoniche chiave
- Composizione dei layer: L’architettura di Jamba consiste in blocchi che combinano Mamba e layer Transformer in un rapporto specifico (ad esempio, 1:7, che significa un layer Transformer per ogni sette layer Mamba). Questo rapporto è stato ottimizzato per prestazioni ed efficienza.
- Integrazione MoE: I layer MoE sono applicati ogni pochi layer, con 16 esperti disponibili e i primi 2 esperti attivati per token. Questa configurazione consente a Jamba di scalare in modo efficace, gestendo i compromessi tra utilizzo della memoria e efficienza computazionale.
- Normalizzazione e stabilità: Per garantire la stabilità durante l’addestramento, Jamba incorpora RMSNorm nei layer Mamba, il che aiuta a mitigare problemi come picchi di attivazione grandi che possono verificarsi a scala.
Prestazioni e benchmarking di Jamba
Jamba è stato sottoposto a una serie di test approfonditi su una vasta gamma di benchmark, dimostrando prestazioni competitive in tutti i campi. Le sezioni seguenti evidenziano alcuni dei principali benchmark in cui Jamba ha eccelso, mostrando le sue forze sia nei compiti di elaborazione del linguaggio naturale generale che in scenari di contesto lungo.
1. Benchmark di elaborazione del linguaggio naturale comuni
Jamba è stato valutato su diversi benchmark accademici, tra cui:
- HellaSwag (10-shot): un compito di ragionamento comune in cui Jamba ha raggiunto un punteggio di prestazione dell’87,1%, superando molti modelli concorrenti.
- WinoGrande (5-shot): un altro compito di ragionamento in cui Jamba ha raggiunto un punteggio dell’82,5%, mostrando ancora una volta la sua capacità di gestire complesse inferenze linguistiche.
- ARC-Challenge (25-shot): Jamba ha dimostrato una forte prestazione con un punteggio del 64,4%, riflettendo la sua capacità di gestire domande multiple scelte impegnative.
Nei benchmark aggregati come MMLU (5-shot), Jamba ha raggiunto un punteggio del 67,4%, indicando la sua robustezza in compiti diversi.
2. Valutazioni di contesto lungo
Una delle caratteristiche di spicco di Jamba è la sua capacità di gestire contesti estremamente lunghi. Il modello supporta una lunghezza di contesto di fino a 256K token, la più lunga tra i modelli pubblicamente disponibili. Questa capacità è stata testata utilizzando il benchmark Needle-in-a-Haystack, in cui Jamba ha mostrato un’eccezionale accuratezza di recupero in vari contesti di lunghezza, fino a 256K token.
3. Velocità di elaborazione ed efficienza
L’architettura ibrida di Jamba migliora significativamente la velocità di elaborazione, in particolare con sequenze lunghe.
In test di confronto della velocità di elaborazione (token al secondo) tra diversi modelli, Jamba ha costantemente superato i suoi pari, in particolare in scenari con batch di grandi dimensioni e contesti lunghi. Ad esempio, con un contesto di 128K token, Jamba ha raggiunto il triplo della velocità di elaborazione di Mixtral, un modello comparabile.
Utilizzo di Jamba: Python
Per gli sviluppatori e i ricercatori ansiosi di sperimentare con Jamba, AI21 Labs ha reso il modello disponibile su piattaforme come Hugging Face, rendendolo accessibile per una vasta gamma di applicazioni. Il seguente snippet di codice dimostra come caricare e generare testo utilizzando Jamba:
<p>from transformers import AutoModelForCausalLM, AutoTokenizer</p> <p>model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1") tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")</p> <p>input_ids = tokenizer("Nel recente Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]</p> <p>outputs = model.generate(input_ids, max_new_tokens=216)</p> print(tokenizer.batch_decode(outputs))
Questo semplice script carica il modello Jamba e il tokenizer, genera testo in base a un prompt di input fornito e stampa il testo generato.
Aggiornamento di Jamba
Jamba è progettato come modello base, il che significa che può essere aggiornato per compiti o applicazioni specifiche. L’aggiornamento consente agli utenti di adattare il modello a domini di nicchia, migliorando le prestazioni in compiti specializzati. L’esempio seguente mostra come aggiornare Jamba utilizzando la libreria PEFT:
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("ai21labs/Jamba-v0.1") model = AutoModelForCausalLM.from_pretrained( "ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)</p> <p>lora_config = LoraConfig(r=8, target_modules=[ "embed_tokens","x_proj", "in_proj", "out_proj", # mamba "gate_proj", "up_proj", "down_proj", # mlp "q_proj", "k_proj", "v_proj" # attention], task_type="CAUSAL_LM", bias="none")</p> <p>dataset = load_dataset("Abirate/english_quotes", split="train") training_args = SFTConfig(output_dir="./results", num_train_epochs=2, per_device_train_batch_size=4, logging_dir='./logs', logging_steps=10, learning_rate=1e-5, dataset_text_field="quote") trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args, peft_config=lora_config, train_dataset=dataset, ) trainer.train()
Questo snippet di codice aggiorna Jamba su un set di dati di citazioni in inglese, adattando i parametri del modello per adattarsi meglio al compito specifico di generazione di testo in un dominio specializzato.
Distribuzione e integrazione
AI21 Labs ha reso la famiglia Jamba ampiamente accessibile attraverso varie piattaforme e opzioni di distribuzione:
- Piattaforme cloud:
- Disponibile su principali provider cloud, tra cui Google Cloud Vertex AI, Microsoft Azure e NVIDIA NIM (NVDA ).
- In arrivo su Amazon Bedrock, Databricks Marketplace e Snowflake Cortex.
- Framework di sviluppo AI:
- Integrazione con framework popolari come LangChain e LlamaIndex (in arrivo).
- AI21 Studio:
- Accesso diretto attraverso la piattaforma di sviluppo di AI21.
- Hugging Face:
- Modelli disponibili per il download e l’esperimentazione.
- Distribuzione on-premises:
- Opzioni per la distribuzione privata, on-site per le organizzazioni con esigenze di sicurezza o conformità specifiche.
- Soluzioni personalizzate:
- AI21 offre servizi di personalizzazione e aggiornamento del modello per i clienti aziendali.
Funzionalità per gli sviluppatori
I modelli Jamba sono dotati di diverse funzionalità integrate che li rendono particolarmente attraenti per gli sviluppatori:
- Chiamata di funzioni: integra facilmente strumenti e API esterni nei flussi di lavoro di AI.
- Output JSON strutturato: genera strutture di dati pulite e parsabili direttamente da input di linguaggio naturale.
- Elaborazione di oggetti di documento: elabora ed comprende in modo efficiente strutture di documento complesse.
- Ottimizzazioni RAG: funzionalità integrate per migliorare le pipeline di generazione con recupero.
Queste funzionalità, combinate con la finestra di contesto lunga e l’elaborazione efficiente del modello, rendono Jamba uno strumento versatile per una vasta gamma di scenari di sviluppo.
Considerazioni etiche e AI responsabile
Mentre le capacità di Jamba sono impressionanti, è fondamentale approcciare il suo utilizzo con una mentalità di AI responsabile. AI21 Labs sottolinea diversi punti importanti:
- Natura del modello base: i modelli Jamba 1.5 sono modelli base pre-addestrati senza allineamento o adattamento specifico.
- Mancanza di meccanismi di moderazione integrati: i modelli non hanno meccanismi di moderazione innati.
- Distribuzione accurata: si consiglia di implementare adattamenti e meccanismi di sicurezza aggiuntivi prima di utilizzare Jamba in ambienti di produzione o con utenti finali.
- Privacy dei dati: quando si utilizza la distribuzione cloud, è necessario essere consapevoli della gestione dei dati e dei requisiti di conformità.
- Consapevolezza dei pregiudizi: come tutti i grandi modelli linguistici, Jamba può riflettere i pregiudizi presenti nei suoi dati di addestramento. Gli utenti dovrebbero essere consapevoli di ciò e implementare le dovute mitigazioni.
Tenendo presente questi fattori, gli sviluppatori e le organizzazioni possono sfruttare le capacità di Jamba in modo responsabile ed etico.
Un nuovo capitolo nello sviluppo di AI?
L’introduzione della famiglia Jamba da parte di AI21 Labs segna un importante traguardo nell’evoluzione dei grandi modelli linguistici. Combinando le forze di Transformer e modelli di stato, integrando tecniche di mixture of experts e spingendo i limiti della lunghezza del contesto e della velocità di elaborazione, Jamba apre nuove possibilità per le applicazioni di AI in vari settori.
Mentre la comunità di AI continua a esplorare e costruire su questa architettura innovativa, possiamo aspettarci ulteriori progressi nell’efficienza del modello, nella comprensione del contesto lungo e nella distribuzione pratica di AI. La famiglia Jamba rappresenta non solo un nuovo insieme di modelli, ma anche un potenziale cambio di paradigma nel modo in cui progettiamo e implementiamo i grandi sistemi di AI.
