Transformers 中的 Doge 小语言模型完全指南：动态掩码注意力（DMA）与跨域 MoE（CDMoE）架构解析与实战

Doge 是由 SmallDoge 团队训练、于 2025-07-08 贡献给 Hugging Face Transformers 的一系列小语言模型（SLM），其设计目标是在状态空间模型（SSM）与自注意力机制之间取长补短，通过零阶保持（zero-order hold）方式从缓存的价值状态（value states）中动态计算掩码，解决主流语言模型在长文本中“迷失上下文”的问题。本文以 官方模型文档 为骨架，结合本仓库的 配置实现、建模实现 与 测试套件，系统讲解 Doge 的架构原理、配置参数与开箱即用的推理/微调实践，读完你可以在 Transformers 中直接加载并运行 Doge 模型，也能理解它每一层模块的底层工作方式。

Doge 是什么：面向“上下文迷失”问题的小模型架构

Doge 家族基于小 Doge（SmallDoge）架构，是一组专注于小规模参数下的语言模型。根据 模型文档 的 Overview 描述，Doge 的出发点非常明确：

• 融合两种序列建模范式：Doge 结合状态空间算法与自注意力算法的优点，期望在小模型上获得更稳定的长文本建模能力；
• 用动态掩码解决“迷失在上下文中”：Doge 使用 wsd_scheduler 学习率调度器在 smollm-corpus 语料上预训练，并且可以从稳定阶段的 checkpoint（stable stage checkpoints）继续训练，或在其基础上追加稀疏激活的前馈网络（sparse activation feedforward networks），从而在不重训整个模型的前提下持续迭代、降低成本。

从架构上拆解，Doge 的每个解码器层包含两条核心变换路径（对应 DogeDecoderLayer.forward）：

1. 序列变换（sequence transformation）：使用 Dynamic Mask Attention（DMA）。训练阶段使用与 value states 相关的自注意力，推理阶段则退化为“不带过去状态衰减的状态空间”形式，用来规避传统 Transformer 或 SSM 在长文本上的“上下文迷失”；
2. 状态变换（state transformation）：使用 Cross Domain Mixture of Experts（CDMoE），它由稠密线性层与稀疏嵌入层组成。由于稀疏参数可以被额外追加到稠密权重 checkpoint 之上，模型可以从已有的稠密权重继续训练，无需整体重训。

此外，Doge 在层归一化与残差连接上引入了可学习参数（RMSNorm + 带参数残差），用于适配深层模型的梯度量级。整体架构图可参考模型文档中的 doge_architecture.png（远程图片，本仓库未包含本地副本）。

从源码理解两大核心模块

modeling_doge.py 是从 modular_doge.py 自动生成的建模文件（文件头注释明确说明两者由 CI 保证同步）。从 modular 文件的 import 可以看出 Doge 站在了成熟模型肩膀上：RMSNorm、RoPE 旋转位置编码、MLP 与 eager attention 直接复用 Llama 的实现，而模型主体与 CausalLM 头部参照 Mixtral（MixtralModel、MixtralForCausalLM）。这也解释了为什么 Doge 原生就带 MoE 相关的输出类型与辅助损失。

动态掩码注意力（DMA）：从缓存 value states 生成稀疏掩码

DMA 的核心思想在 prepare_dynamic_mask 的注释中写得非常直白：动态计算注意力掩码，遮蔽应该被遮蔽的 token，从而形成稀疏注意力。结合 dt_states 与 attention_mask 生成最终的 attn_mask。

其前向链路（DogeAttention.forward）如下：

1. 投影与归一化：q_proj/k_proj/v_proj 完成线性投影后，Q、K 额外经过 q_norm/k_norm（按 head_dim 做的 RMSNorm），再施加旋转位置编码 apply_rotary_pos_emb；

2. KV 缓存：若传入 past_key_values，则调用 past_key_values.update(...) 更新缓存（DogeModel 中默认使用 DynamicCache）；

3. 由 value states 计算动态掩码：

   • dt_states = dt_proj(value_states) —— 用一个轻量线性层把价值状态投影到 num_key_value_heads 维；
   • dt_states = exp(A * softplus(dt_states)) —— 其中 A 是每个 KV 头一个标量的可学习参数（形状 [num_key_value_heads]，初始化全 0），softplus 保证非负、exp 得到衰减/强度因子；
   • prepare_dynamic_mask 将 dt_states（形状 [batch, heads, key_len]）广播为 [batch, heads, query_len, key_len]，把 padding 位置填为 min_dtype，随后执行核心稀疏化：当 key 序列长度超过 keep_window_size 时，仅保留按分值最大的 keep_window_size 个位置（用 torch.topk + scatter 实现），其余一律遮蔽。这正对应文档所说的“零阶保持 + 动态窗口稀疏化”。

   值得注意，dt_states 是对缓存过的 value states做投影得到的，因此训练阶段可以端到端学习“该掩掉哪些历史 token”，而推理阶段这种按 value 状态的稀疏化天然构成一种无显式衰减的状态空间路径。

