Nested Learning 架构设计文档

Nested Learning 架构设计文档

本文档通过图文并茂的方式,详细解析 Nested Learning (HOPE) 项目的整体架构、核心组件和数据流,帮助开发者快速理解项目并进行二次开发。

1. 系统整体架构

1.1 架构概览

层级学习系统

记忆系统

输出层

HOPE Blocks (N layers)

嵌入层

输入层

输入 Tokens
[batch, seq_len]

Embedding Layer
vocab_size → dim

HOPE Block 0

HOPE Block 1

HOPE Block N-1

LayerNorm

LM Head
dim → vocab_size

输出 Logits
[batch, seq_len, vocab]

TITAN Memory
Long-term Memory

CMS (Continual Memory System)
Multi-level Memory

LevelClock
更新调度

Teach Signal
教学信号

Surprise Detection
惊喜检测

1.2 核心设计理念

Nested Learning (HOPE) 的核心创新在于层级化的记忆系统多频率更新机制

  1. TITAN Memory: 长期记忆模块,负责存储和检索知识
  2. CMS (Continual Memory System): 持续记忆系统,包含多个更新频率不同的记忆层
  3. LevelClock: 层级时钟,控制不同层的更新时机
  4. Teach Signal: 教学信号,引导模型进行有意义的更新
  5. Surprise Detection: 惊喜检测,过滤无意义的更新

2. HOPE Block 变体架构

2.1 四种 Block 变体

变体选择

Attention only

With TITAN

With Self-Modifying

Baseline

输入 x

Block Variant

hope_attention
Attention → CMS

hope_hybrid
TITAN → CMS

hope_selfmod
Self-Modifying Titans → CMS

transformer
Attention → MLP

输出

输出

输出

输出

2.2 HOPEBlock 详细结构 (hope_hybrid)

输出

CMS 路径

注意力路径

TITAN 路径

输入

Update Signal

Update Signal

x: [batch, seq, dim]

teach_signal (可选)

TITAN Memory

Residual Connection

Update Controller

Surprise Check

SelfAttention

Residual Connection

CMS Level 1 (fast)

CMS Level 2 (slow)

CMS Level N

Multi-level Fusion

输出: [batch, seq, dim]

2.3 HOPEAttentionBlock 结构 (hope_attention)

更新系统

核心组件

输入

Controlled Update

x: [batch, seq, dim]

teach_signal (可选)

SelfAttention

CMS (无 TITAN)

LayerNorm

LevelClock

Optimizer Manager

输出

2.4 HOPESelfModBlock 结构 (hope_selfmod)

CMS 系统

计算单元

Self-Modifying Titans

输入

Fast State Update

Fast State Update

Fast State Update

Fast State Update

Fast State Update

Fast State Update

x: [batch, seq, dim]

teach_signal (可选)

M_k (Key Memory)

M_v (Value Memory)

M_q (Query Memory)

M_eta (Eta Memory)

M_alpha (Alpha Memory)

M_memory (Memory)

Feed Forward

Update Rule
(Eq. 90/93 DGD)

CMS Blocks

输出

3. 数据流详解

3.1 训练时前向传播

LevelClockCMSTITAN MemoryHOPEBlockHOPEModel用户LevelClockCMSTITAN MemoryHOPEBlockHOPEModel用户遍历所有层级loop[Each Level]收集更新统计信息alt[有 teach_signal]forward(tokens, teach_signal)_run_blocks(tokens, teach_signal)SelfAttention(x)forward(query)forward(attn_out, return_intermediates=True)CMSBlock.forward(x)tick()should_update(level)?cms_out, cms_inputs, cms_outputs_update_cms(cms_inputs, cms_outputs, teach_signal)update(key, value, error_signal)block_outputLayerNorm + LM Headlogits

3.2 Fast State 模式 (推理/测试时适应)

更新流程

init_fast_state()

memorize_tokens(inputs, fast_state)

输出更新后的 fast_state

原始模型保持不变

Fast State 模式

克隆参数到上下文

在克隆上计算更新

不修改原始权重

低开销,适合推理时适应

标准训练模式

参数更新
直接修改模型权重

持久化变更

高开销

3.3 LevelClock 更新调度

