Neural Tensor Network in PyTorch

Hi all, I am trying to implement Neural Tensor Network (NTN) layer proposed by Socher.

Basically, it aims to learn the relationship between two vectors. The sequence looks like below:

o = u’ f(x’ W y + V[x, y] + b)

where u, W, V, and b are the parameters. In my case, it looks like below:

  • u is in k (where k is a slice size)
  • f is tanh
  • x in b x n x p (b: batch size, n: sequence length, p: hidden dimension)
  • y in b x q (b: batch size, q: hidden dimension)
  • W in k x p x q (k: slice size)
  • V is multiplied by the stack of x and y, which is in k x (p + q)
  • b is a bias in k

I am expecting to see a list of scalars in batch size. (b x n)
I’ve tried to implement this layer using bmm, nn.Linear, nn.Bilinear with no success.
I’d be so happy if someone can help me doing this. Thanks.

1 Like

Here is what I would try

  • implement x’ W y using nn.Bilinear(p, q, k)
  • implement V[x, y] using nn.Linear(p+q, k) and torch.cat((x, y), -1) as [x, y]
  • both of the above provide a bias term by default so that takes care of b
  • implement u' * output_of_f using nn.Linear(k, 1, bias=False)

reshape w when needed

import torch
torch.matmul(torch.mm(x,w), y)