Training Compute-Optimal Vision Transformers for Brain Encoding
作者: Sana Ahmadi, Francois Paugam, Tristan Glatard, Pierre Lune Bellec
分类: eess.IV, cs.CV, cs.LG, q-bio.NC
发布日期: 2024-10-17
💡 一句话要点
针对大脑编码，研究计算量最优的视觉Transformer训练策略
🎯 匹配领域: 支柱八：物理动画 (Physics-based Animation)
关键词: 大脑编码 视觉Transformer 模型缩放 数据缩放 计算优化 VideoGPT Ridge回归
📋 核心要点
- 大脑编码模型的训练受限于模型大小、数据规模和计算资源，如何平衡三者是核心问题。
- 利用VideoGPT提取视频时空特征，结合Ridge模型预测大脑活动，探索不同模型和数据规模的影响。
- 实验表明，增加隐藏层维度和数据规模能有效提升大脑编码性能，16位精度训练可加速训练过程。
📝 摘要（中文）
本研究探讨了视觉Transformer在大脑编码中的最优训练策略，重点关注数据规模、模型规模和计算资源的影响。研究使用VideoGPT提取视频中的时空特征，并训练Ridge模型以基于这些特征预测大脑活动。通过改变数据规模（1万、10万、100万、600万）和GPT-2模型的配置（包括隐藏层维度、层数和注意力头数）进行基准实验。同时评估了32位和16位浮点表示的训练效果。结果表明，增加隐藏层维度能显著提高大脑编码性能。数据规模的扩大比模型规模的扩大更能提升大脑编码性能。此外，使用16位精度训练在保持与32位精度相同的大脑编码准确率的同时，将训练时间缩短了1.17倍。
🔬 方法详解
问题定义：论文旨在解决大脑编码任务中，如何高效训练视觉Transformer模型的问题。现有方法在模型大小、数据规模和计算资源之间缺乏有效的权衡，导致训练效率低下或性能不佳。
核心思路：论文的核心思路是通过系统性地研究数据规模、模型规模和计算精度对大脑编码性能的影响，找到计算量最优的训练策略。通过调整GPT-2模型的不同参数，并结合不同规模的训练数据，来评估它们对大脑编码效果的影响。
技术框架：整体框架包括以下几个主要步骤：1) 使用VideoGPT从视频中提取时空特征；2) 使用不同配置的GPT-2模型（调整隐藏层维度、层数和注意力头数）学习这些特征；3) 使用Ridge回归模型将学习到的特征映射到大脑活动；4) 使用Pearson相关系数评估大脑编码的性能。
关键创新：论文的关键创新在于系统性地研究了数据规模、模型规模和计算精度对大脑编码性能的综合影响，并量化了它们各自的贡献。此外，论文还验证了使用16位浮点精度进行训练的可行性，从而加速了训练过程。
关键设计：关键设计包括：1) 使用VideoGPT作为特征提取器，以获得高效的时空特征表示；2) 针对GPT-2模型，重点调整了隐藏层维度、层数和注意力头数等关键参数；3) 使用Ridge回归模型作为大脑活动预测器；4) 采用Pearson相关系数作为大脑编码性能的评估指标；5) 比较了32位和16位浮点精度对训练时间和性能的影响。
🖼️ 关键图片
📊 实验亮点
实验结果表明，增加隐藏层维度能显著提高大脑编码性能，数据规模的扩大比模型规模的扩大更能提升大脑编码性能。使用600万数据集训练的模型取得了最高的Pearson相关系数。此外，使用16位精度训练在保持与32位精度相同的大脑编码准确率的同时，将训练时间缩短了1.17倍，验证了其在高计算性能任务中的有效性。
🎯 应用场景
该研究成果可应用于神经科学领域，帮助研究人员更有效地构建大脑编码模型，从而更好地理解大脑的工作机制。此外，该研究提出的计算量优化策略也可推广到其他视觉任务中，提升模型的训练效率和性能。未来，该研究或可用于开发个性化的大脑解码模型，用于辅助诊断和治疗神经系统疾病。
📄 摘要（原文）
The optimal training of a vision transformer for brain encoding depends on three factors: model size, data size, and computational resources. This study investigates these three pillars, focusing on the effects of data scaling, model scaling, and high-performance computing on brain encoding results. Using VideoGPT to extract efficient spatiotemporal features from videos and training a Ridge model to predict brain activity based on these features, we conducted benchmark experiments with varying data sizes (10k, 100k, 1M, 6M) and different model configurations of GPT-2, including hidden layer dimensions, number of layers, and number of attention heads. We also evaluated the effects of training models with 32-bit vs 16-bit floating point representations. Our results demonstrate that increasing the hidden layer dimensions significantly improves brain encoding performance, as evidenced by higher Pearson correlation coefficients across all subjects. In contrast, the number of attention heads does not have a significant effect on the encoding results. Additionally, increasing the number of layers shows some improvement in brain encoding correlations, but the trend is not as consistent as that observed with hidden layer dimensions. The data scaling results show that larger training datasets lead to improved brain encoding performance, with the highest Pearson correlation coefficients observed for the largest dataset size (6M). These findings highlight that the effects of data scaling are more significant compared to model scaling in enhancing brain encoding performance. Furthermore, we explored the impact of floating-point precision by comparing 32-bit and 16-bit representations. Training with 16-bit precision yielded the same brain encoding accuracy as 32-bit, while reducing training time by 1.17 times, demonstrating its efficiency for high-performance computing tasks.