用 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”世界模型”理念的一次靠拢。

项目特点

1. 全 R 语言实现

完全基于 R 语言的 torch 和 luz 生态，不依赖任何 Python 深度学习框架。这在 LLM 领域极为少见，展示了 R 语言在深度学习前沿探索中的可能性。

2. JEPA + VQ 核心架构

不同于标准 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 码本，无需额外参数
• 推理时可结合连续空间操作（如插值、编辑）后再量化

3. JEPA vs 标准 Decode-Only 架构对比

标准 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 按时租用）

核心技术细节

分词器 — 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> 时截断

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），消除长尾噪音

训练流水线

数据准备（src/ 目录）与模型训练（scripts/ 目录）分离，各自独立运行。

数据准备

本项目仅需两个数据源：

• wiki_focused_semantic.jsonl：117 万篇中文文档，用于预训练。基于 Minimind 原始语料，先清洗 + 语义过滤 (Sentence-BERT)，得到与 AI/大数据/数据科学主题相关的精华文本。
• sft_data.jsonl：3000 条带有思考链指令的微调数据。通过调用 DeepSeek V4 Flash API（开启 thinking 模式）蒸馏获得，主题同样聚焦 AI/大数据/数据科学。

注：使用特定领域的数据在保证训练效果的同时降低了数据量，因此模型问答能力也限定在 AI/大数据/数据科学相关主题。

模型训练

项目分为 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）

代码设计特点

尽管项目可在 4090 上全速训练，但代码设计优先考虑了可读性和教学性：

• 每个脚本约 100-200 行，主流程总计约 700 行代码
• 核心架构（JEPA_model.R）在 utils 中统一维护，消除重复定义
• 共享配置（config.R）集中管理环境检测与超参，各脚本各司其职
• 不依赖混合精度、断点续训、多卡并行等复杂工程技巧——聚焦算法本质
• 完整的 R6 类 + nn_module 封装，结构清晰
• 带有详细的中文注释

快速开始

环境要求

• NVIDIA RTX 4090 24GB（或任何 CUDA 显卡，调整 batch size 即可）
• CUDA 12.x
• R >= 4.5

安装依赖

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 方向上的交流与启发