Optax v0.2.5 版本发布:优化器库的重大更新与功能增强
Optax 是 Google DeepMind 开发的一个基于 JAX 的梯度处理和优化库,它提供了丰富的优化算法实现和梯度变换工具。作为 JAX 生态系统中最重要的优化器库之一,Optax 以其简洁的 API 设计和高效的实现赢得了广大机器学习研究者和开发者的青睐。最新发布的 v0.2.5 版本带来了多项重要更新和功能增强,本文将深入解析这些新特性及其技术价值。
核心优化器增强
本次更新中,Optax 对多个核心优化器进行了重要改进:
-
Sophia-H 优化器:新增了 Sophia-H 优化器实现,这是一种二阶优化器,特别适合处理大规模语言模型的训练任务。Sophia-H 通过利用 Hessian 信息的近似,能够更有效地调整学习率,在保持训练稳定性的同时加速收敛。
-
Muon 优化器:贡献模块中新增了 Muon 优化器,这是一种结合了动量与自适应学习率特性的新型优化器。Muon 通过独特的参数更新机制,在传统动量方法基础上引入了额外的缩放因子,能够更好地处理不同参数尺度的梯度变化。
-
AdEMAMix 优化器:新增了 AdEMAMix 及其简化版本,这是一种结合了 EMA(指数移动平均)和自适应矩估计的混合优化器。它特别适合处理非平稳目标函数和噪声较大的优化问题,在强化学习等领域表现出色。
-
L-BFGS 改进:修复了 L-BFGS 优化器对复数支持的问题,现在可以正确处理复数参数的优化问题,扩展了在信号处理等领域的应用场景。
梯度处理与损失函数改进
在梯度处理和损失函数方面,v0.2.5 版本也带来了多项重要更新:
-
三重边际损失修复:修复了三重边际损失(triplet margin loss)的实现问题,确保其数学正确性。这种损失函数在度量学习和嵌入学习中被广泛使用。
-
Sigmoid 交叉熵改进:改进了 sigmoid_binary_cross_entropy 和 sigmoid_focal_loss 函数,现在可以正确处理非数组类型的标签输入,提高了API的健壮性和易用性。
-
梯度集中化增强:简化了 centralize 函数的实现,并改进了其文档说明,使梯度中心化操作更加清晰易懂。梯度中心化是一种有效的训练技巧,能够提高模型的泛化能力。
工具函数与实用功能
-
树操作增强:
- 新增 tree_batch_shape 函数,用于获取参数树的批量形状
- 统一了 tree_l1_norm、tree_l2_norm 和 tree_linf_norm 为 tree_norm 函数
- 增加了 tree_scale 和 tree_add_scale 函数(原 tree_scalar_mul 和 tree_add_scalar_mul)
- 为 tree_clip 函数添加了默认的最小值和最大值参数
-
参数冻结工具:新增了 freeze 和 selective_optimizer 等实用工具,支持通过前缀掩码选择性地冻结模型参数,为迁移学习和微调提供了便利。
-
投影操作改进:优化了 projections 模块的实现,包括非负最小二乘(NNLS)求解器等,为约束优化问题提供了更好的支持。
基础设施与代码质量提升
-
测试与CI增强:
- 添加了对 Python 3.12 的支持
- 移除了对 Python 3.9 的官方支持
- 引入了 Ruff 静态分析工具
- 添加了预提交钩子(pre-commit)支持
- 增加了许可证检查
-
文档改进:
- 修复了多处文档错误和拼写问题
- 重组了预热调度器的文档结构
- 为多个函数添加了更详细的数学描述和示例
- 改进了 power_iteration 等函数的文档字符串
-
类型稳定性:通过 jax.lax.cond 测试确保了优化器状态的类型稳定性,提高了代码的可靠性。
向后兼容性说明
本次更新包含了一些破坏性变更,需要用户注意:
- 移除了长期弃用的函数
- 将 multi_transform 标记为已弃用,推荐使用 partition 替代
- 更改了随机数生成器的接口,使用 jax.random.key 替代 PRNGKey
- 重命名了部分树操作函数以保持命名一致性
Optax v0.2.5 版本的发布标志着这个优化器库在功能完备性和代码质量上又迈出了重要一步。新加入的优化算法和增强的工具函数将帮助研究人员和开发者更高效地解决各类优化问题,而基础设施的改进则确保了项目的长期可维护性。无论是训练最新的深度学习模型,还是解决传统的数值优化问题,Optax 都提供了强大而灵活的工具支持。
PaddleOCR-VLPaddleOCR-VL 是一款顶尖且资源高效的文档解析专用模型。其核心组件为 PaddleOCR-VL-0.9B,这是一款精简却功能强大的视觉语言模型(VLM)。该模型融合了 NaViT 风格的动态分辨率视觉编码器与 ERNIE-4.5-0.3B 语言模型,可实现精准的元素识别。Python00- DDeepSeek-OCR暂无简介Python00
openPangu-Ultra-MoE-718B-V1.1昇腾原生的开源盘古 Ultra-MoE-718B-V1.1 语言模型Python00
HunyuanWorld-Mirror混元3D世界重建模型,支持多模态先验注入和多任务统一输出Python00
AI内容魔方AI内容专区,汇集全球AI开源项目,集结模块、可组合的内容,致力于分享、交流。03
Spark-Scilit-X1-13BFLYTEK Spark Scilit-X1-13B is based on the latest generation of iFLYTEK Foundation Model, and has been trained on multiple core tasks derived from scientific literature. As a large language model tailored for academic research scenarios, it has shown excellent performance in Paper Assisted Reading, Academic Translation, English Polishing, and Review Generation, aiming to provide efficient and accurate intelligent assistance for researchers, faculty members, and students.Python00
GOT-OCR-2.0-hf阶跃星辰StepFun推出的GOT-OCR-2.0-hf是一款强大的多语言OCR开源模型,支持从普通文档到复杂场景的文字识别。它能精准处理表格、图表、数学公式、几何图形甚至乐谱等特殊内容,输出结果可通过第三方工具渲染成多种格式。模型支持1024×1024高分辨率输入,具备多页批量处理、动态分块识别和交互式区域选择等创新功能,用户可通过坐标或颜色指定识别区域。基于Apache 2.0协议开源,提供Hugging Face演示和完整代码,适用于学术研究到工业应用的广泛场景,为OCR领域带来突破性解决方案。00- HHowToCook程序员在家做饭方法指南。Programmer's guide about how to cook at home (Chinese only).Dockerfile013
Spark-Chemistry-X1-13B科大讯飞星火化学-X1-13B (iFLYTEK Spark Chemistry-X1-13B) 是一款专为化学领域优化的大语言模型。它由星火-X1 (Spark-X1) 基础模型微调而来,在化学知识问答、分子性质预测、化学名称转换和科学推理方面展现出强大的能力,同时保持了强大的通用语言理解与生成能力。Python00- PpathwayPathway is an open framework for high-throughput and low-latency real-time data processing.Python00