Yi: Open Foundation Models by 01.AI

Yi: Open Foundation Models by 01.AI.

01.AI

Code: https://github.com/01-ai/Yi

Model: https://huggingface.co/01-ai

We introduce the Yi model family, a series of language and multimodal models with strong multi-dimensional capabilities. Yi is based on 6B and 34B pretrained language models, then extended to chat, 200K long context, depth-upscaled, and vision-language variants. Base models obtain strong performance on benchmarks like MMLU, and chat models obtain strong human preference rate on AlpacaEval and Chatbot Arena.

Built on a scalable super-computing infrastructure and the classical Transformer architecture, we attribute Yi's performance primarily to data quality from our data-engineering efforts. For pretraining we construct 3.1T tokens of English-Chinese corpora through a cascaded deduplication and quality-filtering pipeline. For finetuning we polish a small (<10K) instruction set over many iterations, every instance verified by our ML engineers. We extend context length to 200K through lightweight continual pretraining, and we further show that depth-upscaling continual pretraining further improves performance.

Our vision for large language models is to make them the next-generation computational platform. Yi-6B and Yi-34B are pretrained from scratch on 3.1T highly-engineered tokens and then finetuned on a small but meticulously polished alignment set. Due to data quality from substantial engineering, Yi achieves near GPT-3.5 benchmark scores and human preference rate.

Design choices along model scale, data scale, and data quality: (1) we want a model small enough to run inference on consumer GPUs like the RTX 4090 with 24GB memory, yet large enough to exhibit complex reasoning and emergent abilities — 34B is the sweet spot. (2) Since 34B is smaller than Chinchilla / LLaMA's 70B, we increase pretraining tokens to 3.1T to compensate; this puts us in the post-Chinchilla-optimal regime where the model is overtrained on more tokens than compute-optimal, with the upside being lower serving cost: after int4 quantization the 34B chat model fits in 24GB GPU memory with almost no performance drop. (3) Our data principle is quality over quantity for both pretraining and finetuning. (4) For finetuning we handcraft <10K instructions over multiple iterations from user feedback, deviating from FLAN / UltraChat-style scaling and aligning with LIMA-style handcrafting.

Our pretraining data cleaning system features a cascaded filtering pipeline based on language, heuristic textual features, perplexity, semantics, topic, and safety, together with cascaded deduplication based on paragraph MinHash and sub-document exact matching. The resulting removal ratio is significantly higher than existing pipelines like CCNet, RefinedWeb, and RedPajama. Architecturally we use a standard Transformer with Grouped-Query Attention (GQA), SwiGLU, and RoPE with adjusted base frequency (RoPE ABF), following the lineage of GPT-3, Chinchilla, LLaMA, Baichuan, and Qwen.

We extend Yi along three dimensions: (i) context scaling, (ii) vision-language adaptation, and (iii) depth-upscaling. To reach 200K context we continue pretraining on ~5B length-upsampled tokens; for vision-language we add a ViT encoder and train multi-stage alignment to the LLM's semantic space; for depth-upscaling we make the model deeper via continual pretraining and confirm consistent gains.

For finetuning, our <10K dataset is annotated directly by our ML engineers and polished over many iterations of user feedback. Thanks to the manageable size, we run extensive grid search to identify optimal data composition, promote diversity, and pick hyperparameters. After 8-bit / 4-bit quantization the final chat model deploys on consumer-grade GPUs with almost no performance drop versus bf16.

Our infrastructure supports full-stack development from pretraining to finetuning to serving. For pretraining we build cross-cloud elastic task scheduling, automatic failure recovery, and topology-aware resource allocation — running jobs against real-time available GPU nodes across clusters with limited switching overhead. For finetuning we build a hierarchical scheduling framework supporting different distributed backends for different models — e.g. Megatron for the policy model and DeepSpeed for the reward model. For inference we use 4-bit model + 8-bit KV cache quantization, combined with PagedAttention and Dynamic Batching.

Empirically, Yi-34B matches GPT-3.5 on both performance and efficiency. On standard benchmarks like MMLU (base) and LMSys ELO Rating (chat), Yi-34B reaches scores comparable to GPT-3.5; after parameter + KV-cache quantization the inference cost lets a wide range of community deploy the model on cost-effective devices.

