Modely a platformy AI
Jamba: AI21 Labs’ Nová Hybridní Transformer-Mamba Jazykový Model
Přidejte Unite.AI mezi své preferované zdroje na Google
Jazykové modely prošly rychlým vývojem, přičemž architektury založené na Transformeru vedou v oblasti zpracování přirozeného jazyka. Nicméně, jak modely rostou, tak se problémy související s处理em dlouhých kontextů, efektivitou paměti a propustností stávají stále výraznějšími.
AI21 Labs představilo nové řešení s Jambou, špičkovým velkým jazykovým modelem (LLM), který kombinuje silné stránky obou Transformeru a Mamba architektur v hybridním rámci. Tento článek podrobně popisuje architekturu Jambly, její výkon a potenciální aplikace.
Přehled Jambly
Jamba je hybridní velký jazykový model vyvinutý AI21 Labs, který využívá kombinaci Transformer vrstev a Mamba vrstev, integrovaných s Mixture-of-Experts (MoE) modulem. Tato architektura umožňuje Jambě vyvážit využití paměti, propustnost a výkon, čímž se stává silným nástrojem pro širokou škálu úkolů NLP. Model je navržen tak, aby se vešel do jediného 80GB GPU, nabízí vysokou propustnost a malou stopu paměti, zatímco udržuje špičkový výkon na různých benchmarcích.
Architektura Jambly
Architektura Jambly je základem jejích schopností. Je postavena na novém hybridním designu, který prokládá Transformer vrstvy s Mamba vrstvami, začleněnými MoE moduly, aby zvýšila kapacitu modelu bez významného zvýšení výpočetních požadavků.
1. Transformer Vrstvy
Architektura Transformeru se stala standardem pro moderní LLM, díky své schopnosti efektivně zpracovávat paralelní zpracování a zachycovat dlouhodobé závislosti v textu. Nicméně, její výkon je často omezený vysokými požadavky na paměť a výpočetní zdroje, zejména při zpracování dlouhých kontextů. Jamba řeší tyto omezení integrací Mamba vrstev, které budeme dále zkoumat.
2. Mamba Vrstvy
Mamba je nedávný model prostoru (SSM), navržen pro efektivnější zpracování dlouhých vztahů v sekvencích než tradiční RNN nebo dokonce Transformer. Mamba vrstvy jsou obzvláště efektivní při snižování stopy paměti spojené s ukládáním mezipaměti klíčových hodnot (KV) v Transformerech. Prokládáním Mamba vrstev s Transformer vrstvami snižuje Jamba celkové využití paměti, zatímco udržuje vysoký výkon, zejména v úkolech vyžadujících zpracování dlouhých kontextů.
3. Mixture-of-Experts (MoE) Moduly
MoE modul v Jambě představuje flexibilní přístup ke škálování kapacity modelu. MoE umožňuje modelu zvýšit počet dostupných parametrů bez proporcionálního zvýšení aktivních parametrů během inferencing. V Jambě je MoE aplikován na některé z MLP vrstev, s router mechanismem, který vybírá horní experty pro aktivaci pro každý token. Tato selektivní aktivace umožňuje Jambě udržet vysokou efektivitu, zatímco zpracovává komplexní úkoly.
Níže uvedený obrázek demonstruje funkčnost indukční hlavy v hybridním modelu Attention-Mamba, klíčové funkce Jambly. V tomto příkladu je attention head zodpovědný za předpověď labelů, jako je “Positive” nebo “Negative” v reakci na úkoly sentiment analýzy. Zvýrazněné slova ilustrují, jak je pozornost modelu silně zaměřena na label tokeny z few-shot příkladů, zejména v kritickém okamžiku před předpovědí konečného labelu. Tento mechanismus pozornosti hraje zásadní roli v schopnosti modelu provádět v-kontextové učení, kde model musí odvodit příslušný label na základě daného kontextu a few-shot příkladů.
Vylepšení výkonu nabízená integrací Mixture-of-Experts (MoE) s hybridní architekturou Attention-Mamba jsou zdůrazněna v Tabulce. Používáním MoE zvyšuje Jamba svou kapacitu bez proporcionálního zvýšení výpočetních nákladů. To je zvláště patrné ve významném zvýšení výkonu napříč různými benchmaky, jako je HellaSwag, WinoGrande a Natural Questions (NQ). Model s MoE nejen dosahuje vyšší přesnosti (například 66,0 % na WinoGrande ve srovnání s 62,5 % bez MoE), ale také демонстриuje lepší log-probability napříč různými doménami (například -0,534 na C4).
Klíčové Architektonické Funkce
- Složení Vrstev: Architektura Jambly se skládá z bloků, které kombinují Mamba a Transformer vrstvy v určitém poměru (například 1:7, což znamená jednu Transformer vrstvu pro sedm Mamba vrstev). Tento poměr je optimalizován pro nejlepší výkon a efektivitu.
- Integrace MoE: MoE vrstvy jsou aplikovány každých několik vrstev, s 16 experty dostupnými a dvěma nejlepšími experty aktivovanými pro každý token. Tato konfigurace umožňuje Jambě škálovat efektivně, zatímco spravuje kompromisy mezi využitím paměti a výpočetní efektivitou.
- Normalizace a Stabilita: Pro zajištění stability během trénování Jamba zahrnuje RMSNorm v Mamba vrstvách, což pomáhá zmírnit problémy, jako jsou velké aktivační špičky, které mohou nastat při škálování.
Výkon Jambly a Benchmarking
Jamba prošla náročnými testy proti širokému spektru benchmarkek, demonstrujícím konkurenceschopný výkon napříč všemi. Následující sekce zdůrazňují některé z klíčových benchmarkek, kde Jamba excelovala, ukazujíc její silné stránky v obou obecných úkolech NLP a scénářích s dlouhými kontexty.
1. Obecné Benchmarky NLP
Jamba byla hodnocena na několika akademických benchmarkech, včetně:
- HellaSwag (10-shot): Úkol rozumného uvažování, kde Jamba dosáhla výkonu 87,1 %, překonávající mnoho konkurenčních modelů.
- WinoGrande (5-shot): Další úkol rozumného uvažování, kde Jamba dosáhla 82,5 %, opět ukazujíc její schopnost zpracovávat komplexní lingvistické uvažování.
- ARC-Challenge (25-shot): Jamba demonstrovala silný výkon se skórem 64,4 %, odrážejícím její schopnost zvládat náročné multiple-choice otázky.
V agregovaných benchmarkech, jako je MMLU (5-shot), Jamba dosáhla skóre 67,4 %, ukazujíc její robustnost napříč různými úkoly.
2. Hodnocení Dlouhých Kontextů
Jedním z význačných rysů Jambly je její schopnost zpracovávat extrémně dlouhé kontexty. Model podporuje kontextovou délku až 256K tokenů, což je nejdelší mezi veřejně dostupnými modely. Tato schopnost byla testována pomocí benchmarkek Needle-in-a-Haystack, kde Jamba ukázala výjimečnou přesnost vyhledávání napříč různými kontextovými délkami, včetně až 256K tokenů.
3. Propustnost a Efektivita
Hybridní architektura Jambly významně zlepšuje propustnost, zejména u dlouhých sekvencí.
V testech srovnávajících propustnost (tokeny za sekundu) napříč různými modely Jamba konzistentně překonávala své konkurenty, zejména ve scénářích s velkými dávkami a dlouhými kontexty. Například s kontextem 128K tokenů Jamba dosáhla trojnásobné propustnosti Mixtralu, srovnatelného modelu.
Používání Jambly: Python
Pro vývojáře a výzkumníky, kteří chtějí experimentovat s Jambou, AI21 Labs poskytlo model na platformách, jako je Hugging Face, což jej činí dostupným pro širokou škálu aplikací. Následující kódový úsek demonstruje, jak načíst a generovat text pomocí Jambly:
<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("In the recent 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))
Tento jednoduchý skript načte model Jambly a tokenizér, vygeneruje text na základě zadaného vstupního promptu a vytiskne vygenerovaný výstup.
Jemné Upravování Jambly
Jamba je navržena jako základní model, což znamená, že ji lze jemně upravit pro specifické úkoly nebo aplikace. Jemné upravování umožňuje uživatelům přizpůsobit model pro úzké domény, zlepšujíc tak výkon na specializovaných úkolech. Následující příklad ukazuje, jak jemně upravit Jambu pomocí knihovny 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()
Tento kódový úsek jemně upravuje Jambu na datasetu anglických citátů, upravujíc parametry modelu pro lepší přizpůsobení specifickému úkolu generování textu v úzké doméně.
Nasazení a Integrace
AI21 Labs zpřístupnilo Jambu široké veřejnosti prostřednictvím různých platforem a možností nasazení:
- Cloud Platforms:
- Dostupné na hlavních cloudových poskytovatelích, včetně Google Cloud Vertex AI, Microsoft Azure a NVIDIA NIM (NVDA ).
- Brzy dostupné na Amazon Bedrock, Databricks Marketplace a Snowflake Cortex.
- Rámců pro vývoj AI:
- Integrace s populárními rámci, jako je LangChain a LlamaIndex (příští).
- AI21 Studio:
- Přímý přístup prostřednictvím vývojové platformy AI21.
- Hugging Face:
- Modely dostupné ke stažení a experimentování.
- Místní Nasazení:
- Možnosti pro soukromé, místní nasazení pro organizace s konkrétními bezpečnostními nebo dodržovacími požadavky.
- Vlastní Řešení:
- AI21 nabízí přizpůsobenou úpravu modelu a jemné upravování služeb pro podnikové klienty.
Vývojářské Funkce
Jamba modely přicházejí s několika vestavěnými funkcemi, které je činí zvláště atraktivními pro vývojáře:
- Funkční Volání: Snadno integrujte externí nástroje a API do svých AI pracovních postupů.
- Strukturovaný JSON Výstup: Generujte čisté, analyzovatelné datové struktury přímo z přirozeného jazykového vstupu.
- Trávení Dokumentů: Efektivně zpracovávejte a rozumějte komplexním dokumentovým strukturám.
- Optimalizace RAG: Vestavěné funkce pro zlepšení pipeline retrieval-augmented generace.
Tyto funkce, v kombinaci s dlouhou kontextovou oknem a efektivním zpracováním, činí Jambu univerzálním nástrojem pro širokou škálu vývojových scénářů.
Etické Úvahy a Zodpovědné AI
Zatímco schopnosti Jambly jsou působivé, je důležité přistupovat k jejímu použití se zodpovědným AI přístupem. AI21 Labs zdůrazňuje několik důležitých bodů:
- Příroda Základní Modelu: Jamba 1.5 modely jsou předtrénované základní modely bez specifické orientace nebo instrukční úpravy.
- Chybějící Vestavěné Bezpečnostní Mechanismy: Modely nemají v sobě vestavěné moderátorské mechanismy.
- Přičinlivé Nasazení: Další úpravy a bezpečnostní mechanismy by měly být implementovány před použitím Jambly v produkčních prostředích nebo s koncovými uživateli.
- Ochrana Dat: Při použití cloudových nasazení buďte si vědomi zpracování a dodržování požadavků na data.
- Povědomí o Zaujatosti: Jako všechny velké jazykové modely, Jamba může odrážet zaujatosti přítomné ve svém tréninkovém datu. Uživatelé by si měli být vědomi tohoto a implementovat příslušná zmírnění.
Při zachování těchto faktorů mohou vývojáři a organizace využít schopnosti Jambly zodpovědně a eticky.
Nová Kapitola ve Vývoji AI?
Představení rodiny Jambly AI21 Labs představuje významný milník ve vývoji velkých jazykových modelů. Kombinací silných stránek Transformeru a state space modelů, integrací technik mixture of experts a posunem hranic kontextové délky a zpracování rychlosti, Jamba otevírá nové možnosti pro AI aplikace napříč průmysly.
Jakmile AI komunita bude pokračovat ve výzkumu a budování na této inovativní architektuře, můžeme očekávat další pokroky ve výkonnosti modelu, dlouhém kontextovém chápání a praktickém nasazení AI. Rodina Jambly reprezentuje nejen novou sadu modelů, ale potenciální posun v přístupu k návrhu a implementaci velkých AI systémů.