4. Attention 后端选择：Doge 通过 AttentionInterface 注册了名为 doge_flex_attention 的 flex attention 实现（modeling_doge.py#L248-L249），并通过 config._attn_implementation 在 flex attention 与 eager attention 之间切换；测试所需的注意力打分使用 head_dim**-0.5 缩放，softmax 在 fp32 上完成以保持数值稳定。

从 DogePreTrainedModel 的类属性可以确认后端能力边界：_supports_flash_attn = False、_supports_sdpa = True、_supports_flex_attn = True，即官方明确支持 SDPA 与 Flex Attention，但不声称支持 Flash Attention。其 _init_weights 也体现了两个特殊初始化约定：注意力参数 A 用 zeros_、残差参数 input_residual/post_attention_residual 用 ones_ 初始化，保证初始化时行为与普通残差网络一致。

带参数的残差与归一化：稳定深层小模型的梯度

DogeDecoderLayer 中的两条变换路径都被 RMSNorm + 可学习残差包裹（modeling_doge.py#L447-L492）：

# 序列变换
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
hidden_states, _ = self.self_attn(...)
hidden_states = self.input_residual * residual + hidden_states   # 可学习残差

# 状态变换
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = self.post_attention_residual * residual + hidden_states  # 可学习残差

其中 input_residual、post_attention_residual 均为形状 [hidden_size] 的 nn.Parameter，初始化为全 1。DogeRMSNorm 即 T5LayerNorm 的实现（逐元素可学习缩放 weight，先转 fp32 计算方差再还原精度）。残差系数可学习意味着模型能自适应地缩放“跨层传递的信号”与“本层新学习信号”的比例，这正是应对深层小模型梯度波动的手段之一。

跨域 MoE（CDMoE）：在稠密 checkpoint 上“叠加”稀疏专家

DogeCDMoE（modeling_doge.py#L390-L444）把 FFN 拆成两部分：

• 共享稠密专家：gate_proj/up_proj/down_proj，与 DogeMLP 完全一致的 SiLU 门控结构，当 is_moe=False 时解码器直接使用普通 DogeMLP；
• 路由检索专家：router_gate 先把 hidden states 映射为 num_keys * 2 的 logits，随后分别沿 x、y 两个“域”做 top-k（num_keys = floor(sqrt(num_experts))），再通过笛卡尔积合成最终专家索引——这正是“跨域（Cross Domain）”的含义：把二维的专家平面分解为两个独立的一维选择，减少路由维度爆炸；
• 专家本身用嵌入表实现：down_embed（形状 [num_experts, hidden_size]）与 up_embed，对选中的 top-k 专家索引查表后与 hidden states 做两次矩阵乘完成“下投影 → 激活 → 上投影”，最终加到共享专家的输出上。

由于被路由的专家是稀疏的嵌入参数，训练新数据/新任务时可以只冻结共享部分、扩展专家嵌入规模继续训练，这正是文档所述的“从稠密 checkpoint 以低成本追加稀疏参数继续训练”的实现基础。当 config.is_moe=True 且需要加载稠密 checkpoint 继续训练时，MoE 会“继承 MLP 来初始化”（见 DogeConfig.is_moe 说明）。

Doge 的 MoE 还复用了 Switch Transformer 风格的负载均衡辅助损失 load_balancing_loss_func：每个解码器层把 [2, batch*seq, num_keys] 的路由 logits 汇总，统计各专家被选中的 token 比例与路由概率，惩罚路由不均衡。默认系数为 router_aux_loss_coef=0.001，通过 output_router_logits=True 打开，并在有 labels 时叠加进总损失。

在 Transformers 中快速上手

用 Doge-Base 做文本生成

Base 版本直接使用因果语言建模接口即可（对应 模型文档 的 Usage 示例）：

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("SmallDoge/Doge-20M")
model = AutoModelForCausalLM.from_pretrained("SmallDoge/Doge-20M", device_map="auto")
inputs = tokenizer("Hey how are you doing?", return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.batch_decode(outputs))

Doge 已接入 Auto 体系（auto_mappings.py 中注册 "doge" -> DogeConfig，modeling_auto.py 分别注册了 DogeModel、DogeForCausalLM、DogeForSequenceClassification），因此无需手工指定类名。所有 SmallDoge 官方发布的 Doge 模型 checkpoint 都可以通过 Hub 上的模型 ID 直接加载，系列从 20M 级别起步（如 SmallDoge/Doge-20M），文档中的配置默认值对应 SmallDoge/Doge-320M。

