exploreGPT - 纯Python实现的 Hybrid 架构 (TTT + Transformer)

纯Python实现的 Hybrid 架构，结合 Test Time Training (TTT) 和 Attention 两种记忆机制。

快速开始

python main.py

首次运行会自动从 GitHub 下载 input.txt 数据集（32033个人名）。

运行示例

num docs: 32033
vocab size: 27
num params: 1224
=== 开始 Hybrid 预训练 ===
Step  100 / 500 | Loss: 3.2978
Step  200 / 500 | Loss: 2.4301
Step  300 / 500 | Loss: 2.4887
Step  400 / 500 | Loss: 2.2724
Step  500 / 500 | Loss: 2.3534

=== 推理生成 (temperature=0.5) ===
sample  1: manli
sample  2: dia
sample  3: apellellle
...
sample 20: bakili

架构特点

Hybrid 记忆机制

本项目核心创新：TTT（长程记忆） + Attention（短程查表） 双层架构。

┌─────────────────────────────────────────────────────────────────┐
│                    Hybrid 架构                                   │
├─────────────────────────────────────────────────────────────────┤
│                                                                  │
│   Token Embedding (wte)                                         │
│         │                                                        │
│         ▼                                                        │
│   ┌─────────────────────────────────────┐                      │
│   │ Layer 0: TTT (长程记忆矩阵)          │                      │
│   │   - 状态矩阵 W 随输入动态更新         │                      │
│   │   - 自监督：重建输入 x               │                      │
│   │   - 适合：长期依赖、序列模式          │                      │
│   └─────────────────────────────────────┘                      │
│         │                                                        │
│         ▼                                                        │
│   ┌─────────────────────────────────────┐                      │
│   │ Layer 1: Attention (短程查表)        │                      │
│   │   - KV Cache 缓存历史信息            │                      │
│   │   - Q·K^T 计算注意力分数            │                      │
│   │   - 适合：局部上下文、精确检索        │                      │
│   └─────────────────────────────────────┘                      │
│         │                                                        │
│         ▼                                                        │
│   ┌─────────────────────────────────────┐                      │
│   │ Layer 2: MLP (特征变换)              │                      │
│   │   - 两层线性 + ReLU                  │                      │
│   │   - 非线性特征提取                   │                      │
│   └─────────────────────────────────────┘                      │
│         │                                                        │
│         ▼                                                        │
│   lm_head → logits → softmax → 预测                             │
│                                                                  │
└─────────────────────────────────────────────────────────────────┘

TTT vs Attention 对比

特性	TTT 层	Attention 层
记忆机制	状态矩阵 W	KV Cache
更新方式	在线学习（梯度下降）	追加缓存
适合场景	长程依赖、序列模式	局部上下文、精确检索
计算复杂度	O(d²) per token	O(n·d) per token

整体执行流程

┌─────────────────────────────────────────────────────────────────┐
│                        exploreGPT 完整执行流程                    │
└─────────────────────────────────────────────────────────────────┘

┌──────────────────┐
│  1. 数据准备阶段   │
│                  │
│  下载names.txt   │──────→ ┌─────────────────┐
│ 提取字符表uchars  │        │ 创建词汇表映射   │
│ 定义BOS标记      │         │ 字符↔Token ID   │
│ 打乱数据顺序     │         └─────────────────┘
└──────────────────┘                 ↓
                              ┌─────────────────┐
                              │  Value类        │
                              │  (自动微分引擎)  │
                              │ - 前向: 计算图  │
                              │ - 反向: 链式法则│
                              └─────────────────┘
                                       ↓
┌─────────────────────────────────────────────────────────────────┐
│  2. 训练循环 (500步, Adam优化器)                                  │
├─────────────────────────────────────────────────────────────────┤
│                                                                  │
│  ┌────────────┐     ┌────────────┐     ┌────────────┐          │
│  │ 取一个名字   │────→ │  前向传播    │────→│  计算损失  │          │
│  │ "anna"     │     │  forward() │     │  -log(p)   │          │
│  └────────────┘     └────────────┘     └────────────┘          │
│         ↑                 ↓                     ↓              │
│         │           ┌────────────┐     ┌────────────┐          │
│         │           │ TTT层      │     │  反向传播  │          │
│         │           │ Attention层│     │  backward()│          │
│         │           │ MLP层      │     │  计算梯度  │          │
│         │           └────────────┘     └────────────┘          │
│         │                  ↓          ┌────────────┐          │
│         │           ┌────────────┐    │ Adam优化器 │          │
│         │           │ 残差连接   │───→│ 更新参数   │          │
│         │           └────────────┘    │ 学习率衰减 │          │
│         │                             └────────────┘          │
└─────────┼────────────────────────────────────────────────────┘
          ↓
┌─────────────────────────────────────────────────────────────────┐
│  3. 推理生成 (20个样本, 温度采样)                                 │
├─────────────────────────────────────────────────────────────────┤
│                                                                  │
│  ┌─────────┐   ┌─────────┐   ┌─────────┐   ┌─────────┐          │
│  │ BOS开始 │──→│ forward │──→│ softmax │──→│ 采样    │          │
│  └─────────┘   │ 前向传播│   │ 温度缩放│   │ 下一个  │          │
│                └─────────┘   └─────────┘   │ token   │          │
│                                            └─────────┘          │
│                                  ↓                              │
│                            ┌─────────┐                          │
│                            │是否BOS? │──否──→ 继续循环           │
│                            └─────────┘                          │
│                                  ↓ 是                           │
│                            ┌─────────┐                         │
│                            │ 输出名字 │                         │
│                            └─────────┘                         │
└─────────────────────────────────────────────────────────────────┘

