Universal transformers
Methods, systems, and apparatus, including computer programs encoded on computer storage media, for implementing a sequence to sequence model that is recurrent in depth while employing self-attention to combine information from different parts of sequences.
1. A system implemented by one or more computers, the system comprising:
an encoder configured to receive an input sequence of elements each having a respective initial input representation and to revise the input representations by repeatedly applying a same series of encoding operations to all the elements of the sequence in parallel for each of multiple time steps of an encoding process, including revising the representations of the elements with each time step in the multiple time steps of the encoding process, for at most a predetermined maximum number of time steps; and
a decoder configured to decode a target sequence of symbols y=(y1, . . . , yn) autoregressively while at every time step of multiple time steps of a decoding process conditioning on previous symbols of the decoding process and on a final output of the encoder for the sequence.
2. The system of claim 1 , wherein each of the input representations corresponds to a respective input element.
3. The system of claim 2 , wherein the input elements comprise words and the input representations comprise respective word embeddings of the words, wherein word embeddings are vector representations of words.
4. The system of claim 1 , wherein:
the input sequence is of a length m and the representations are d-dimensional; and
the encoder is configured:
to initialize a matrix H 0 ∈ m×d with m rows, one for each item of the sequence, with the d elements of the representation of the item in the d columns of the matrix in the row; and
to compute representations H t at time step t, for time steps t from 1 through T, a depth of iteration, iteratively, by applying a multihead dot product self-attention mechanism followed by a recurrent transition function.
5. The system of claim 4 , wherein:
H t =LayerNorm( A t-1 +Transition( A t ))
where A t =LayerNorm( H t-1 +MultiHeadSelfAttention( H t-1 +P t )),
and the P t are two-dimensional (position, time) coordinate embeddings, obtained by computing vertically and horizontally the sinusoidal position embedding vectors
P pos,2i t =sin(pos/10000 2i/d )
P pos,2i+1 t =cos(pos/10000 2i/d )
for the position and the time step separately (pos) for each dimension (i), and summing these component-wise before applying self-attention; and
the multihead self-attention is computed with k heads as
MultiheadSelfAttention( H )=Concat(head 1 , . . . ,head k ) W O
where head 1 =Attention( HW i Q ,HW i K ,HW i V );
the projections are the following learned parameter matrices:
W Q ∈ d×d/k ,W K ∈ d×d/k ,W V ∈ d×d/k and W O ∈ d×d
and the attention is a scaled dot product attention according to
Attention
(
Q
,
K
,
V
)
=
softmax
(
QK
T
d
)
V
where d is the number of columns of Q, K and V and the dimension of the symbol representations.
6. The system of claim 5 , wherein the transition function is a fully-connected neural network that consists of a single rectified-linear activation function in-between two linear transformations, applied individually to each row of A t .
7. The system of claim 4 , wherein the transition function is a separable convolution.
8. The system of claim 4 , wherein:
after the T time steps of iteration, the encoder yields a matrix of vector representations H T for the input sequence as the final output of the encoder for the sequence.
9. The system of claim 8 , wherein the decoder and the encoder have the same recurrent structure in depth.
10. The system of claim 8 , wherein the decoder is configured:
to attend to the encoder representation of the input sequence using a multihead dot product attention function with queries Q obtained from projecting the decoder representations and keys K and values and V obtained from projecting the encoder representations.
11. The system of claim 8 , wherein the decoder is configured to obtain the per-symbol target distribution at position 1≤pos≤n by applying an affine transformation O from the final state to an output vocabulary size, followed by the softmax:
p ( y pos |y [1:pos-1] ,H T )=Softmax( OH T ).
12. The system of claim 8 , wherein the decoder is configured during training with the decoder self-attention distributions masked so that the decoder can only observe symbols to the left of any predicted symbol.
13. The system of claim 8 , wherein the system applies an Adaptive Computation Time (ACT) halting mechanism at each position and copies the state to the next time step after a block is halted, until all blocks are halted or the predetermined maximum number of time steps have been performed.
14. A method performed by one or more computers, the method comprising:
receiving, by an encoder, an input sequence of elements each having a respective initial input representation;
revising, by the encoder, the input representations by repeatedly applying a same series of encoding operations to all the elements of the sequence in parallel for each of multiple time steps of an encoding process, including revising the representations of the elements with each time step in the multiple time steps of the encoding process, for at most a predetermined maximum number of time steps; and
decoding, by a decoder, a target sequence of symbols y=(y1, . . . , yn) autoregressively while at every time step of multiple time steps of a decoding process conditioning on previous symbols of the decoding process and on a final output of the encoder for the sequence.
15. The method of claim 14 , wherein each of the input representations corresponds to a respective input element.
16. The method of claim 15 , wherein the input elements comprise words and the input representations comprise respective word embeddings of the words, wherein word embeddings are vector representations of words.
17. The method of claim 14 , wherein:
the input sequence is of a length m and the representations are d-dimensional; and further comprising:
initializing, by the encoder, a matrix H 0 ∈ m×d with m rows, one for each item of the sequence, with the d elements of the representation of the item in the d columns of the matrix in the row; and
computing, by the encoder, representations H t at time step t, for time steps t from 1 through T, a depth of iteration, iteratively, by applying a multihead dot product self-attention mechanism followed by a recurrent transition function.
18. The method of claim 17 , wherein:
H t =LayerNorm( A t-1 +Transition( A t ))
where A t =LayerNorm( H t-1 +MultiHeadSelfAttention( H t-1 +P t )),
and the P t are two-dimensional (position, time) coordinate embeddings, obtained by computing vertically and horizontally the sinusoidal position embedding vectors
P pos,2i t =sin(pos/10000 2i/d )
P pos,2i+1 t =cos(pos/10000 2i/d )
for the position and the time step separately (pos) for each dimension (i), and summing these component-wise before applying self-attention; and
the multihead self-attention is computed with k heads as
MultiheadSelfAttention( H )=Concat(head 1 , . . . ,head k ) W O
where head 1 =Attention( HW i Q ,HW i K ,HW i V );
the projections are the following learned parameter matrices:
W Q ∈ d×d/k ,W K ∈ d×d/k ,W V ∈ d×d/k and W O ∈ d×d
and the attention is a scaled dot product attention according to
Attention
(
Q
,
K
,
V
)
=
softmax
(
QK
T
d
)
V
,
where d is the number of columns of Q, K and V and the dimension of the symbol representations.
19. The method of claim 18 , wherein the transition function is a fully-connected neural network that consists of a single rectified-linear activation function in-between two linear transformations, applied individually to each row of A t .
20. The method of claim 17 , wherein the transition function is a separable convolution.
21. The method of claim 17 , further comprising:
after the T time steps of iteration, yielding, by the encoder, a matrix of vector representations H T for the input sequence as the final output of the encoder for the sequence.
22. The method of claim 21 , wherein the decoder and the encoder have the same recurrent structure in depth.
23. The method of claim 21 , further comprising:
attending, by the decoder, to the encoder representation of the input sequence using a multihead dot product attention function with queries Q obtained from projecting the decoder representations and keys K and values and V obtained from projecting the encoder representations.
24. The method of claim 21 , further comprising obtaining, by the decoder, the per-symbol target distribution at position 1≤pos≤n by applying an affine transformation O from the final state to an output vocabulary size, followed by the softmax:
p ( y pos |y [1:pos-1] ,H T )=Softmax( OH T ).
25. The method of claim 21 , wherein the decoder is configured during training with the decoder self-attention distributions masked so that the decoder can only observe symbols to the left of any predicted symbol.
26. The method of claim 21 , further comprising applying an Adaptive Computation Time (ACT) halting mechanism at each position and copying the state to the next time step after a block is halted, until all blocks are halted or the predetermined maximum number of time steps have been performed.
27. One or more non-transitory computer storage media encoded with computer program instructions that when executed by one or more computers cause the one or more computers to perform operations comprising:
receiving, by an encoder, an input sequence of elements each having a respective initial input representation;
revising, by the encoder, the input representations by repeatedly applying a same series of encoding operations to all the elements of the sequence in parallel for each of multiple time steps of an encoding process, including revising the representations of the elements with each time step in the multiple time steps of the encoding process, for at most a predetermined maximum number of time steps; and
decoding, by a decoder, a target sequence of symbols y=(y1, . . . , yn) autoregressively while at every time step of multiple time steps of a decoding process conditioning on previous symbols of the decoding process and on a final output of the encoder for the sequence.