用 Doge-Instruct 走聊天模板对话

Instruct 版本需要应用 chat template，配合 GenerationConfig 与 TextStreamer 逐 token 流式输出（以下修正了原示例中变量名笔误，保证可运行）：

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, TextStreamer

tokenizer = AutoTokenizer.from_pretrained("SmallDoge/Doge-20M-Instruct")
model = AutoModelForCausalLM.from_pretrained("SmallDoge/Doge-20M-Instruct", device_map="auto")

generation_config = GenerationConfig(
    max_new_tokens=100,   # 最多新生成的 token 数
    use_cache=True,       # 开启 KV 缓存加速
    do_sample=True,       # 采样式解码
    temperature=0.8,      # 采样温度，越低越保守
    top_p=0.9,            # nucleus 采样累积概率阈值
    repetition_penalty=1.0  # 重复惩罚，1.0 表示不惩罚
)
streamer = TextStreamer(tokenizer=tokenizer, skip_prompt=True)

prompt = "Hi, how are you doing today?"
conversation = [
    {"role": "user", "content": prompt}
]
inputs = tokenizer.apply_chat_template(
    conversation=conversation,
    tokenize=True,
    return_tensors="pt",
)

outputs = model.generate(
    inputs,
    tokenizer=tokenizer,
    generation_config=generation_config,
    streamer=streamer
)

Doge 完整支持 GenerationMixin（DogeForCausalLM 同时继承 DogePreTrainedModel 与 GenerationMixin），所以 generate、beam search、流式输出等通用生成能力都可用。由于 keys_to_ignore_at_inference = ["past_key_values"]，在生成推理阶段 past_key_values 不会参与设备放置的 key 匹配，由 DynamicCache 正常管理缓存（见 DogeModel.forward）。

加载后训练 / 继续训练

Doge 预训练采用 wsd_scheduler 在 smollm-corpus 上进行；由于文档明确了“可以在新数据集上继续训练，或从稳定阶段 checkpoint 追加稀疏激活 FFN”，实际工作中常见的路径有两种：

• 普通继续预训练/微调：用 AutoModelForCausalLM.from_pretrained("SmallDoge/Doge-320M") 加载后，配合 Trainer 在自有语料上以较小的学习率继续训练即可；
• 从稠密权重升级为 CDMoE：将配置中 is_moe=True 打开，MoE 模块会继承原 MLP 的权重完成初始化（DogeConfig.is_moe 文档说明），从而在不重训共享层的前提下扩展稀疏专家规模。

DogeConfig 全参数详解

Doge 的核心配置类 DogeConfig 继承自 PreTrainedConfig，model_type = "doge"，并内置了张量并行（TP）与流水线并行（PP）默认切分方案。下表整理了 configuration_doge.py#L73-L100 中声明的全部默认值与含义：

参数	默认值	含义
vocab_size	32768	词表大小
hidden_size	1024	隐藏层维度
intermediate_size	2048	FFN/MoE 中间层维度
num_hidden_layers	32	解码器层数
num_attention_heads	8	注意力头数
num_key_value_heads	None	KV 头数，为空时在 __post_init__ 中回退为 num_attention_heads（即当前非 GQA 分组）
head_dim（可选）	由 hidden_size // num_attention_heads 推导	每头维度
max_position_embeddings	2048	位置编码最大长度
rope_parameters	None	RoPE 参数（如 rope_type、rope_theta），支持 Transformers 统一的动态 RoPE 更新机制 dynamic_rope_update
hidden_dropout	0.0	层内 dropout
attention_dropout	0.0	注意力 dropout
hidden_act	"silu"	激活函数
attention_bias / mlp_bias	False	各投影层是否带偏置
initializer_range	0.02	初始化范围
rms_norm_eps	1e-6	RMSNorm 的 epsilon
use_cache	True	是否启用 KV 缓存
tie_word_embeddings	False	是否绑定输入/输出嵌入权重
sliding_window	None	滑动窗口尺寸，None 表示使用完整因果掩码（源码依据 modeling_doge.py#L572：按此值选择 create_causal_mask 或 create_sliding_window_causal_mask）
keep_window_size	2048	DMA 关键参数：不被动态掩码的“常驻窗口”大小，仅当序列长度超过该值后才执行动态 top-k 稀疏掩码
is_moe	False	是否启用 Cross Domain MoE；为 True 时解码器 FFN 替换为 DogeCDMoE 且继承 MLP 初始化
num_experts	16384	MoE 专家总数
num_experts_per_tok	64	每个 token 激活的专家数（等效 top-k）
norm_topk_prob	False	是否对路由权重做归一化（按 top-k 求和缩放）
output_router_logits	False	是否输出路由 logits（用于计算辅助损失）
router_aux_loss_coef	0.001	路由负载均衡辅助损失系数
pad_token_id/bos_token_id/eos_token_id	None	特殊 token ID