We design a cascaded data-processing pipeline (Fig. ①) targeting quality and diversity. Starting from Common Crawl web documents, we use the CCNet pipeline for language identification and perplexity scoring, then apply a combination of filtering and deduplication.

Heuristic Rule Filters remove low-quality text via (i) URL / domain / word blocklists and garbled-text filters; (ii) document length, ratios of special symbols and short / consecutive / incomplete lines; (iii) repeated words, n-grams, or paragraphs. Thresholds are set via statistical analysis of large document samples. We also identify and anonymize PII such as emails and phone numbers.

Learned Filters address nuanced cases beyond heuristics. Chinese content from Common Crawl is especially tricky (more pornography / gambling). We integrate four learned scorers: (1) a Perplexity Scorer via KenLM (CCNet-style) discards documents with perplexity well above average; (2) a Quality Scorer classifies pages by similarity to Wikipedia-quality and drops low-quality ones; (3) a Document Coherence Scorer flags incoherent documents made of disparate sentences for segmentation or removal; (4) a Safety Scorer removes toxic content (violence, pornography, political propaganda).

Cluster-based Filters group documents by unsupervised semantic clustering, allowing efficient batch annotation; clusters identified as low-quality (via automatic + manual check) are excluded.

Deduplication. Following Penedo et al. (RefinedWeb), we apply document-level MinHash deduplication together with sub-document exact-match dedup. A topic model further classifies pages (news / ads / knowledge), and we down-sample low-information categories like advertisements to ensure information density. Final composition is shown in Fig. ②.

Tokenization. We use BPE via SentencePiece. Vocabulary size is set to 64,000 to balance computational efficiency and word comprehension. We split numbers into individual digits to ease numeric understanding; rare characters fall back to unicode-byte encoding for fault tolerance. We use an identity tokenizer to avoid forcing all punctuation to half-width. English-prioritizing LLMs commonly add a dummy whitespace prefix; we do not, because the assumption fails for sentences starting with quotation marks and shows no positive effect in Chinese.

Model architecture. Yi uses a modified decoder-only Transformer based on LLaMA's code. Yi-6B: hidden $4096$, $32$ Q-heads, $4$ KV-heads, $32$ layers, $4096$ pretrain seq len, max LR $3 \times 10^{-4}$. Yi-34B: hidden $7168$, $56$ Q-heads, $8$ KV-heads, $60$ layers, $4096$ pretrain seq len, max LR $1.5 \times 10^{-4}$.

Attention. Whereas LLaMA-2 uses GQA only on the 70B variant (full attention on 7B/13B), we apply GQA on both Yi-6B and Yi-34B. GQA splits query heads into $G$ groups sharing one K and V head per group, substantially reducing training and inference cost versus full Multi-Head Attention. We observe no performance degradation from applying GQA to our 6B model.

Activation. We use SwiGLU as the post-attention layer, reducing its activation size from $4h$ to $\tfrac{8}{3}h$ ($h$ = hidden size). This adjustment also compensates for the parameter reduction from GQA, keeping total parameter count comparable to existing 7B / 34B models.

Positional embedding & long context. We use RoPE; following Xiong et al. we adjust its base frequency (RoPE ABF) to support up to 200K context, while the base model itself is trained on 4K context length. To adapt to longer context we continue pretraining on 10B tokens from the pretraining mixture with slightly upsampled long sequences (mostly books). Only 1–2B tokens suffice for the model to converge to low loss across 4K–200K lengths, and a lightweight finetuning further unlocks near-perfect long-context retrieval. We therefore view modeling longer dependency than the pretrain length as an intrinsic capability that post-training simply releases.

Our finetuning emphasizes quality over quantity. Unlike data-intensive approaches like FLAN / UltraChat (millions of entries), we align with LIMA / DEITA — focusing on data selection rather than scaling. With <10K entries we can examine and optimize every data point.

Quality is All You Need. The dataset consists of <10K multi-turn instruction-response dialogs, each constructed and polished across multiple iterations from user feedback. Our preliminary results show that a smaller, manually annotated set beats hundreds of thousands of open-source entries — consistent with Gemini, LLaMA-2, LIMA.

Techniques: (1) for prompt distribution, drawing from WizardLM, we develop compound instructions and progressively evolve their complexity — drastically shrinking SFT data needed; (2) for response formatting we use a LIMA-extended default — introduction–body–conclusion, body typically a bullet list; (3) for chain-of-thought we use a "Step-Back" pattern (Zheng et al.), abstracting to a higher-level solution before tackling the concrete question. To reduce hallucination we verify that response knowledge is within the model and remove items that would induce memorization; to reduce repetition we rewrite repetitive turns often overlooked in SFT data.

