在单张3090显卡上训练出高效高精度CLIP模型
Simplifying CLIP: Unleashing the Power of Large-Scale Models on Consumer-level Computers
- 简化Transformer结构，结合权重继承与多阶段知识蒸馏
- 用合成标签增强数据，设计配对匹配损失提升性能
- 适合资源有限的研究者在消费级设备上复现大模型
对比语言-图像预训练（CLIP）因其出色的零样本性能和良好的迁移能力受到广泛关注。然而，训练这类大规模模型通常需要大量计算和存储资源，对普通用户构成障碍。本文旨在仅使用一张Nvidia RTX3090 GPU和1TB存储空间，实现具有竞争力的性能。一方面，我们简化了Transformer模块结构，并结合权重继承与多阶段知识蒸馏（WIKD），降低参数量并提升训练与部署时的推理速度；另一方面，针对小规模数据集带来的收敛挑战，我们为每个样本生成合成描述作为数据增强，并设计了一种新的配对匹配（PM）损失，充分挖掘正负图像-文本对之间的区分性。大量实验表明，所提模型在数据规模-参数量-准确率之间实现了新的最优权衡，有望进一步推动CLIP模型在研究社区中的普及。
原文摘要 · Abstract (English)
Contrastive Language-Image Pre-training (CLIP) has attracted a surge of attention for its superior zero-shot performance and excellent transferability to downstream tasks. However, training such large-scale models usually requires substantial computation and storage, which poses barriers for general users with consumer-level computers. Motivated by this observation, in this paper we investigate how to achieve competitive performance on only one Nvidia RTX3090 GPU and with one terabyte for storing dataset. On one hand, we simplify the transformer block structure and combine Weight Inheritance with multi-stage Knowledge Distillation (WIKD), thereby reducing the parameters and improving the inference speed during training along with deployment. On the other hand, confronted with the convergence challenge posed by small dataset, we generate synthetic captions for each sample as data augmentation, and devise a novel Pair Matching (PM) loss to fully exploit the distinguishment among positive and negative image-text pairs. Extensive experiments demonstrate that our model can achieve a new state-of-the-art datascale-parameter-accuracy tradeoff, which could further popularize the CLIP model in the related research community.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。