(Neural Turning Machine)神經圖靈機理解和pytorch實現

一般的神經網絡不具有記憶功能，輸出的結果只基於當前的輸入；而LSTM網絡的出現則讓網絡有了記憶：能夠根據之前的輸入給出當前的輸出。但是，LSTM的記憶程度並不是那麼理想，對於比較長的輸入序列，LSTM的最終輸出只與最後的幾步輸入有關，也就是long dependency問題，當然這個問題可以由注意力機制解決，然而卻不能從根本上解決長期記憶的問題，原因是由於LSTM是假設在時間序列上的輸入輸出：由t-1時刻得到t時刻的輸出，然後再循環輸入t時刻的結果得到t+1時刻的輸出，這樣勢必會使處於前面序列的輸入被淹沒，導致這部分記憶被“丟掉“。

神經圖靈機通過引入外部記憶解決了這個問題。 舉個簡單的例子，我們人類在記憶一些事情的時候，除了用腦袋記，還會寫在備忘錄上，當我們想不起來的時候，就可以去翻閱備忘錄，從而獲得相關的記憶。神經圖靈機模仿人類記憶的過程：其中的控制器(controller)相當於我們人類的大腦，用於把輸入事物的特徵提取出來；外部記憶(memory)相當於我們的備忘錄，把事物的特徵記錄在上面，那麼完整的過程就是：控制器將當前輸入轉化爲特徵，寫入記憶，再讀取與當前輸入特徵有關的記憶作爲最後的輸出。整個過程與圖靈機的讀寫很像，只不過神經圖靈機這裏讓所有的讀寫操作都可微分化，因此可以用神經網絡誤差後向傳播的方式去訓練模型。

那麼問題就來了，當獲得一個輸入的特徵後，我們如何確定在記憶中儲存的位置，而且如何從記憶中獲取與當前輸入相關的信息呢？這就是接下來要分析的神經圖靈機主要工作。

Neural Turning Machine

1. 讀記憶 (Read Heads)

我們把記憶看作是一個N×M$N\times M$ 的矩陣Mt$M_t$ ，t表示當前時刻，表示記憶會隨着時間發生變化。我們的讀過程就是生成一個定位權值向量wt$w_t$ ，長度爲N$N$ ，表示N個位置對應的記憶權值大小，最後讀出的記憶向量rt$r_t$ 爲：

rt=∑iNwt(i)Mt(i)$$r_t=\sum _i^N{w}_t(i)M_t(i)$$

其中權值向量的和爲1：∑iwt(i)=1$\sum _i{w}_t(i)=1$ ，本質上是一個對N條記憶進行一個加權求和的思想。

2. 寫記憶(Write Heads)

神經圖靈機的寫過程參考了LSTM的門的概念：先用輸入門決定增加的信息，再用遺忘門決定要丟棄的信息，最後用更新門加上增加的信息並減去丟棄的信息。具體來說，神經圖靈機會生成一個擦除向量et$e_t$ (erase vector) 和一個增加向量at$a_t$ (add vector)，長度都爲N$N$ ，向量中每個元素的值大小範圍從0到1，表示要增加或者刪除的信息。對於寫記憶過程，神經圖靈機首先執行一個擦除操作，擦除程度的大小同樣由向量wt$w_t$ 決定：

M′t(i)=Mt−1(i)(1−wt(i)et(i))$$M_t^{\prime }(i)=M_{t-1}(i)(1-w_t(i)e_t(i))$$

這個操作表示從t−1$t-1$ 時刻的記憶中丟棄了一些信息，若wt$w_t$ 和et$e_t$ 同時爲0，則表示記憶沒有丟棄信息，當前記憶與t−1$t-1$ 時刻保持不變。執行完擦除後，然後執行增加操作：

Mt(i)=M′t(i)+wt(i)at(i)$$M_t(i)=M_t^{\prime }(i)+w_t(i)a_t(i)$$

