超曲面约束动态权重更新实验 [P]

大家好！我想分享一下我作为副项目运行的一个小型语言模型架构实验的结果。最终目标是测试一种减少模型训练参数数量的方法，因为训练的主要瓶颈是显存。核心思想与Universal Transformer非常相似。让我们只取一个解码器块，并在循环中迭代地将输入传递L次。但是，我的模型不是用循环迭代深度信息更新输入，而是动态更新基础层的权重。 它是如何实现的？模型使用一组学习的超曲面从超曲面的横截面生成权重增量（𝛥Wl）。因此，给定层l的权重矩阵构建为Wl = W0 + 𝛥Wl。超曲面由一系列周期函数定义。到目前为止，我得到的最佳结果是使用三角波。模型学习这些波在坐标维度上的振幅、频率和相位。给定这个集合E的大小，总共有3*E*dim个参数。最初，我尝试纯粹从超曲面生成完整权重，但这太受限，无法收敛。因此，我转向了当前的方法，其中模型有一个基础解码器层正在被更新。我还在前向传播中添加了一个状态向量，该向量是使用门控线性注意力从输入序列计算的。它用于调制超曲面的几何形状，使生成的权重增量𝛥Wl具有序列感知性。 预训练结果 我在FineWeb-Edu数据集的10B标记样本上运行了预训练实验。我希望保持训练参数计数最小，因此我使用了来自GPT-2的预训练、冻结的嵌入层。我也不使用任何位置编码，遵循NoPE方法。序列长度为1024个标记，批大小为16，训练运行了10,000步。 我使用了3个基线： - 一个简单的单层仅解码器Transformer - 一个跨24个循环迭代展开的单层Transformer - 一个标准的24层仅解码器Transformer 实验模型： - 一个使用正弦波表面增量的单个循环块，无上下文调制 - 一个使用三角波表面增量和上下文调制的单个循环块 - 一个使用三个堆叠循环块的模型，每个块利用三角波表面增量和上下文调制 这里是一个训练损失图表： 训练结果 模型参数计数 - 标准24层仅解码器Transformer：169,906,944 - 3个循环块+三角波+上下文调制：27,162,624（约基线模型大小的16%） 关键要点： 虽然经典的仅解码器架构仍然产生最佳的绝对损失，但三角表面+上下文模型显示出相对于标准未展开基线的真实性能提升。该模型引入了全新的参数和超参数集。我目前正在进行三角波函数集大小和上下文向量大小的实验。我还尝试了超曲面参数的不同初始化策略。 （非常模糊）目标不一定是击败全参数模型，而是找到一个在损失方面足够好，同时在硬件资源上大大减轻的配置。另一种可能性是超曲面表示具有不同的收敛动力学，需要更多的步骤和更多的标记来收敛。所以我的下一步是在FineWeb-Edu数据集的整个10B标记样本上训练一个具有3个循环块、更大上下文向量和扩展函数基的模型。 GitHub仓库：https://github.com/morgengramlich/LoopSLM