【RustyML入门】3.7. 循环层

3.7. 循环层RustyML 提供 3 种循环层SimpleRNN、LSTM、GRU。它们都位于rustyml::neural_network::layers::recurrent并通过 prelude 重新导出。三者共用同一套约定每个层都接收一个三维序列张量沿时间轴从左到右跑一遍递归返回最后的隐藏状态。如果你熟悉 Keras可以把这些层理解成把return_sequencesFalse固定死的SimpleRNN、LSTM、GRU。这个固定设置会影响你如何堆叠这些层。搭建深层循环网络之前请先读 3.7.7 节。3.7.1. 输入/输出约定每个循环层都要求一个三维输入张量形状为(batch_size, timesteps, features)输出则是二维的(batch_size, units)。features轴必须等于你传给构造函数的input_dimunits就是你设定的隐藏层宽度。递归会把时间步轴吃掉所以它不会出现在输出里只有最后一个隐藏状态h_T会留下来。usendarray::Array;userustyml::neural_network::layers::activation::Tanh;userustyml::neural_network::layers::recurrent::SimpleRNN;userustyml::neural_network::traits::Layer;fnmain(){// input_dim 4 个每步特征units 3 个隐藏神经元letrnnSimpleRNN::new(4,3,Tanh::new()).unwrap();// (batch 2, timesteps 5, features 4)letxArray::zeros((2,5,4)).into_dyn();// predict 只跑递归不记录反向传播所需的缓存letoutrnn.predict(x).unwrap();// 时间步轴消失了只有最后一个隐藏状态留了下来。println!(output shape: {:?},out.shape());// [2, 3] (batch, units)}output shape: [2, 3]这 3 种层共用同一个构造函数签名new(input_dim, units, activation) - ResultSelf, Error。激活参数的类型是impl IntoActivation你可以传一个Activation枚举变体比如Activation::Tanh也可以传某个轻量的层包装比如Tanh::new()、ReLU::new()、Sigmoid::new()、Linear::new()、Softmax::new()。每一种包装都会转换成同一个枚举。完整的激活函数目录见 3.2. 全连接层与激活函数。input_dim或units为 0 时new返回Error::InvalidParameter。RustyML 没有return_sequences、return_state、bidirectional也没有 cell 内 dropout 选项。每个层永远只返回最后的状态永远沿时间正向推进永远使用同一套稠密实现。二维或四维输入是硬错误不会被悄悄 reshape只要输入不是三维forward和predict就返回Error::InvalidInput。层会缓存前向激活值供反向传播使用所以在forward之前调用backward会返回Error::NeuralNetwork(NnError::ForwardPassNotRun(SimpleRNN))对 LSTM、GRU 则是LSTM、GRU。这些错误变体如何配合见 1.6. 错误处理。RustyML 和 Keras 有一处不同activation参数只控制候选值或输出那一处非线性不涉及门。在SimpleRNN里这个激活作用于每一个隐藏状态在LSTM和GRU里它作用于候选值在 LSTM 里还会作用于输出门之前的细胞状态。门本身永远使用 sigmoid。RustyML 不像 Keras 那样提供独立的recurrent_activation选项门的非线性是固定的。Tanh 是默认的激活函数也是几乎所有已发表架构采用的选择。3.7.2. SimpleRNN 及门控单元存在的理由SimpleRNN就是教科书里的 Elman 递归从全零隐藏状态出发每个时间步用 2 个权重矩阵和 1 个偏置把当前输入和上一步的隐藏状态混合起来h_0 0 h_t activation( x_t W h_{t-1} U b ) for t 1..T output h_T这里W是输入核(input_dim, units)U是递归核(units, units)b是偏置(1, units)是沿 batch 维度进行的矩阵乘法。RustyML 用 Xavier/Glorot 均匀分布初始化W用正交矩阵Gram-Schmidt初始化U。正交递归核是刻意为之它让状态转移在初始化时保持范数不变这是延缓下文要讲的梯度消失问题的一种省事办法。这个问题就是梯度消失以及它的对立面梯度爆炸。从h_T反向传播到h_1每一步都要把上游梯度乘上一个新的雅可比矩阵大致是grad_{t-1} (activation(h_t) * grad_t) U^T。把T步这样的运算串起来相当于给一个矩阵求T次幂如果它的有效量级小于 1梯度就会指数衰减网络学不到往前超过几步的依赖关系如果量级大于 1梯度反而会爆炸。正交的U让U^T这个因子保持范数不变tanh 1又把乘积控制在有界范围内。即便如此这个乘积在长序列上依然会趋向于零。这正是SimpleRNN只能应付短序列至多几十步、碰到长程依赖就失效的原因也是 LSTM 和 GRU 存在的原因。3.7.3. LSTM一条加性的记忆高速路LSTM 增加了第二个状态细胞状态c_t它的更新是加性的。LSTM 用 3 个 sigmoid 门来决定写入什么、保留什么、读出什么。RustyML 把 4 个权重块并排融合存放列的顺序和 Keras 一致[input | forget | cell | output]记作[i | f | g | o]i_t sigmoid( x_t W_i h_{t-1} U_i b_i ) 输入门写入多少候选值 f_t sigmoid( x_t W_f h_{t-1} U_f b_f ) 遗忘门保留多少旧细胞状态 g_t act( x_t W_g h_{t-1} U_g b_g ) 候选值cell gate o_t sigmoid( x_t W_o h_{t-1} U_o b_o ) 输出门暴露多少细胞状态 c_t f_t * c_{t-1} i_t * g_t 细胞状态加性更新 h_t o_t * act(c_t) 隐藏状态这里*是逐元素乘法act是可配置的激活函数默认 Tanh同时作用于候选值和细胞状态。关键的一行是c_t f_t * c_{t-1} i_t * g_t它对c_{t-1}的梯度就是f_t一次逐元素乘法没有反复的矩阵乘法。当遗忘门打开也就是f接近 1 时细胞状态几乎无损地把梯度往回传是一条近乎恒等的高速路加性项则不断汇入这条路。记忆之所以留得住、梯度之所以流得动是因为主干路径是加法而不是反复的矩阵相乘。RustyML 把遗忘门的偏置初始化为 1.0其余偏置全部为零这给这条记忆高速路一个先发优势。在训练把门调整成形之前f就已经偏向打开于是细胞状态连同它的梯度从第一个 epoch 起就能存活下来。集成测试lstm_forget_bias_is_one_not_zero就是用来检验这一行为的。LSTM 的参数量相当于 4 个门param_count 4 * (input_dim * units units * units units)。3.7.4. GRU同一思路门更少GRU 保留了 LSTM 的加性混合思路把输入门和遗忘门合并成一个更新门并去掉了独立的细胞状态所以只需要 3 个权重块而不是 4 个。RustyML 把它们融合存放顺序是[update | reset | candidate]记作[z | r | h]这个顺序和 Keras 一致z_t sigmoid( x_t W_z h_{t-1} U_z b_z ) 更新门 r_t sigmoid( x_t W_r h_{t-1} U_r b_r ) 重置门 n_t act( x_t W_h (r_t * h_{t-1}) U_h b_h ) 候选值 h_t z_t * h_{t-1} (1 - z_t) * n_t 隐藏状态更新门z_t做的是一次凸混合。当z接近 1 时层会把上一步的隐藏状态原封不动地拷过来这是一条梯度高速路和关闭的 LSTM 遗忘门效果一样。当z接近 0 时层会用新的候选值替换掉上一步的状态。这是 Keras 的约定。有些资料用的是它的补从那类资料里搬过来的z需要翻转。测试gru_update_gate_one_keeps_previous_hidden和gru_update_gate_zero_takes_the_candidate分别检验这两个极端测试gru_fused_kernel_first_block_is_the_update_gate则检验列的顺序。要留意重置门作用的确切位置RustyML 在候选值的递归矩阵乘法之前就算好了r_t * h_{t-1}即(r_t * h_{t-1}) U_h。这对应 Cho 等人最初的公式也就是 Keras 的reset_afterFalse而不是 CuDNN 的reset_afterTrue变体后者在矩阵乘法之后才施加重置并且每个门需要 2 套偏置。这里每个门只有一套偏置。GRU 的param_count 3 * (input_dim * units units * units units)是同等宽度 LSTM 的四分之三。实践中 GRU 训练略快在很多任务上和 LSTM 打平当任务需要一段长而精准可控的记忆时LSTM 有时表现更好。两者的输入核都用 Xavier/Glorot 初始化用的是每个门各自的扇入input_dim units而不是融合后的宽度每个门的递归块都各自初始化成一个独立的正交矩阵。3.7.5. 权重、形状以及手动设置可训练张量及其形状层kernelrecurrent_kernelbias融合的列块SimpleRNN(input_dim, units)(units, units)(1, units)无LSTM(input_dim, 4 * units)(units, 4 * units)(1, 4 * units)[i | f | g | o]GRU(input_dim, 3 * units)(units, 3 * units)(1, 3 * units)[z | r | h]把每个门融合进一个矩阵不只是好看它让输入投影和递归投影在每个时间步都能跑成一个大 GEMM而不是每个门各跑一个这在缓存和 SIMD 上带来实打实的收益。你可以用get_weights()检视这些实时数组它返回一个LayerWeight::{SimpleRNN,LSTM,GRU}值里面带着借用的kernel、recurrent_kernel、bias字段userustyml::neural_network::layers::activation::Tanh;userustyml::neural_network::layers::layer_weight::LayerWeight;userustyml::neural_network::layers::recurrent::LSTM;userustyml::neural_network::traits::Layer;fnmain(){// input_dim 4units 8。with_random_state 让初始化可复现。letlstmLSTM::new(4,8,Tanh::new()).unwrap().with_random_state(42);matchlstm.get_weights(){LayerWeight::LSTM(w){// 4 个门并排融合在一起宽度 4 * units。println!(kernel {:?},w.kernel.shape());// [4, 32]println!(recurrent_kernel {:?},w.recurrent_kernel.shape());// [8, 32]println!(bias {:?},w.bias.shape());// [1, 32]}_unreachable!(),}}kernel [4, 32] recurrent_kernel [8, 32] bias [1, 32]想要可复现的初始化调用with_random_state(seed)它会以确定的方式重跑核与递归核的采样遗忘门偏置为 1.0 的规则依然生效。不传种子的话权重就从全局种子或系统熵取种。见 7.1. 可复现性与随机种子。要手动装入权重比如从别的框架移植过来或者对一段精确的递归做单元测试每个层都提供set_weights(kernel, recurrent_kernel, bias)直接接收融合后的矩阵。LSTM 和 GRU 还额外提供set_gate_weights(...)它按门接收一组(kernel, recurrent_kernel, bias)三元组并替你拼接成融合布局。按门传参的顺序LSTM 是(input, forget, cell, output)共 12 个数组GRU 是(reset, update, candidate)共 9 个数组。任何对不上的形状都会返回Error::NeuralNetwork(NnError::WeightShape { .. })。保存或加载整个模型时这些数组会原样序列化。见 3.9. 权重保存与加载。3.7.6. BPTT 与开销模型训练用的是随时间反向传播BPTT而且是完整展开没有截断窗口。前向传播时层会把反向传播需要的一切都缓存下来SimpleRNN存每一个隐藏状态前面补上h_0 0LSTM还会存细胞状态、activation(c_t)以及每个时间步的 4 个门激活值GRU存重置门、更新门、候选值以及r_t * h_{t-1}这个乘积。predict给这些缓存传None跳过记录和相应的克隆这就是推理比训练时的forward更省的原因。内存开销随timesteps线性增长长序列吃的是内存不只是时间。backward反向遍历时间步。它要求一个二维的上游梯度(batch, units)也就是损失对最终隐藏状态的梯度这是层唯一需要的梯度因为层对外吐出来的就只有这个最终状态。每一步backward都会算出该时间步的预激活梯度把grad_hLSTM 还有grad_c串回上一步这部分天生沿时间串行。接着backward把权重梯度的归约批量化把逐时间步的dz在(batch, timesteps)上折叠起来kernel 和 bias 梯度各自从一个大 GEMM 里算出。递归核的梯度通常也是这样但 GRU 例外它有 2 个不同的递归输入要用 2 个 GEMM 才能算出递归核的梯度。RustyML 以替换语义存储梯度层内部不做裁剪如果需要梯度裁剪用一个提供该功能的优化器。见 3.4. 优化器。要记住的性能形态是沿时间轴串行沿 batch 轴和融合的门轴并行。输入投影x W不依赖递归RustyML 会一次性把它算完作为一个跨所有时间步的批量 GEMM。只有h_{t-1} U这一项必须一步一步来而其中每一步本身又是一个 batch 并行的 GEMM。GRU 在这上面还能再省一点它把重置门和更新门的递归投影融合成一个 GEMM因为两者都要读h_{t-1}。只有候选值的递归投影仍然单独计算因为它的输入r_t * h_{t-1}依赖刚算好的重置门。每一次矩阵乘法都交给 gemmkit 后端见 6.2. 矩阵乘法。gemmkit 会根据工作量自行决定并行规模宽层和大 batch 会自动获得线程并行小层则保持串行以避免额外开销。一个时间步里融合的门投影总会在计算乘积的同一趟里施加偏置。只有SimpleRNN搭配ReLU时激活函数才会融合进这一趟里LSTM、GRU以及其他任何激活函数都不会融合。实用的结论是更多序列也就是更大的 batch能很好地并行更多时间步则不行因为那个轴是一条串行的依赖链。线程池的各项设置见 7.3. 性能调优与并行。3.7.7. 堆叠与构建模型循环层只返回最后一个隐藏状态一个二维的(batch, units)张量。你不能把这个输出直接喂给另一个循环层循环层要的是三维的(batch, timesteps, features)输入面对二维张量只会返回Error::InvalidInput。RustyML 没有return_sequences选项层没法吐出一条逐时间步的序列所以 Keras 意义上那种循环层的深层堆叠在这里搭不出来。请提前为这个限制做好规划。标准的写法是拿一个循环层当序列编码器后面接一个稠密头把最终状态映射到目标上。这样组合起来很干净循环层把(batch, timesteps, features)变成(batch, units)Dense要的恰好就是这个二维形状。下面的示例学的是一个真实的序列任务预测一条长度为 3 的标量序列之和。它用LSTM编码器、一个Dense读出层、Adam、均方误差搭起来能在几百个全批量 epoch 内收敛回想 3.1. Sequential模型 中提到fit每个 epoch 只走一步全批量梯度usendarray::Array;userustyml::neural_network::sequential::Sequential;userustyml::prelude::*;fnmain(){// 4 条序列3 个时间步1 个特征。目标 3 个标量之和。letxArray::from_shape_vec((4,3,1),vec![0.1,0.2,0.1,// 和为 0.40.3,0.1,0.2,// 和为 0.60.0,0.2,0.2,// 和为 0.40.2,0.2,0.1,// 和为 0.5],).unwrap().into_dyn();letyArray::from_shape_vec((4,1),vec![0.4,0.6,0.4,0.5]).unwrap().into_dyn();letmutmodelSequential::new();model// 循环特征提取器(batch, 3, 1) - (batch, 16).add(LSTM::new(1,16,Tanh::new()).unwrap().with_random_state(42))// 把最后的隐藏状态读出成单个标量。.add(Dense::new(16,1,Linear::new()).unwrap()).compile(Adam::new(0.01,0.9,0.999,1e-8,0.0).unwrap(),MeanSquaredError::new(),);model.fit(x,y,300).unwrap();letpredmodel.predict(x).unwrap();println!(target : {:?},y.as_slice().unwrap());println!(prediction : {:?},pred.as_slice().unwrap());}300 个 epoch 之后4 个预测值都落在[0.4, 0.6, 0.4, 0.5]附近的几个百分点以内。LSTM 学会了累加这条序列稠密头再把累加结果读出来。把LSTM换成GRU或SimpleRNN同一套代码照样能编译、能训练在这条短序列上3 种层都会收敛这恰恰说明门控单元在梯度消失上的优势只有在长序列上才会显现出来。model.summary()会把每个循环层的输出形状打印成(None, units)其中None代表动态的 batch 维度。数据集只要比玩具规模大一些就应该用fit_with_batches代替fit这样每个 epoch 会走好几步小批量并重新洗牌数据。