這步表示在丟棄一些信息後需要新增的信息，同樣，若wt$w_t$ 和at$a_t$ 都爲0，表示當前記憶無新增，與擦除後的記憶保持一致。其中，et$e_t$ 和at$a_t$ 都是由控制器給出，而控制器基本上由神經網絡實現，可以是LSTM，也可以是MLP。

由於整個過程都是都是矩陣的加減乘除，所有的讀寫操作都是可微分的，因此我們可以用梯度下降法訓練整個參數模型。但是接下來，我們需要確定wt$w_t$ 定位向量，由於這個向量直接決定着當前輸入與記憶的相關性，因此神經圖靈機在生成wt$w_t$ 向量上做了很多工作。

3. 定位機制(Addressing Mechanism)

關於決定其相關性的方法有很多，主要分爲兩大類： 基於內容的(content-based)和基於位置的(location-based)。神經圖靈機結合了這兩個方法提出一個定位機制用於生成定位向量wt$w_t$ ，具體來說，先用基於內容的方法，再用基於位置的方法。

3.1 Content-based Addressing

基於內容的定位計算主要基於餘弦相似度：首先控制器給出一個kt$k_t$ 向量作爲查詢的key，然後計算kt$k_t$ 與Mt$M_t$ 中各個記憶向量的餘弦相似度，最後經過一個softmax操作得到基於內容的定位向量wct$w_t^c$ ：

wct(i)=exp(βtK[kt,Mt(i)])∑jexp(βtK[kt,Mt(j)])$$w_t^c(i)=\frac{\exp ({\beta }_{t}K[k_t,M_t(i)])}{\sum _j\exp ({\beta }_{t}K[k_t,M_t(j)])}$$

其中K[.,.]$K[.,.]$ 是餘弦相似度計算：

K[u,v]=u⋅v||u||⋅||v||$$K[u,v]=\frac{u\cdot v}{||u||\cdot ||v||}$$

3.2 Location-based Addressing

3.2.1. Interpolation(插值)

控制器生成一個閾值gt$g_t$ 對當前的內容定位向量wct$w_t^c$ 與t-1時刻的定位向量wt−1$w_{t-1}$ 進行一個插值操作，插值的結果即爲輸出值wgt$w_t^g$ ：

wgt=gtwct+(1−gt)wt−1$$w_t^g=g_t{w}_t^c+(1-g_t)w_{t-1}$$

這裏的插值操作可以理解爲LSTM的更新門，結合過去的w$w$ 權值計算新的w$w$

3.2.2. shift(偏移)

對於wgt$w_t^g$ 中的每個位置元素wgt(i)$w_t^g(i)$ ，我們考慮它相鄰的k個偏移元素，認爲這k個元素與wgt(i)$w_t^g(i)$ 相關，如當k=3時，三個相鄰元素分別是：wgt(i)$w_t^g(i)$ 本身和位置偏移爲1的元素wgt(i−1)$w_t^g(i-1)$ 和wgt(i+1)$w_t^g(i+1)$ ，此時，我們希望新的位置爲i的元素能包含這三個元素，因此用一個長度爲3的偏移權值向量st$s_t$ 來表示這三個元素的權重，然後權值求和得到輸出值w′t$w_t^{\prime }$ :

w′t(i)=∑j=−11wgt(i+j)s(j+1)$$w_t^{\prime }(i)=\sum _{j=-1}^1{w}_t^g(i+j)s(j+1)$$

這裏的偏移操作在原文中用的是循環卷積(circular convolution)公式表示的，我們可以理解爲把向量wgt$w_t^g$ 首尾相連形成一個環狀，然後在環中用st$s_t$ 作爲卷積核做一維卷積操作。本質上是假設當前元素與相鄰的偏移元素相關。

3.2.3. Sharping(重塑)

