TinyLlama是一个由新加坡科技设计大学（SUTD）研究人员推出的小型语言模型。它具有11亿参数量，并在大约3万亿个token上进行了预训练。TinyLlama基于Llama 2的架构和分词器构建，利用了开源社区的各种进展，例如FlashAttention，以实现更好的计算效率。尽管模型尺寸相对较小，但TinyLlama在一系列下游任务中表现出色，明显优于现有的具有相似尺寸的开源语言模型。
TinyLlama的训练数据集主要包含两类：自然语言数据和编程语言数据。自然语言数据源自SlimPajama数据集，这是一个大型的开源语料库，基于RedPajama创建，包含超过1.2万亿个token。编程语言数据则来自Starcoderdata数据集，包括大约2500亿个token，涵盖86种编程语言。
TinyLlama的模型架构与Llama2完全一样，但只有11亿参数，这使得它在体积小巧的同时，适用于需要限制计算和内存占用的多种应用。例如，TinyLlama的4比特量化版本的模型权重只需要550MB的内存，使其适合在边缘设备上运行，如智能手机、物联网设备和嵌入式系统。
此外，TinyLlama项目还提供了即插即用的特性，因为它采用了和Llama2相同的分词器，这意味着它可以在许多基于Llama的开源项目中使用。模型的检查点和代码可以在GitHub上公开获取。
“TinyLlama: An Open-Source Small Language Model”由Peiyuan Zhang、Guangtao Zeng、Tianduo Wang和Wei Lu撰写，介绍了TinyLlama模型，这是一个紧凑的11亿参数语言模型，在约1万亿词元上进行了长达3个epoch的预训练，在一系列下游任务中表现出色，超过了现有同等规模的开源语言模型，其模型检查点和代码已在GitHub上公开。
-
模型概述
- 研究背景：自然语言处理（NLP）的进展主要由语言模型规模的扩大推动，但训练较小模型与更大数据集的潜力尚未充分探索。TinyLlama旨在研究小模型在大规模训练下的表现。
- 模型特点：基于Llama 2的架构和分词器，采用多种优化技术提高计算效率，如FlashAttention、Lit-GPT等，在保持较小规模的同时展现出优异性能。
-
预训练
- 数据：使用SlimPajama和StarCoder训练数据集的混合数据，经过处理后共约9500亿词元，训练约3万亿词元，采样比例约为7:3。
- 架构：基于仅解码器的Transformer架构，包含旋转位置嵌入（RoPE）、预归一化（Pre-norm）、RMSNorm、SwiGLU激活函数和分组查询注意力（Grouped-query Attention）等组件。
- 速度优化：通过Fully Sharded Data Parallel（FSDP）、FlashAttention-2和xFormers等技术提高训练速度，与Pythia - 1.0B和MPT - 1.3B相比，训练效率显著提升。
- 训练过程：基于lit - gpt框架，使用自回归语言建模目标，采用AdamW优化器、余弦退火学习率调度等超参数设置，使用16个A100 - 40G GPU进行训练。
- 版本1.1改进：修复学习率调度和数据加载问题，减少通信开销，将预训练词元数降至2万亿，引入三阶段预训练过程（基础预训练、特定领域持续预训练和冷却阶段），产生三个模型变体（TinyLlama v1.1、TinyLlama v1.1 Math&Code和TinyLlama v1.1 Chinese）。
-
实验结果
- 常识推理任务：在Hellaswag、OpenBookQA等多个常识推理任务中，TinyLlama模型表现优于基线模型，如OPT - 1.3B、Pythia - 1.0B和Pythia - 1.4B，平均得分最高。
- 问题解决任务：使用InstructEval基准评估，在MMLU、BBH、DROP和HumanEval等任务中，TinyLlama比现有模型展现出更好的问题解决能力，其中TinyLlama v1.1 Math&Code在HumanEval和DROP任务中提升显著。
- 中文任务评估：在xwinograd、xstorycloze等中文理解和推理任务中，TinyLlama v1.1 Chinese表现出色，TinyLlama v1.0和TinyLlama v1.1 Math&Code也取得了较好成绩，可能得益于训练数据中约3.5%的Python语料库和约4%的Jupyter语料库包含中文文本。
-
结论与展望：TinyLlama作为开源小模型，有望在移动设备端实现应用，为语言模型研究提供轻量级平台。研究团队通过发布相关信息，旨在推动开源语言模型预训练社区的发展，并为未来研究提供参考。