Diversity & Mixture. We include a wide spectrum of open-source prompts: QA, creative writing, dialogue, reasoning, math, coding, safety, bilingual capability, etc. Inspired by InsTag we build an instruction-tagging system and a diversity-focused sampler to balance distribution across tags. Following Dong et al. we approximate grid-search over capability ratios in $\{1, \tfrac{1}{2}, \tfrac{1}{4}, \tfrac{1}{8}, \tfrac{1}{16}, \tfrac{1}{32}, \tfrac{1}{64}\}$, guided by validation results and in-house human eval sets.

ChatML format. We use ChatML-style structured format so the model can differentiate among system config, user inputs, and assistant responses.

Training method. Next-word prediction loss; we compute loss only on responses, not on system / user instructions. AdamW with $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$. Sequence length $4096$, batch size $64$, $300$ training steps, constant LR $1 \times 10^{-5}$, weight decay $0.1$, gradient clip $1.0$. NEFTune noise scale: $45$ for Yi-34B-Chat, $5$ for Yi-6B-Chat.

Compute scheduling. Pretraining spans months on thousands of GPUs, so we build an efficient multi-cloud task scheduler for pretrain / SFT / RLHF jobs of different priorities, plus an in-house training framework that can elastic-scale pretrain jobs across node sizes following GPU availability — with all training hyper-parameters scaled seamlessly. For reliability we apply automated node inspection / prediction / labeling for software / hardware errors (tainted nodes auto-quarantined), a task queue with pre-checks and fast automatic recovery, and a user-friendly multi-task console.

Memory & comm efficiency. ZeRO-1 partitions optimizer states across DP ranks; tensor parallel + pipeline parallel are confined within a node to avoid inter-node bottleneck; the 3D-parallel strategy avoids activation checkpointing and minimizes pipeline bubbles. Kernel fusion uses FlashAttention and JIT kernels to cut redundant global-memory access; topology-aware ranking minimizes cross-switch communication in fat-tree topologies.

Finetuning framework. Unlike pretraining, finetuning may require orchestrating multiple models (e.g., DPO, PPO). A typical step: the reference / reward model predicts a batch, then the target model uses that batch to compute loss and update. We thus build a multi-model scheduling framework supporting different distributed backends in a single job — e.g., simultaneously using Megatron for the policy and DeepSpeed for the reward model. For DPO, intermediate ref-model results can be cached and reused so wall-clock cost approaches that of supervised finetuning.

Inference. We use 4-bit weight quantization plus 8-bit KV-cache quantization, with <1% accuracy drop on MMLU / CMMLU. Dynamic batching minimizes response time; PagedAttention improves memory utilization. For long context we implement compute-comm overlap, sequence parallelism, and communication compression to support up to 200K continual pretraining and finetuning. We do not modify the architecture (no sparse / local / sliding-window attention) — the model uses full attention even at 200K.

Safety (RAISE). A Responsible-AI Safety Engine covers pretraining and alignment. In pretraining we filter personal identifiers and reduce sexual / violent / extremist content. In alignment we build a comprehensive safety taxonomy (privacy, hate speech, illegal activity, self-harm, sexual content, mental health, cybersecurity, etc.), curate corresponding datasets, mix them with the dialog SFT data, and add prompts simulating attack scenarios to improve robustness against misuse.

Base model performance. Across benchmarks (MMLU, BBH, C-Eval, CMMLU, Gaokao, commonsense reasoning, reading comprehension, code, math), Yi-34B reaches scores on par with or above existing 70B / 180B open-source models. Yi-34B obtains MMLU 76.3, C-Eval 81.4, CMMLU 83.7, Gaokao 82.8 — surpassing LLaMA-2-70B (69.7 / 50.1 / 53.3 / 23.3) and Falcon-180B (70.4 / 57.8 / 58.0 / 59.0). The 3.1T pretraining tokens (vs. typical $\le$ 2T) yield substantial gains. However, math (MATH) and code (HumanEval / MBPP) still lag; we intentionally limited math / code content in the initial pretrain mix and plan to release math / code-enhanced versions.

