中文 | English
用 R torch 从零构建的 JEPA (Joint Embedding Predictive Architecture) 大语言模型，含 BPE 分词器训练、预训练、SFT 微调、推理全流程，总计约 700 行主脚本代码，单卡 RTX 4090 可跑，训练成本约 4 块钱。
最初版本基于标准 GPT 架构。后来和好友 Vivian Zhang（NYC Data Science Academy 创始人）聊到 Yann LeCun 团队在 JEPA 方向的工作，很受启发，于是将整个模型用 JEPA 框架重新实现了一遍——用 InfoNCE 对比损失替代交叉熵，用VQ 最近邻量化替代线性投影 head。这既是一次对 JEPA 在语言建模中的可行性验证，也是向 LeCun 那套"世界模型"理念的一次靠拢。
完全基于 R 语言的 torch 和 luz 生态实现，不依赖任何 Python 深度学习框架。这在 LLM 领域极为少见，展示了 R 语言在深度学习前沿探索中的可能性。
不同于标准 GPT 式语言模型（在 logits 上计算交叉熵损失），本项目采用 JEPA (Joint Embedding Predictive Architecture)：
- 隐空间预测：模型不是直接预测下一个 token 的 ID，而是在 embedding 语义空间中预测下一个 token 的连续向量表示
- InfoNCE 对比损失：使用对比学习目标函数，让预测的隐向量在语义空间中接近真实下一 token 的 embedding，远离其他 token
- VQ (Vector Quantization) 解码：推理时通过最近邻搜索，在 token embedding 矩阵（同时作为码本）中找到最接近预测向量的 token ID
- 直通估计器 (STE)：VQ 量化过程通过 STE 保持梯度通路，可端到端反向传播
这种架构相比传统 LLM 有以下特点：
- 使用对比学习让模型学习到 token 之间的语义关系，而非机械的 n-gram 统计
- Token embedding 矩阵天然作为 VQ 码本，无需额外参数
- 在推理时可结合连续空间操作（如插值、编辑）后再量化
本项目最核心的技术决策是用 JEPA 替代了主流的 decode-only（GPT 式）架构。以下是两者的关键差异：
| 维度 | 标准 Decode-Only (GPT) | 本项目 JEPA + VQ |
|---|---|---|
| 损失函数 | 交叉熵损失 (Cross-Entropy) | InfoNCE 对比损失 |
| 监督信号 | 在词表维度上做分类，预测正确的 token ID | 在 embedding 空间中拉近正样本、推开负样本 |
| 输出头 | 线性 LM Head | VQ 最近邻搜索，复用 token embedding 矩阵 |
| 模型容量 | LM Head 贡献约 5.2M 参数 | 无输出头参数，全部容量分配给 Transformer 主干 |
| 梯度回传 | 标准 softmax 交叉熵，梯度直接回传 | VQ 量化通过直通估计器 (STE) 绕路回传 |
| 解码方式 | 取概率最高的 token | 取最近的 embedding |
| 学习目标 | 学习 token 的共现频率分布 (n-gram 模式) | 学习 token 间的语义相似度（语义空间中的相对位置） |
| 输出空间 | 离散词表上的概率分布（唯一 ground truth） | 连续 embedding 空间（可插值、编辑后再量化） |
直观理解：标准 GPT 像一个"词汇选择题"——模型在 8192 个词中选一个填进去，选错了就纠正。JEPA 则像一个"语义连线题"——模型在向量空间中预测下一个词应该在哪，看预测位置和真实位置的远近。
两种架构各有所长。标准 GPT 在开放域文本生成上更成熟稳定，而 JEPA 在语义理解和连续空间操作上有独特优势。本项目选择了 JEPA，是对 LeCun 世界模型理念的一次实践验证。
- 模型架构：约 50M 参数（10 层 Transformer / 10 个注意力头 / 640 隐藏层维度 / 8192 词表）
- 上下文长度：256（预训练）→ 512（SFT）
- 训练数据：
- 预训练：117 万篇 / 787MB（基于 Minimind 语料清洗过滤）
- SFT：3000 条 CoT 思维链指令（DeepSeek V4 Flash 蒸馏）
- 算力与成本：单卡 RTX 4090 (24GB)，总耗时 2 小时（PT 1.5h + SFT 0.5h），约 4 元（AutoDL 按时租用）
数据准备（src/ 目录）与模型训练（scripts/ 目录）分离，各自独立运行：
A.关于数据
本项目仅需要两个数据源：
wiki_focused_semantic.jsonl：117 万篇中文文档，用于预训练。基于 Minimind 原始语料，先清洗 + 语义过滤(Sentence-BERT)，得到与 AI/大数据/数据科学主题相关的精华文本。sft_data.jsonl：3000 条带有思考链指令的微调数据。数据通过调用 DeepSeek V4 Flash API（开启 thinking 模式）蒸馏获得，主要也是 AI/大数据/数据科学相关主题。
注：仅仅使用特定领域的数据，降低数据量，提高训练效率。因此，问答也仅限于 AI/大数据/数据科学相关主题。
B.模型训练流水线
项目分为 5 个独立脚本，共享同一份环境检测与超参配置 (scripts/utils/config.R)，每个阶段可独立运行：
01_bpe_tokenizers.R # BPE 分词器训练 + 评估
02_JEPA_VQ.R # JEPA 预训练（两阶段）
03_VQ_inference.R # 预训练模型推理测试
04_SFT_train.R # 指令微调（冻结 1-7 层 + Loss Mask）
05_SFT_inference.R # SFT 后推理（带温度/惩罚/Top-K）分词器 - 01_bpe_tokenizers.R
- 词表大小 8192 (2^13)，对 Tensor Core 友好：8192 ÷ 128 = 64，无底层算力损耗
- UNK 率仅 0.16%（每 1000 token 约 1.6 个未知词）
- 压缩率 53.34（每 100 中文字符消耗 ~53 个 token）
- 流式读取 JSONL 采样，不加载全部数据到内存
预训练 - 02_JEPA_VQ.R
- 两阶段训练策略：第一阶段用 one-cycle LR 调度 + 分组学习率（embedding 层 1e-4，其他层 1e-3），第二阶段用 warm restart + 余弦退火（lr 5e-4）
- 使用
torch_scaled_dot_product_attention调用底层 FlashAttention 核函数，自动处理 causal mask - 支持混合精度训练 (AMP)，自动检测 CUDA 可用性
- 自动环境识别：Mac 本地测试模式使用小 batch、单进程、关闭 AMP；Linux GPU 模式全速运行
SFT 指令微调 - 04_SFT_train.R
- Loss Mask：自动对 User 侧的 prompt 做 pad 掩码，仅对 Assistant 侧的回复和思考过程计算损失
- 层冻结策略：冻结 embedding 层和前 7 层 Transformer，仅微调后 3 层 + predictor，有效防止灾难性遗忘
- DeepSeek-R1 式思考链：支持
<think>/</think>标签，赋予模型长链推理能力 - 位置编码动态扩展：预训练位置编码可自动缝合扩展到更长序列
- 原生 torch 训练循环（不依赖 luz），提供更高的控制力
推理（预训练 — 03_VQ_inference.R）
- VQ 贪心解码：模型输出的连续预测向量通过最近邻搜索，在 token embedding 码本中找到最近的 token ID，作为下一个 token
- EOS 提前停止：检测到 EOS token id 或解码文本为
<EOS>时截断
VQ_inference.mp4
推理（SFT — 05_SFT_inference.R）
- Logits 重建：用
pred @ emb.T将 predictor 输出的连续向量投影到词表空间，重建伪 logits - 温度控制：
logits = logits / temperature，默认 0.4，降低温度减少随机性 - 重复惩罚：对已出现的 token 做自适应惩罚（出现过的 token 分数除以惩罚系数 1.05）
- Top-K 截断：仅保留概率最高的 K 个 token（默认 10），消除长尾噪音
SFT_inference.mp4
尽管项目可在 4090 上全速训练，但代码设计时优先考虑了可读性和教学性：
- 每个脚本约 100-200 行，主流程总计约 700 行代码
- 核心架构（JEPA_model.R）在 utils 中统一维护，消除重复定义
- 共享配置（config.R）集中管理环境检测与超参，各脚本各司其职
- 不依赖混合精度、断点续训、多卡并行等复杂工程技巧——聚焦算法本质
- 完整的 R6 类 +
nn_module封装，结构清晰 - 带有详细的中文注释
- NVIDIA RTX 4090 24GB（或任何 CUDA 显卡，调整 batch size 即可）
- CUDA 12.8
- R >= 4.5
- R torch 0.7.0
install.packages(c("jsonlite", "tokenizers.bpe", "R6", "purrr", "torch", "luz"))# 1. BPE 分词器训练
source("scripts/01_bpe_tokenizers.R")
# 2. JEPA 预训练
source("scripts/02_JEPA_VQ.R")
# 3. 预训练模型推理测试
source("scripts/03_VQ_inference.R")
# 4. SFT 指令微调
source("scripts/04_SFT_train.R")
# 5. SFT 后推理
source("scripts/05_SFT_inference.R")本项目数据可以在 HuggingFace 下载，放入 data/raw/ 目录：
预训练数据：每行 JSON 包含 text 字段
{"text": "深度学习是机器学习的一个子集..."}SFT 数据：每行 JSON 包含 instruction 和 output 字段，output 中包含 <think> 标签，表示模型的思考过程。
{"instruction":"Git 用于对代码进行什么控制？","output":"<think>\n用户询问Git的功能，这是关于代码版本控制的工具"}scripts/
├── 01_bpe_tokenizers.R BPE 分词器训练与评估
├── 02_JEPA_VQ.R JEPA 预训练（可选两阶段）
├── 03_VQ_inference.R 预训练模型推理
├── 04_SFT_train.R 指令微调（Loss Mask + 层冻结）
├── 05_SFT_inference.R SFT 后推理（带温度/惩罚/Top-K）
└── utils/
├── config.R 共享环境检测与超参配置
├── JEPA_model.R 核心模型定义（复用于 02-05）
├── RtomicBPETokenizer.R BPE 分词器 R6 封装
└── count_parameters.R 参数量统计
data/raw/ JSONL 原始数据（含 787MB 预训练语料）
models/
└── rtomic_bpe.model 训练好的 BPE 分词器
checkpoints/ 训练的模型权重
- 增加 MoE 架构，提高模型的并行计算能力
- 增加 decode-only 模型的对比实验，评估其在推理时的性能
- 预训练的原始语料来自 Minimind 项目
- 语义过滤基于 Sentence-BERT 模型，可以快速过滤相似文本
- 参考了 Yann LeCun 等人提出的 JEPA (Joint Embedding Predictive Architecture) 思想
- 特别感谢 Vivian Zhang 在 JEPA 方向上的交流与启发