當偏移操作中的權值比較平均的時候，上述的卷積操作會導致數據的分散(dispersion)和泄漏(leakage)，就像把一個點的信息分散在三個點中，權值如果太平均會使三個點包含的值太模糊（個人理解），因此需要把權值大小的區別進行強化，也就是sharping。具體來說，控制器生成一個參數γt>1${\gamma }_t>1$ ，然後對各個權值進行γt${\gamma }_t$ 指數然後歸一化：

wt(i)=w′t(i)γt∑jw′t(j)γt$$w_t(i)=\frac{w_t^{\prime }(i{)}^{{\gamma }_t}}{\sum _j{w}_t^{\prime }(j{)}^{{\gamma }_t}}$$

最後我們得出了最終的wt$w_t$ 用於提取和儲存記憶。

Pytorch實現

這裏代碼基於的是pytorch-ntm，代碼寫的相當工整，可讀性很高，這裏只分析一些重要的步驟：

讀過程

讀過程就是從控制器（LTSM）輸出的值提取我們需要的k, beta, g, s, gama值，然後調用_address_memory獲得當前的定位權值向量w, 再用矩陣乘法獲得讀過程的輸出

def forward(self, embeddings, w_prev):
    """NTMReadHead forward function.

    :param embeddings: input representation of the controller.
    :param w_prev: previous step state
    """
    o = self.fc_read(embeddings)
    k, beta, g, s, gama = _split_cols(o, self.read_lengths)

    # Read from memory
    w = self._address_memory(k, beta, g, s, gama, w_prev)
    r = self.memory.read(w)

    return r, w

 def read(self, w):
     """Read from memory (according to section 3.1)."""
     return torch.matmul(w.unsqueeze(1), self.memory).squeeze(1)

寫過程

寫過程同樣是獲得定位機制需要的k，beta, g, s, gama以及需要擦除的向量e和增加的向量a，然後調用_address_memory獲得定位向量w，然後根據e和a計算得出最後的寫入向量

def forward(self, embeddings, w_prev):
    """NTMWriteHead forward function.

    :param embeddings: input representation of the controller.
    :param w_prev: previous step state
    """
    o = self.fc_write(embeddings)
    k, beta, g, s, gama, e, a = _split_cols(o, self.write_lengths)

    # e should be in [0, 1]
    e = F.sigmoid(e)

    # Write to memory
    w = self._address_memory(k, beta, g, s, gama, w_prev)
    self.memory.write(w, e, a)

    return w

def write(self, w, e, a):
    """write to memory (according to section 3.2)."""
    self.prev_mem = self.memory
    self.memory = Variable(torch.Tensor(self.batch_size, self.N, self.M))
    erase = torch.matmul(w.unsqueeze(-1), e.unsqueeze(1))
    add = torch.matmul(w.unsqueeze(-1), a.unsqueeze(1))
    self.memory = self.prev_mem * (1 - erase) + add

Addressing Mechanism

定位機制的計算非常直觀，首先_similarity方法計算餘弦相似讀獲得wc，然後調用_interpolate與過去的w_prev進行插值操作，接着_shift偏移操作，這裏實際上調用的是_convolve循環卷積方法，最後進行_sharpen操作獲得最終的w

def address(self, k, beta, g, s, gama, w_prev):

    # Content focus
    wc = self._similarity(k, beta)

    # Location focus
    wg = self._interpolate(w_prev, wc, g)
    w1 = self._shift(wg, s)
    w = self._sharpen(w1, gama)

    return w

def _similarity(self, k, beta):
    k = k.view(self.batch_size, 1, -1)
    w = F.softmax(beta * F.cosine_similarity(self.memory + 1e-16, k + 1e-16, dim=-1), dim=1)
    return w

def _interpolate(self, w_prev, wc, g):
    return g * wc + (1 - g) * w_prev

def _shift(self, wg, s):
    result = Variable(torch.zeros(wg.size()))
    for b in range(self.batch_size):
        result[b] = _convolve(wg[b], s[b])
    return result