核心数据流

输入: "anna"
  ↓ Tokenizer
[BOS, 0, 13, 13, 0, BOS]
  ↓ Token Embedding (wte)
  [向量1, 向量2, ...]
  ↓
  ┌──────────────────────────────────────┐
  │ Layer 0: TTT                         │
  │  1. RMSNorm                          │
  │  2. linear(x, W) → y                 │
  │  3. 自监督更新 W (重建 x)            │
  │  4. 残差连接: y + x                  │
  └──────────────────────────────────────┘
  ↓
  ┌──────────────────────────────────────┐
  │ Layer 1: Attention                   │
  │  1. RMSNorm                          │
  │  2. Q, K, V 投影                     │
  │  3. KV Cache 追加                    │
  │  4. Q·K^T / √d → 注意力分数          │
  │  5. softmax → 加权求和               │
  │  6. 残差连接                         │
  └──────────────────────────────────────┘
  ↓
  ┌──────────────────────────────────────┐
  │ Layer 2: MLP                         │
  │  1. RMSNorm                          │
  │  2. linear → ReLU → linear           │
  │  3. 残差连接                         │
  └──────────────────────────────────────┘
  ↓
  lm_head → logits
  ↓ softmax (温度缩放)
probs: [0.15, 0.10, 0.35, ...]
  ↓
预测下一个 token

关键技术点

1. 自动微分 (Value类): 无需手动求导，自动构建计算图并反向传播梯度
2. TTT 层: 状态矩阵在推理时动态更新，实现长程记忆
3. Attention 层: KV Cache 缓存历史，实现短程精确检索
4. 残差连接: 梯度顺畅流动，支持更深网络
5. 温度采样: 控制生成文本的随机性和创造性
6. Adam 优化器: 动量 + 自适应学习率，训练更稳定

---

核心模块详解

1. 神经元：Value 类 (自动求导引擎)

在深度学习中，模型需要知道"我离正确答案差多少"以及"我该往哪边改"。

• data: 存储当前的数值
• grad: 存储梯度（改进方向），通过 backward() 计算
• 数学运算: 重载 +, *, **, log, exp 等，构建计算图
• backward(): 从 Loss 开始，利用链式法则反向传播梯度

2. 分词器：Tokenizer

模型不认识字符，只认识数字。

uchars = sorted(set(''.join(docs)))  # ['a', 'b', 'c', ...]
BOS = len(uchars)                     # 26 (特殊标记)
vocab_size = 27                       # 26字母 + 1个BOS

把 "anna" 变成 [BOS, 0, 13, 13, 0, BOS]

3. TTT 层：Test Time Training

核心创新：状态矩阵 W 在推理时继续学习

def ttt_step(x, W_hidden, ttt_lr=0.2):
    y = linear(x, W_hidden)           # 输出
    x_hat = linear(x, W_hidden)       # 重建输入
    loss = sum((xh - xi)**2)          # 自监督损失
    loss.backward()                    # 计算梯度
    W_hidden -= ttt_lr * grad          # 更新 W
    return y

为什么有效：

• 每个 token 都会让 W 学到一点东西
• 长序列中，W 累积了整个序列的"记忆"
• 类似 RNN，但更新规则是梯度下降

4. Attention 层：KV Cache

传统注意力，缓存历史 K 和 V：

def attention_step(q, k_cache, v_cache):
    scores = [q·k / √d for k in k_cache]  # 注意力分数
    weights = softmax(scores)               # 归一化
    return Σ(weights[t] * v_cache[t])      # 加权求和

5. 训练循环

for step in range(500):
    # 1. 前向传播
    logits = forward(token_id, ttt_states, kv_cache)
    
    # 2. 计算损失
    loss = -probs[target].log()
    
    # 3. 反向传播
    loss.backward()
    
    # 4. Adam 更新（带学习率衰减）
    lr_t = lr * (1 - step / 500)
    p.data -= lr_t * m_hat / (v_hat**0.5 + eps)

6. 推理生成

for _ in range(block_size):
    logits = forward(token_id, ttt_states, kv_cache)
    probs = softmax([l / temperature for l in logits])
    token_id = random.choices(range(vocab_size), weights=probs)[0]
    if token_id == BOS: break
    sample.append(it_to_char[token_id])

Temperature 参数：

• 0.1: 保守，生成常见名字
• 0.5: 平衡，适度创造性
• 1.0: 狂野，生成奇怪组合

---

模型参数

n_embd = 12          # 嵌入维度
block_size = 16      # 最大上下文长度
vocab_size = 27      # 26字母 + BOS
num_params = 1224    # 总参数量

数据来源

• 数据集: karpathy/makemore 的 names.txt
• 样本数: 32033 个人名
• 自动下载: 首次运行自动获取

---

文件结构

exploreGPT/
├── main.py           # 核心实现 (~200行)
├── input.txt         # 人名训练数据 (自动下载)
├── README.md         # 本文档
├── production.md     # 项目全局文档
└── schema/           # 任务文档目录
    └── archive/      # 已完成任务归档

扩展建议

1. 增加层数: 添加更多 TTT/Attention 层
2. 增大嵌入维度: n_embd = 64 或更大
3. 更长训练: num_steps = 2000+
4. 学习率调度: 余弦衰减、预热等
5. 多层 TTT: 多个 TTT 层串联，不同时间尺度