Contrastors: 高效的对比学习工具包

Contrastors是一个功能强大的对比学习工具包,旨在帮助研究人员和工程师高效地训练和评估对比模型。作为一个开源项目,Contrastors为对比学习领域的研究和应用提供了坚实的基础。

Contrastors示例图

主要特性

Contrastors具有以下突出特点:

基于Flash Attention构建,实现快速高效的训练
支持多GPU训练,充分利用硬件资源
集成GradCache,可在有限内存环境下进行大批量训练
支持Huggingface模型,轻松加载常用模型如Pythia/GPTNeoX、BERT等
支持掩码语言建模(MLM)预训练
实现Matryoshka表示学习,提供灵活的嵌入大小
支持CLIP和LiT风格的对比学习
可加载流行的ViT模型(如timm)

这些特性使Contrastors成为进行对比学习研究和应用的理想工具。无论是学术研究还是工业应用,Contrastors都能提供强大的支持。

研究成果

Contrastors项目已经产出了一些重要的研究成果:

Nomic Embed: 训练可复现的长上下文文本嵌入器 - 由Zach Nussbaum、Jack Morris、Andriy Mulyar和Brandon Duderstadt完成。
Nomic Embed Vision: 扩展潜在空间 - 由Zach Nussbaum、Brandon Duderstadt和Andriy Mulyar完成。

这些研究展示了Contrastors在实际应用中的潜力和效果。

快速入门

要开始使用Contrastors,您需要先设置环境并安装必要的依赖:

确保您的系统安装了CUDA 11.8+。
创建并激活Python虚拟环境。
安装PyTorch。
安装Flash Attention及其自定义内核。
安装其他依赖项和Contrastors包。

详细的安装步骤请参考项目README。

数据访问

Contrastors提供了对nomic-embed-text-v1数据集的访问。您需要在atlas.nomic.ai创建账户,然后使用nomic包登录并获取访问凭证。

训练模型

Contrastors支持多种训练模式:

掩码语言建模预训练
对比预训练和微调
视觉模型对齐

每种模式都有详细的训练命令和配置说明。

预训练模型

Contrastors提供了多个预训练模型,包括:

nomic-embed-text-v1
nomic-embed-vision-v1
nomic-embed-text-v1.5
nomic-embed-vision-v1.5
等等

这些模型可以在Hugging Face上找到并直接使用。

社区参与

Contrastors是一个开放的社区项目,欢迎更多人参与:

许可证

Contrastors采用Apache 2.0许可证。各个模型可能有单独的许可证,请查看相应的模型卡片。

致谢

Contrastors项目要感谢Tri Dao在Flash Attention上的工作,OpenCLIP团队的优秀仓库,以及Huggingface团队在transformers库上的杰出工作。

引用

如果您在研究中使用了Contrastors,请引用以下论文:

@misc{nussbaum2024nomic,
      title={Nomic Embed: Training a Reproducible Long Context Text Embedder}, 
      author={Zach Nussbaum and John X. Morris and Brandon Duderstadt and Andriy Mulyar},
      year={2024},
      eprint={2402.01613},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

@misc{nussbaum2024nomicembedvisionexpanding,
      title={Nomic Embed Vision: Expanding the Latent Space}, 
      author={Zach Nussbaum and Brandon Duderstadt and Andriy Mulyar},
      year={2024},
      eprint={2406.18587},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2406.18587}, 
}

Contrastors为对比学习研究和应用提供了一个强大而灵活的工具包。无论您是研究人员还是工程师,Contrastors都能帮助您更高效地进行对比学习相关工作。我们期待看到更多基于Contrastors的创新研究和应用!