在常见的神经网络中，神经元通常都是无序的，比如遗忘门 f t f_t ft是一个向量，向量的各个元素的位置没有什么规律。
ON-LSTM则试图把这些神经元排个序，并且用这个序来表示一些特定的结构，从而把神经元的序信息利用起来。比如说，c有6个维度，按顺序编号为0,1,2,3,4,5。接近0的维度存储高层级的语义，接近5的维度则存储低层次语义。
举个例子：
上图是一个句子的语法树结构，一共有5层，所以我们用一个5维的c来存储上述结构（姑且用一个值就能表示语义特征）。每一层都有其相应的语义信息，即便是语句层。比如说，你这句话是不是一条完整的语句，这是语句层可以提供的信息。
当我们输入The的时候。语句层的语义是“这不是一条完整的语句”。即便输入professor trains Jack后语义依然没有更新，直到接收到"."之后才更新为“这是一条完整的语句”。所以我们可以发现，高层级的语义通常更新较缓慢
最后一层是单词序列的语义层，显然大多数时候每输入一个新单词语义都会更新，因此低层级的语义更新频率较大
ON-LSTM正是依据高层级语义更新频率小，低层级的语义更新频率大这一特点，学习出了具有层次的记忆。
ON-LSTM首先引入了一个变量 d i d_i di，表示当前信息所包含语义的最高层次。比如说，当我们输入The的时候， d i d_i di为4，表示我们需要更新全部的历史信息。当输入professor的时候， d i d_i di为1，表示只需要更新历史信息的第0层和第1层。因为当前信息影响不到第2层及以上的语义，所以这些层不需要更新。
我们可能觉得疑惑的地方是，为什么输入The和输入professor，包含语义层级的差别这么大？原因就在于，我们评估当前信息所包含语义的最高层级不仅需要考虑当前信息，还需要考虑历史信息。其实这很合理，我们判断当前信息的层级的办法相当于是将当前信息与历史信息的每个层级的语义计算相关性，如果相关，那么说明当前信息就包含该层级，否则就不包含。开始输入The的时候，历史信息为空，我们默认都相关，所以 d i = 4 d_i=4 di=4。输入professor的时候，The和冠词都与professor相关，所以 d i = 1 d_i=1 di=1
除了 d i d_i di之外，ON-LSTM还引入了一个新的变量 d f d_f df，表示历史信息所包含语义的最低层级。
这是什么意思呢？我举个例子，trains这个词和Jack从单词本身来看并没有联系，真正将他们联系起来的是他们的语法角色。trains是一个动词，Jack是一个名词，二者正好组成一个动词短语。所以说，trains对Jack的影响没有单词层面的语义影响，而是在更高层次上。而这种情况是很常见的。这也就是说，历史信息可以把不需要的底层的信息遗忘掉，只需要保留高层语义即可。
ON-LSTM正是设计了一个符合上述过程的机制，通过变量 d f d_f df来确定哪些层级的语义是跟当前输入不相关的，从而把它遗忘掉。
首先，根据当前历史信息和当前信息计算出 d i d_i di和 d f d_f df以及传统LSTM的三个门 f t , i t , o t f_t,i_t,o_t ft,it,ot和输入 u t u_t ut
如果 d i ≥ d f d_i\ge d_f di≥df，表示当前信息包含语义最高层级高于历史信息包含语义最低层级，对于高于 d f d_f df的部分，保留历史信息不用更新，对于低于 d i d_i di的部分，由于历史信息没有意义了，所以直接用当前信息相应层级覆盖掉，而对于 d f d_f df到 d i d_i di之间的交集部分，表示既有历史信息，也有当前输入，所以按照原始更新方式进行，即融合历史信息与当前信息作为新的历史信息。
如果 d i < d f d_i< d_f di<df，对于高于 d f d_f df的部分，保留历史信息不用更新，对于低于 d i d_i di的部分，用当前信息相应层级覆盖掉，对于 d i d_i di到 d f d_f df的部分，因为既没有信息输入，原来的历史信息又没用，所以就置零，这样当有新信息进来的时候就会默认相关，存储进新的历史信息。而如果不做处理的话这部分空间可能永远都得不到更新了。
我们要如何用公式来表示上面的过程？
如果直接预测 d i d_i di和 d f d_f df两个整数，不太容易限定范围。所以我们用预测一个长度与c相同的one-hot张量来代表 d i d_i di和 d f d_f df
比如
d
i
=
3
d_i=3
di=3，
d
f
=
1
d_f=1
df=1就相当于
d
i
=
[
0
,
1
,
0
,
0
,
0
]
d
f
=
[
0
,
0
,
0
,
1
,
0
]
d_i = [0,1,0,0,0] \\ d_f = [0,0,0,1,0]
di=[0,1,0,0,0]df=[0,0,0,1,0]
当前信息的层级区域和历史信息的层级区域就可以表示为
i
~
t
=
c
u
m
s
u
m
→
d
i
=
[
0
,
1
,
1
,
1
,
1
]
f
~
t
=
c
u
m
s
u
m
←
d
f
=
[
1
,
1
,
1
,
1
,
0
]
\begin{aligned} \tilde i_t &= \operatorname{\overrightarrow{cumsum}}d_i = [0,1,1,1,1] \\ \tilde f_t &= \operatorname{\overleftarrow{cumsum}}d_f = [1,1,1,1,0] \end{aligned}
i~tf~t=cumsum
di=[0,1,1,1,1]=cumsum
df=[1,1,1,1,0]
当前信息和历史信息的层级交集可表示为
w
t
=
i
~
t
∗
f
~
t
=
[
0
,
1
,
1
,
1
,
0
]
w_t = \tilde i_t * \tilde f_t = [0,1,1,1,0]
wt=i~t∗f~t=[0,1,1,1,0]
当前信息除去交集部分表示为
i
~
t
−
w
t
=
[
0
,
0
,
0
,
0
,
1
]
\tilde i_t-w_t = [0,0,0,0,1]
i~t−wt=[0,0,0,0,1]
历史信息除去交集部分表示为
f
~
t
−
w
t
=
[
1
,
0
,
0
,
0
,
0
]
\tilde f_t-w_t = [1,0,0,0,0]
f~t−wt=[1,0,0,0,0]
我们会发现上述表示跟
d
i
d_i
di和
d
f
d_f
df的大小无关，这里就不赘述了。
这样一来，ON-LSTM的c的更新式为：
c
t
=
w
t
⋅
(
f
t
∘
c
t
−
1
+
i
t
∘
u
t
)
+
(
f
~
t
−
w
t
)
⋅
c
t
−
1
+
(
i
~
t
−
w
t
)
⋅
u
t
c_t = w_t \cdot (f_t \circ c_{t-1}+i_t \circ u_t) + (\tilde f_t-w_t)\cdot c_{t-1}+(\tilde i_t-w_t)\cdot u_t
ct=wt⋅(ft∘ct−1+it∘ut)+(f~t−wt)⋅ct−1+(i~t−wt)⋅ut
现在的问题之一是
d
i
d_i
di和
d
f
d_f
df都是one-hot，不方便求梯度，所以我们用softmax软化
p
i
=
softmax
(
W
i
~
x
t
+
U
i
~
h
t
−
1
+
b
i
~
)
p
f
=
softmax
(
W
f
~
x
t
+
U
f
~
h
t
−
1
+
b
f
~
)
\begin{aligned} p_i &= \operatorname{softmax}\left(W_{\tilde{i}} x_{t}+U_{\tilde{i}} h_{t-1}+b_{\tilde{i}}\right) \\ p_f &= \operatorname{softmax}\left(W_{\tilde{f}} x_{t}+U_{\tilde{f}} h_{t-1}+b_{\tilde{f}}\right) \\ \end{aligned}
pipf=softmax(Wi~xt+Ui~ht−1+bi~)=softmax(Wf~xt+Uf~ht−1+bf~)
然后我们定义
cummax
=
cumsum
(
softmax
)
\operatorname{cummax} = \operatorname{cumsum}(\operatorname{softmax})
cummax=cumsum(softmax)，于是有
i
~
t
=
c
u
m
m
a
x
→
(
W
i
~
x
t
+
U
i
~
h
t
−
1
+
b
i
~
)
=
[
0
,
1
,
1
,
1
,
1
]
f
~
t
=
c
u
m
m
a
x
←
(
W
f
~
x
t
+
U
f
~
h
t
−
1
+
b
f
~
)
=
[
1
,
1
,
1
,
1
,
0
]
\begin{aligned} \tilde i_t &= \operatorname{\overrightarrow{cummax}}\left(W_{\tilde{i}} x_{t}+U_{\tilde{i}} h_{t-1}+b_{\tilde{i}}\right) = [0,1,1,1,1] \\ \tilde f_t &= \operatorname{\overleftarrow{cummax}}\left(W_{\tilde{f}} x_{t}+U_{\tilde{f}} h_{t-1}+b_{\tilde{f}}\right) = [1,1,1,1,0] \end{aligned}
i~tf~t=cummax
(Wi~xt+Ui~ht−1+bi~)=[0,1,1,1,1]=cummax
(Wf~xt+Uf~ht−1+bf~)=[1,1,1,1,0]
其它公式保持不变，这样一来，ON-LSTM的全部更新式如下
f
t
=
σ
(
W
f
x
t
+
U
f
h
t
−
1
+
b
f
)
i
t
=
σ
(
W
i
x
t
+
U
i
h
t
−
1
+
b
i
)
o
t
=
σ
(
W
o
x
t
+
U
o
h
t
−
1
+
b
o
)
u
t
=
tanh
(
W
c
x
t
+
U
c
h
t
−
1
+
b
c
)
f
~
t
=
c
u
m
m
a
x
←
(
W
f
~
x
t
+
U
f
~
h
t
−
1
+
b
f
~
)
i
~
t
=
c
u
m
m
a
x
→
(
W
i
~
x
t
+
U
i
~
h
t
−
1
+
b
i
~
)
ω
t
=
f
~
t
∘
i
~
t
c
t
=
ω
t
∘
(
f
t
∘
c
t
−
1
+
i
t
∘
u
t
)
+
(
f
~
t
−
ω
t
)
∘
c
t
−
1
+
(
i
~
t
−
ω
t
)
∘
u
t
h
t
=
o
t
∘
tanh
(
c
t
)
\begin{aligned} f_{t} &=\sigma\left(W_{f} x_{t}+U_{f} h_{t-1}+b_{f}\right) \\ i_{t} &=\sigma\left(W_{i} x_{t}+U_{i} h_{t-1}+b_{i}\right) \\ o_{t} &=\sigma\left(W_{o} x_{t}+U_{o} h_{t-1}+b_{o}\right) \\ u_{t} &=\tanh \left(W_{c} x_{t}+U_{c} h_{t-1}+b_{c}\right) \\ \tilde{f}_{t} &=\operatorname{\overleftarrow {cummax}}\left(W_{\tilde{f}} x_{t}+U_{\tilde{f}} h_{t-1}+b_{\tilde{f}}\right) \\ \tilde{i}_{t} &=\operatorname{\overrightarrow {cummax}}\left(W_{\tilde{i}} x_{t}+U_{\tilde{i}} h_{t-1}+b_{\tilde{i}}\right) \\ \omega_{t} &=\tilde{f}_{t} \circ \tilde{i}_{t} \\ c_{t} &=\omega_{t} \circ\left(f_{t} \circ c_{t-1}+i_{t} \circ u_{t}\right)+\left(\tilde{f}_{t}-\omega_{t}\right) \circ c_{t-1}+\left(\tilde{i}_{t}-\omega_{t}\right) \circ u_{t} \\ h_{t} &=o_{t} \circ \tanh \left(c_{t}\right) \end{aligned}
ftitotutf~ti~tωtctht=σ(Wfxt+Ufht−1+bf)=σ(Wixt+Uiht−1+bi)=σ(Woxt+Uoht−1+bo)=tanh(Wcxt+Ucht−1+bc)=cummax
(Wf~xt+Uf~ht−1+bf~)=cummax
(Wi~xt+Ui~ht−1+bi~)=f~t∘i~t=ωt∘(ft∘ct−1+it∘ut)+(f~t−ωt)∘ct−1+(i~t−ωt)∘ut=ot∘tanh(ct)
上面的关于 f ~ t \tilde f_t f~t的计算跟苏剑林大神保持一致，原论文是
f ~ t = 1 − cummax ( W f ~ x t + U f ~ h t − 1 + b f ~ ) \begin{aligned} \tilde f_t &= 1-\operatorname{cummax}\left(W_{\tilde{f}} x_{t}+U_{\tilde{f}} h_{t-1}+b_{\tilde{f}}\right) \end{aligned} f~t=1−cummax(Wf~xt+Uf~ht−1+bf~)
个人认为原论文是有一点不合理的，假如我们要表示历史信息的最低层级是第一层，那么就要求 f ~ t \tilde f_t f~t都接近1，也就是说 cummax ( . . . ) \operatorname{cummax}(...) cummax(...)需要全接近0，也就是说 softmax(...) \operatorname{softmax(...)} softmax(...)需要全接近0，而且所有维度的值加起来还要接近0。这显然是不可能的，因为最左维度值一定是1。
也就是说原论文的公式无法表示历史信息包含全部层级的情况，当然应用在实际上可能影响不大。
现在还有第二个问题是通常我们的层级是非常少的，一般4、5层可能就差不多了，而c的维度通常是非常大的，这就导致 f ~ t ⋅ c t − 1 \tilde f_t\cdot c_{t-1} f~t⋅ct−1无法直接计算。那要如何处理呢？难道我们就把LSTM的hidden_size设为4、5层？一个值能表示这么复杂的语义吗？显然是不能的。
所以ON-LSTM的做法是这样的，例如历史信息有8个维度，共划分4个层级，那么 f ~ t \tilde f_t f~t就是一个4维的张量，假如是[0,1,0,0]。我们把它的每个神经元重复2次，即可得到[0,0,1,1,0,0,0,0]，然后用这个新的张量去更新历史信息c。这就相当于我们将整个c的8个维度分成了4份，每一份包含两个维度，我们用每两个维度建模一个层级语义，而这两个维度我们用 f ~ t \tilde f_t f~t的相应层级来控制。
ON-LSTM导出句法树是使用 d f d_f df来导出的，它的主要思想是说，历史信息最低层级发生变化的时候，就表示这是一个新的语法结构的开始，并且历史信息最低层级越高表示新的语法结构的层次越高
这里借用蝈蝈的例子，根据历史层级的变化，x4的层级最高，所以x4处是一个较高层次的新的语法结构的开始。所以可以分成两部分
然后对x1,x2,x3和x4,x5可以做同样的处理，最终得到
然后我们要如何根据
f
~
t
\tilde f_t
f~t获得
d
f
d_f
df呢？
这里本来
d
f
=
a
r
g
m
a
x
k
p
f
(
k
)
d_f = \underset{k}{argmax}\ p_f(k)
df=kargmax pf(k)，上面的公式是argmax软化后的式子，具体可以看苏剑林的博客
但是我在蝈蝈的博客下面评论了一个问题，是我没搞懂的一个地方。就是在输入一句话的第一个词的时候，这个时候没有历史信息，按照博客中说的，把历史信息辐射范围设为0，那么随意输入一串序列，第一个词的历史信息最低层级都为0，那么要如何参与句法树的导出呢？
蝈蝈.《ON-LSTM：能表示语言层次的LSTM》. https://zhuanlan.zhihu.com/p/69314959
苏剑林. (May. 28, 2019). 《ON-LSTM：用有序神经元表达层次结构 》. https://spaces.ac.cn/archives/6621