01234567891011更新 step 0 更新 step 1 更新 step 2 更新 step 2 更新 step 3 更新 step 3 更新 step 4 更新 step 4 更新 step 5 更新 step 5 更新 step 6 更新 step 6 更新 step 6 更新 step 7 更新 step 8 更新 step 8 更新 step 9 更新 step 9 更新 step 10 更新 step 10 更新 step 10 Level: fast (period=1)Level: medium (period=3)Level: slow (period=5)Level: titan (period=2)LevelClock 更新调度示例

4. 核心模块依赖关系

4.1 模块依赖图

学习系统

记忆系统

Transformer 基础

HOPE 模块

核心模块

model.py
HOPEModel

model.py
ModelConfig

hope/block.py
HOPEBlock

hope/block.py
HOPEAttentionBlock

hope/block.py
HOPESelfModBlock

hope/self_mod.py
HOPESelfMod

transformer.py
TransformerBlock

backbones.py
SelfAttention

titan/memory.py
TitanMemory

titan/self_modifying.py
SelfModifyingTitans

cms.py
CMS, CMSBlock

levels.py
LevelSpec, LevelClock

optim/manager.py
LevelOptimizerManager

fast_state.py
ModelFastState

4.2 关键类继承关系

«abstract»

nn.Module

+forward() : Tensor

ModelConfig

+vocab_size: int

+dim: int

+num_layers: int

+block_variant: str

+titan_level: LevelSpec

+cms_levels: LevelSpec[]

LevelSpec

+name: str

+update_period: int

+warmup_steps: int

+jitter: int

HOPEModel

+forward() : Tensor

+forward_with_pre_norm() : Tuple

+init_fast_state() : ModelFastState

HOPEBlock

+titan_memory: TitanMemory

+cms: CMS

+forward() : Tensor

HOPEAttentionBlock

+attn: SelfAttention

+cms: CMS

+forward() : Tensor

HOPESelfModBlock

+selfmod: SelfModifyingTitans

+cms: CMS

+forward() : Tensor

TransformerBlock

+attn: SelfAttention

+ff: FeedForward

+forward() : Tensor

TitanMemory

+forward() : Tensor

+surprise() : Tensor

+update() : None

SelfModifyingTitans

+forward() : Tensor

+update() : None

CMS

+blocks: Dict<str, CMSBlock>

+forward() : Tensor

CMSBlock

+forward() : Tensor

HOPEBlockConfig

HOPEAttentionBlockConfig

HOPESelfModBlockConfig

5. CMS (Continual Memory System) 详解

5.1 CMS 多层级结构

输出

CMS 系统

输入 x

Level: slow (update_period=5)

Level: medium (update_period=3)

Level: fast (update_period=1)

Δ1 (增量)

Δ2 (增量)

grad_clip

grad_clip

grad_clip

[batch, seq, dim]

CMSBlock 1

快速更新
每个 step 更新

CMSBlock 2

中速更新
每 3 steps 更新

CMSBlock 3

慢速更新
每 5 steps 更新

[batch, seq, dim]

5.2 CMSBlock 内部结构

输出

CMSBlock

输入 x

增量 Δ

[batch, seq, dim]

LayerNorm

Linear
dim → dim*multiplier

Activation
GELU/ReLU/SiLU

Linear
dim*multiplier → dim

Grad Clip
norm/clip

x + Δ
[batch, seq, dim]

6. LevelClock 调度机制

6.1 更新决策逻辑

Yes

No

Yes

No

True

False

should_update(level_name)

获取 LevelSpec

获取 LevelState

step < warmup_steps?

返回 False
预热期不更新

计算 delta = step - last_step

获取 update_period

jitter > 0?

period = period + step % (jitter+1)

last_step < 0
OR
delta >= period

返回 True
需要更新

返回 False
不需要更新

6.2 层级配置示例

# 配置示例
titan_level = LevelSpec(
    name="titan",
    update_period=2,      # 每 2 步更新一次
    warmup_steps=0,       # 无预热
    jitter=0,             # 无抖动
)

cms_levels = [
    LevelSpec(
        name="fast",
        update_period=1,  # 每步更新
        warmup_steps=0,
    ),
    LevelSpec(
        name="medium",
        update_period=3,  # 每 3 步更新
        warmup_steps=10,  # 前 10 步不更新
    ),
    LevelSpec(
        name="slow",
        update_period=10, # 每 10 步更新
        warmup_steps=50,  # 前 50 步不更新
    ),
]

7. 快速状态 (Fast State) 机制

7.1 标准模式 vs Fast State 模式