def _sharpen(self, w1, gamma):
    w = w1 ** gamma
    w = torch.div(w, torch.sum(w, dim=1).view(-1, 1) + 1e-16)
    return w

def _convolve(w, s):
    """Circular convolution implementation."""
    assert s.size(0) == 3
    t = torch.cat([w[-1:], w, w[:1]])
    c = F.conv1d(t.view(1, 1, -1), s.view(1, 1, -1)).view(-1)
    return c

訓練過程

首先輸入一系列的數據，每次輸入一個樣本，都先後進行讀和寫過程，然後在不給定輸入的情況下，獲得一系列輸出值，每次獲得一個輸出值時，同樣先後進行着讀和寫過程；只不過輸出的時候控制器接受的是0向量，而輸入數據的時候控制器接受的是樣本x值。我們可以根據輸出的值與樣本label的差距計算loss，對於copy任務來說，輸入樣本和label都是樣本本身，損失可以使用binary entropy loss，最後梯度下降法更新整合模型參數

def train_batch(net, criterion, optimizer, X, Y):
    """Trains a single batch."""
    optimizer.zero_grad()
    inp_seq_len = X.size(0)
    outp_seq_len, batch_size, _ = Y.size()

    # New sequence
    net.init_sequence(batch_size)

    # Feed the sequence + delimiter
    for i in range(inp_seq_len):
        net(X[i])

    # Read the output (no input given)
    y_out = Variable(torch.zeros(Y.size()))
    for i in range(outp_seq_len):
        y_out[i], _ = net()

    loss = criterion(y_out, Y)
    loss.backward()
    clip_grads(net)
    optimizer.step()

    y_out_binarized = y_out.clone().data
    y_out_binarized.apply_(lambda x: 0 if x < 0.5 else 1)

    # The cost is the number of error bits per sequence
    cost = torch.sum(torch.abs(y_out_binarized - Y.data))

    return loss.data[0], cost / batch_size

# 每次調用net(x)或者net()獲得輸出值的forward方法
def forward(self, x, prev_state):
    """NTM forward function.

    :param x: input vector (batch_size x num_inputs)
    :param prev_state: The previous state of the NTM
    """
    # Unpack the previous state
    prev_reads, prev_controller_state, prev_heads_states = prev_state

    # Use the controller to get an embeddings
    inp = torch.cat([x] + prev_reads, dim=1)
    controller_outp, controller_state = self.controller(inp, prev_controller_state)

    # Read/Write from the list of heads
    reads = []
    heads_states = []
    for head, prev_head_state in zip(self.heads, prev_heads_states):
        if head.is_read_head():
            r, head_state = head(controller_outp, prev_head_state)
            reads += [r]
        else:
            head_state = head(controller_outp, prev_head_state)
        heads_states += [head_state]

    # Generate Output
    inp2 = torch.cat([controller_outp] + reads, dim=1)
    o = F.sigmoid(self.fc(inp2))

    # Pack the current state
    state = (reads, controller_state, heads_states)

    return o, state

關於訓練結果，可以去github裏看，目前只有copy和deepcopy兩個任務，應該是分開訓練，但是按照前面分析的，神經圖靈機應該是可以先後訓練多個任務，並且保持新的任務不會覆蓋舊的任務，從理論上分析，如果讓記憶矩陣非常大，那麼就可以把每個任務儲存到記憶中不同的塊中，保持記憶矩陣的稀疏性，是可以做到任務間不互相干涉，因此讓模型達到能學習多個任務的能力。谷歌16年在Nature中提出的DNC其實也就是神經圖靈機，論文裏介紹了一些現在神經圖靈機可以完成的通用任務，想了解神經圖靈機具體應用的可以去看看。下面放出論文地址和代碼地址：

神經圖靈機(NTM)：https://arxiv.org/abs/1410.5401

DNC： https://www.nature.com/articles/nature20101

參考代碼：https://github.com/loudinthecloud/pytorch-ntm