AdaPonderLM: Gated Pondering Language Models with Token-Wise Adaptive Depth

1. 自适应深度计算：AdaPonderLM引入了迭代特定的MLP门控和单调停止掩码，使循环语言模型能够根据token难度动态决定何时停止计算，从而实现token级别的早期退出。
2. KV状态复用：为确保训练-推理一致性并实现实际加速，模型设计了KV复用机制，对于已停止计算的token，其Key/Value状态会从前一迭代缓存中复用。
3. 性能与效率平衡：在Pythia系列模型（70M至2.8B）上，AdaPonderLM在保持与PonderLM相当的语言模型困惑度及下游任务准确率的同时，将推理计算量减少了约10%。
4. 自监督学习行为：分析显示，学习到的门控机制会将更多计算资源分配给高NLL（难）的token，展现出自适应计算时间行为，且在相同FLOPs下优于固定剪枝策略。

模型在Pythia 70M和410M架构上从头预训练，并在Pythia 1.4B和2.8B架构上进行继续预训练。预训练数据为从Pile数据集中采样的26B token，继续预训练使用300B-312B token。基线模型包括PonderLM、Pause Token、Loop Transformer以及Vanilla Pythia，同时与OPT、Bloom、Tinyllama、GPTNeo等大型预训练模型进行下游任务对比。评估指标包括语言模型困惑度（PPL）、损失以及LAMBADA、PIQA、WinoGrande、ARC-E/C、SciQ、HellaSwag、RACE等下游任务的零样本和五样本准确率。AdaPonderLM的推理FLOPs约为基线模型的3.7x-3.8x，而PonderLM等基线为4x。