使用场景

Fast State 模式 (推理/测试时适应)

克隆参数到 GPU 内存

requires_grad=False

低内存开销

临时上下文,不持久化

标准模式 (训练)

参数直接更新

requires_grad=True

高内存开销

持久化变更

训练: 标准模式

推理时适应: Fast State

Few-shot 学习: Fast State

7.2 Fast State 数据结构

ModelFastState

+blocks: List<BlockFastState>

BlockFastState

+titan_params: ParamDict | None

+cms_params: Dict<str, ParamDict>

+level_manager: LevelOptimizerManager

+selfmod_state: SelfModifyingTitansState | None

ParamDict

+Dict<str, Tensor>

SelfModifyingTitansState

+memories: Dict<str, ResidualMLPMemoryState>

+momentum: Dict<str, Tensor>

ResidualMLPMemoryState

+w1: Tensor

+w2: Tensor

+w_skip: Tensor | None

+m_w1: Tensor | None

+m_w2: Tensor | None

+m_w_skip: Tensor | None

8. 配置文件结构

8.1 Hydra 配置层次

CLI 覆盖

配置目录

数据配置

HOPE 变体

基础配置

model.yaml

train.yaml

hope/pilot.yaml

hope/mid.yaml

hope/target.yaml

data/pilot.yaml

data/mid.yaml

--config-name=pilot
--train.device=cuda:0

8.2 ModelConfig 关键参数

参数类型默认值说明
vocab_sizeint必填词汇表大小
dimint必填模型维度
num_layersint必填Transformer 层数
headsint必填注意力头数
block_variantstr“hope_hybrid”块变体类型
titan_levelLevelSpec必填TITAN 层级配置
cms_levelsSequence[LevelSpec]必填CMS 层级列表
teach_scalefloat1.0教学信号缩放
teach_clipfloat0.0教学信号裁剪阈值
gradient_checkpointingboolFalse梯度检查点
surprise_thresholdfloat | NoneNone惊喜阈值
self_mod_lrfloat1e-3自修改学习率
self_mod_chunk_sizeint1自修改分块大小

9. 二次开发指南

9.1 添加新的 Block 变体

添加新变体

创建配置类

继承 nn.Module 实现前向传播

实现 set_surprise_threshold

实现 set_allowed_levels

在 HOPEModel.init 中添加分支

在 init_fast_state 中添加支持

添加配置和测试

9.2 添加新的 Memory 模块

添加新 Memory

创建配置类 Config

继承 nn.Module 实现 forward

实现 surprise 方法

实现 update 方法

在 HOPEBlock 中集成

在 LevelOptimizerManager 中注册

添加测试和文档

9.3 添加新的优化器

添加新优化器

在 optim/factory.py 中添加 build_optimizer 分支

实现优化器类

在 LevelConfig 中配置

在 LevelOptimizerManager 中使用

添加测试

10. 训练流程总结

训练循环

定期

保存检查点

评估

零样本评估

NIAH 测试

持续学习评估

层级更新

LevelClock.tick()

检查每个层级的 should_update

执行优化步骤

记录统计信息

初始化

加载配置 (Hydra)

创建 ModelConfig

初始化 HOPEModel

初始化优化器

获取 batch

前向传播

计算 loss

反向传播

层级更新

11. 关键文件索引

文件路径功能核心类
src/nested_learning/model.py主模型HOPEModel, ModelConfig
src/nested_learning/hope/block.pyHOPE 块实现HOPEBlock, HOPEAttentionBlock, HOPESelfModBlock
src/nested_learning/titan/memory.pyTITAN 记忆TitanMemory, TitanMemoryConfig
src/nested_learning/titan/self_modifying.py自修改 TitansSelfModifyingTitans, SelfModifyingTitansConfig
src/nested_learning/cms.py持续记忆系统CMS, CMSBlock
src/nested_learning/levels.py层级规范LevelSpec, LevelClock
src/nested_learning/optim/manager.py层级优化器LevelOptimizerManager, LevelConfig
src/nested_learning/fast_state.py快速状态ModelFastState, BlockFastState
src/nested_learning/backbones.py注意力机制SelfAttention, AttentionConfig
src/nested_learning/transformer.py基础 TransformerTransformerBlock, FeedForward

文档版本: 1.0
最后更新: 2024年
适用版本: nested-learning 0.1.0+

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值