Discussions. (i) Gain from model scale: Yi-34B improves more on code / math than on commonsense / reading. (ii) Data quality: smaller-but-cleaner models (Yi-34B, Qwen-14B) often beat larger-but-noisier ones (Falcon-180B). (iii) Gap to GPT-4: open-source still lags GPT-4 / GPT-3.5 on many tasks, but bilingual models (Qwen-14B, Yi-34B) match or surpass GPT-4 on Chinese benchmarks (C-Eval, CMMLU, Gaokao).

In-context learning study. We probe whether a model can infer the coefficients of a linear function $y = w_1 x_1 + \ldots + w_n x_n$ from few-shot demonstrations. To dodge the "emergence is a measurement artifact" debate, we use both a continuous metric ($|y - y^*|$, Fig. ③/④) and a discrete one (exact match, Fig. ⑤/⑥). With $[w_1, w_2] = [1, -1]$, Yi-34B and LLaMA-2-70B do best on exact match (Fig. ⑤). Extending to $[1,1,1,1,1]$, only the largest models (LLaMA-2-70B, Mixtral) reach meaningful exact match (Fig. ⑥) — supporting that inferring complex functions in-context is an emergent ability.

Chat model. Yi-34B-Chat reaches MMLU 73.5, CMMLU 81.3, C-Eval 78.5, BBH 71.7, GSM8K 76.0 (5-shot). Crucially, 4-bit-AWQ keeps performance with only marginal degradation. On the 2023 Hungarian high-school math finals (out-of-distribution to GSM8K, not exported here), Yi-34B-Chat performs strongly on both axes; Yi-6B-Chat is weaker on math — small models likely need more SFT data to "activate" math.

Human evaluations. Yi-34B-Chat scores AlpacaEval 94.08, LMSys Chatbot Arena Elo 1110, SuperClue 71.87 (Dec 21, 2023) — proficient bilingual dialogue, well-aligned to user preferences, though still a gap to GPT-4-Turbo. Fig. ⑦ (AlpacaEval 2.0) and Fig. ⑧ (MT-Bench) show that Yi SFT data scales much faster than UltraChat / UltraChat 200K — slope attributed to data quality.

Our long-context recipe has two lightweight stages: continual pretraining and finetuning. We hold the hypothesis that the base model already has the potential to use information anywhere within 200K — continual pretraining unlocks this capability, evidenced by strong Needle-in-a-Haystack performance; the finetuning stage adapts response style to human instructions.

Continual pretraining. We continue pretraining the full-attention model (no sparse / linear attention) with sequence parallelism and distributed attention. Data mixture: (1) original pretraining data; (2) length-upsampled long-context data, mostly from books; (3) multi-document question-answering synthetic data, where the answer first recites the related paragraph. We train for 5B tokens at 4M batch size — about 100 optimization steps — already enough for strong NIAH performance (Fig. ⑨).

Supervised finetuning. We mix short-context SFT data with long-context document QA. The doc-QA is synthesized: randomly concatenate documents, sample one or more paragraphs, ask a chat model to produce question-answer pairs grounded in that paragraph. Crucially, the answer begins with a recitation or paraphrase of the source paragraph — this encourages retrieval behavior and discourages hallucination by biasing the model to use input information instead of internal knowledge.

Results. Fig. ⑨ shows near-all-green NIAH for Yi-34B-200K (two small red blocks near 170K / 194K depths $\approx 41\!-\!56\%$). Table 6: MMLU after 200K adaptation is essentially unchanged from the 4K base model (Yi-34B 4K: 76.32 vs 200K: 75.56) — i.e., context scaling does not hurt short-context capability.

Yi-VL. Yi-VL-6B and Yi-VL-34B are built on top of Yi-6B-Chat / Yi-34B-Chat. The architecture (Fig. 7 in the original paper; not exported here) has three modules: a ViT image encoder (initialized from CLIP ViT-H/14), a Projection Module (two-layer MLP with LayerNorm) that aligns visual features to text feature space, and the LLM (Yi-Chat). To strengthen bilingual multimodal capability we leverage rich bilingual image-text pairs.

