01JAX与NLP初探

• JAX简介
• 为什么选择JAX做NLP
• 环境配置与Hello World

02JAX核心基石

• NumPy兼容API
• 设备数组
• 即时编译(JIT)入门

03自动微分实战

• grad函数
• 梯度计算
• 与PyTorch自动微分的对比

04向量化与并行

• vmap、pmap、jit三剑客
• 批量处理文本数据

05随机数生成

• PRNGKey机制
• 与NumPy随机数的区别
• 在NLP数据增强中的应用

06构建线性模型

• 从零实现Logistic回归
• 用于文本分类

07损失函数与优化器

• 自定义交叉熵损失
• SGD优化器实现

08Flax框架入门

• Flax的Module
• setup与__call__方法

09Flax实战

• 用Flax搭建文本分类器

10Optax优化器库

• 链式优化器
• 学习率调度
• 权重衰减

11数据加载与预处理

• 用tf.data加载文本
• Tokenization实战

12词嵌入层

• 随机初始化
• 预训练词向量加载

13RNN与LSTM

• 用JAX/Flax实现RNN语言模型

14注意力机制

• 从零实现Scaled Dot-Product Attention

15Transformer编码器

• 多头注意力
• 前馈网络
• 层归一化

16Transformer解码器

• Masked Self-Attention
• Cross-Attention实现

17完整Transformer

• 搭建机器翻译Transformer

18训练循环

• 自定义训练循环
• 梯度累积
• 混合精度训练

19模型评估与推理

• BLEU分数计算
• 贪婪搜索与集束搜索

20预训练模型加载

• HuggingFace Transformers
• 与JAX集成

21BERT微调

• 用Flax实现BERT文本分类微调

22GPT风格模型

• 简化自回归语言模型

23Prompt工程与ICL

• 在JAX中实现少样本学习

24模型并行与分布式

• pmap实战
• 多GPU训练NLP模型

25模型量化与部署

• 用JAX量化NLP模型
• 导出为SavedModel

26JAX与TF Serving

• 部署NLP模型到生产环境

27性能调优

• NVTX标注
• Profiling
• 内存优化技巧

28大规模NLP实践

• 在TPU上训练Transformer
• 数据并行策略

29前沿探索

• JAX在LLM微调中的应用
• LoRA、Prefix Tuning

30项目实战

• 端到端情感分析系统
• 从数据处理到模型部署