其中两个最值得留意的参数是 keep_window_size 与 is_moe：前者直接控制 DMA 从“全因果”退化为“局部稀疏”的临界序列长度（由 prepare_dynamic_mask 中的 top-k 分支 实现）；后者决定前馈层是普通 DogeMLP 还是 DogeCDMoE，二者参数形状对齐，保证从稠密权重“继承初始化”成为可能。

自定义初始化配置与直接构建模型的方式：

from transformers import DogeConfig, DogeModel

# 初始化一个 Doge-320M 风格配置
configuration = DogeConfig()

# 由配置构建模型
model = DogeModel(configuration)

# 读取模型配置
configuration = model.config

配置类还预置了并行切分方案：TP 计划中 q_proj/k_proj/v_proj 按列切分、dt_proj/o_proj 按行切分、MoE 的 router_gate 为 colwise_gather_output、down_embed/up_embed 为 rowwise_split_input；PP 计划则将 embed_tokens/layers/norm/lm_head 依序划分。这些计划与 测试套件 中的 test_tp_plan_matches_params 一一对应，验证了 TP 切分与模型参数是吻合的。

提供下游任务的完整 API 与权重转换脚本

根据模型文档，本仓库为 Doge 提供四类公开 API，__all__ 声明在 modeling_doge.py#L816：

• DogeConfig：上文已详解的配置类；
• DogeModel：基础 transformer 解码器，输出 MoeModelOutputWithPast（含 last_hidden_state、past_key_values、可选 router_logits），支持 input_ids/inputs_embeds 二选一输入；
• DogeForCausalLM：语言建模头 + GenerationMixin，输出 MoeCausalLMOutputWithPast；forward 支持 labels（-100 位置的 token 不参与损失）、logits_to_keep（只算最后 N 个位置的 logits，节省显存）以及 output_router_logits。官方文档中的生成示例可复现，例如对 "Hey, are you conscious?..." 类 prompt 正常续写；
• DogeForSequenceClassification：继承 GenericForSequenceClassification，一行实现分类头，测试覆盖单标签/多标签两种场景（test_modeling_doge.py#L295-L334）。

如果你想复现官方 checkpoint 的转换流程，仓库提供了 convert_doge_weights_to_hf.py：它以 STATE_DICT_MAPPING 正则为映射表，把上游训练框架的命名（如 model.word_embed.weight、pre_layernorm、pre_residual、post_layernorm、feed_forward.gate_proj 等）逐一映射为 HF 命名（embed_tokens、input_layernorm、input_residual、post_attention_layernorm、mlp.gate_proj 等），并支持任意 num_hidden_layers 层编号的正则捕获。对照该映射表，可以清晰看到 Doge 模型在实现时对组件命名的最终约定，对理解“如何把自己的 Doge 权重转成可加载的 safetensors”很有参考价值。

测试覆盖与能力边界小结

在 tests/models/doge/test_modeling_doge.py 中，Doge 通过 ConfigTester、ModelTesterMixin、GenerationTesterMixin、PipelineTesterMixin 与 PipelineTesterMixin 系列通用测试获得基础保障，并有三个针对性用例值得关注：

• test_doge_sequence_classification_model(_for_single_label/_for_multi_label)：验证分类头的三种标签模式；
• test_tp_plan_matches_params：校验张量并行切分计划与实际参数一一对应；
• test_Doge_20M_hard（slow 测试）：直接用 SmallDoge/Doge-20M 权重跑硬性回归，是最贴近真实生产的端到端验证。

综合来看，把 Doge 装入 Transformers 体系时需要注意的能力边界是：官方不声明 Flash Attention 支持（_supports_flash_attn = False），支持的快速后端是 SDPA 与 Flex Attention；训练状态支持梯度 checkpoint（GradientCheckpointingLayer 基底与 supports_gradient_checkpointing=True），推理生成默认走 DynamicCache。在此边界内，Doge 可以通过 AutoModelForCausalLM/AutoTokenizer 的通用 API 直接完成加载、生成、微调与评估，是研究“稀疏注意力 + 稀疏专家小模型”组合的一个开箱即用的参考实现。

如果你希望进一步动手验证，可以在本仓库内运行 pytest tests/models/doge/test_modeling_doge.py 观察通用与 Doge 特有用例的执行结果，或阅读 modular_doge.py 与 Llama/Mixtral 对应模块的差异，体会模块化模型开发中“继承 + 覆写”的代码组织方式。