Training has three stages: Stage 1 — train ViT + projection on 100M LAION-400M image-text pairs at $224^2$ resolution to align ViT with the LLM. Stage 2 — scale ViT input to $448^2$, train on 20M LAION-400M + 4.8M from CLLaVA / LLaVAR / Flickr / VQAv2 / RefCOCO / Visual7w. Stage 3 — train the entire model on ~1M curated multimodal chat pairs (GQA / VizWiz / TextCaps / OCR-VQA / Visual Genome / ShareGPT4V / …), with each source capped at 50K to keep data balanced. Stage 1/2: batch $4096$, LR $1\mathrm{e}{-4}$, grad-clip $0.5$, $1$ epoch. Stage 3: batch $256$, LR $2\mathrm{e}{-5}$, grad-clip $1.0$, $2$ epochs. On 128 NVIDIA A100s, total time ~3 days for Yi-VL-6B and ~10 days for Yi-VL-34B. On MMMU at release time, Yi-VL-34B scored 41.6 — second only to GPT-4V (55.7).

Depth Upscaling (Yi-9B). Following Kim et al. (Solar-10.7B), we upscale Yi-6B's 32 layers to a 48-layer Yi-9B by duplicating 16 middle layers (12–28), then continue pretraining the enlarged model. Which layers to duplicate is informed by input/output cosine similarity per layer: high-similarity layers (close to identity) can be duplicated with minimal output disruption (Figure 8 in the original paper; not exported here).

Continual training. ~800B tokens across two stages, ~70% freshly collected and carefully selected; code coverage is increased in the final stage to boost code performance. Constant LR $3\mathrm{e}{-5}$; batch size starts at 4M tokens and is incrementally raised whenever loss plateaus; all other hyperparameters mirror the Yi-6B configuration.

Results (Table 8). Yi-9B improves substantially over Yi-6B: ARC-C 50.3 → 55.6, HellaSwag 74.4 → 76.4, MMLU 63.2 → 68.4, Winogrande 71.3 → 73.0, GSM8K 32.5 → 52.3, MATH 4.6 → 15.9, HumanEval 15.9 → 39.0, MBPP 26.3 → 54.4 — confirming the efficacy of duplicate-and-continue-train.

Yi-34B matches GPT-3.5 in performance, and thanks to 4/8-bit quantization it is deployable on consumer-grade devices — making it ideal for local deployment.

Pretraining takeaways. (1) Training on more tokens than the Chinchilla-optimal point gives clear, consistent gains — recommended for all pretraining teams. We trained on 3.1T tokens and believe the model has not yet saturated. (2) The two most critical factors for pretraining data quality are source (professionally produced text vs. casual social posts) and cleaning details (strength of filtering and deduplication). Since data cleaning is a complex pipeline hard to fully grid-search, there is still room for improvement.

Finetuning takeaway. Heavily iterate on a small amount of data ($\le 10\text{K}$), case by case, over multiple rounds, directly by ML engineers, refined by real user feedback. This clearly deviates from the instruction-scaling approach (FLAN → UltraChat). As our results suggest, reasoning capability — which we view as the core capability for real-world LLM deployment — is strongly correlated with model scale when pretraining data is fixed. Continuing to scale model parameters using thoroughly optimized data should lead to even stronger frontier models in our next versions.

References omitted from this offline reader. The original paper carries a full bibliography of 95 entries.

Key works cited include: GQA (Ainslie et al., 2023) for grouped-query attention; RoPE (Su et al., 2021) and RoPE-ABF (Xiong et al., 2023) for positional encoding and base-frequency adjustment; CCNet (Wenzek et al., 2020), RefinedWeb (Penedo et al., 2023), and RedPajama (Computer, 2023) for pretraining data pipelines that Yi compares its removal ratio against.

Finetuning-related references: LIMA (Zhou et al., 2023) and DEITA (Liu et al., 2023) for the <10K hand-curated SFT philosophy; FLAN (Chung et al., 2022) and UltraChat (Ding et al., 2023) as the scaling counterpoint; WizardLM (Xu et al., 2023) for instruction evolution; Step-Back prompting (Zheng et al.) for CoT formatting; NEFTune (Jain et al., 2023) for noisy-embedding finetuning.

Infrastructure references: FlashAttention (Dao et al., 2022), ZeRO (Rajbhandari et al., 2020), PagedAttention (Kwon et al., 2023) and Megatron / DeepSpeed for training and inference acceleration; AWQ (Lin et al., 2023) for 4-bit weight-only quantization that Yi-34B-Chat relies on for 24GB-GPU deployment.

For the complete bibliography (95 entries) please consult the original paper on arxiv 2403.04652.