Neural networks implementing attention over object embeddings for object-centric visual reasoning
A video processing system configured to analyze a sequence of video frames to detect objects in the video frames and provide information relating to the detected objects in response to a query. The query may comprise, for example, a request for a prediction of a future event, or of the location of an object, or a request for a prediction of what would happen if an object were modified. The system uses a transformer neural network subsystem to process representations of objects in the video.
1 . A video processing system comprising one or more computers and one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to implement:
an object segmentation subsystem that is configured to receive an input comprising a sequence of video frames and to process each video frame of the sequence of video frames to generate an object segmentation output wherein, for each video frame, the object segmentation output comprises a representation of each object of a plurality of objects in the video frame;
a transformer neural network subsystem that is configured to receive an input sequence comprising i) an object input sequence including the object segmentation output for each video frame combined with an encoding of a position of the video frame in the sequence of video frames and ii) a query input sequence comprising a sequence of vectors representing a query, and to process the input sequence in parallel for all video frames in the sequence of video frames by applying a succession of attention neural network layers to generate an output sequence comprising a transformed input element for each element of the input sequence, wherein one or more of the attention neural network layers has an attention layer input for each element of the input sequence, the elements of the input sequence comprising a respective element for each representation of each of the plurality of objects of each of the video frames and for each of the sequence of vectors representing the query, and is configured to attend over representations of each of the plurality of objects in each of the video frames and each of the sequence of vectors representing the query by applying an attention mechanism over the attention layer inputs for the elements of the input sequence that include a respective element for each of the representations for each of the video frames and for each of the sequence of vectors representing the query to generate an attention layer output for each element of the input sequence; and
an output subsystem that is configured to receive one or more elements of the transformed input sequence and to process the one or more elements of the transformed input sequence by applying one or more output neural network layers to generate a video processing system output, wherein the video processing system output represents a response to the query on the objects in the sequence of video frames.
2 . The video processing system of claim 1 wherein the input sequence further comprises a summary token element, wherein the output sequence comprises a transformed summary token element representing a summary of the input sequence, and wherein the output subsystem is configured to receive the transformed summary token element and to process the transformed summary token element to generate the video processing system output.
3 . The video processing system of claim 1 wherein the object input sequence comprises a sequence of elements for each video frame, each comprising one of the representations combined with the encoding of the position of the video frame.
4 . The video processing system of claim 1 wherein the object segmentation subsystem is configured to determine, for each of the objects, a set of parameters parameterizing a distribution of a latent variable vector representing the object, and the object segmentation output comprises a mean value vector of the distribution for each object, where the mean value vector is the representation of the object.
5 . The video processing system of claim 1 further comprising a query embedding neural network subsystem that is configured to receive a sequence of tokens representing the query and to process the sequence of tokens to generate the sequence of vectors representing the query, wherein the vectors representing the query comprise embeddings of the tokens.
6 . The video processing system of claim 5 wherein the tokens comprise words or wordpieces in a natural language.
7 . The video processing system of claim 1 wherein elements of the input sequence each include a label to identify whether the element belongs to the object input sequence or to the query input sequence.
8 . The video processing system of claim 1 , wherein the input sequence further comprises a summary token element, wherein the output sequence comprises a transformed summary token element representing a summary of the input sequence, and wherein the output subsystem is configured to receive the transformed summary token element and to process the transformed summary token element to generate the video processing system output, and wherein the summary token element belongs to the query input sequence.
9 . The video processing system of claim 1 wherein one or more of the attention neural network layers is configured to apply a self-attention mechanism over the attention layer input, followed by a feed-forward neural network layer to generate the attention layer output.
10 . The video processing system of claim 1 wherein the attention mechanism is configured to apply each of a query transformation, a key transformation, and a value transformation to the attention layer input for each element of the input sequence to derive a respective query vector, key vector, and value vector; to apply each query vector to each key vector to determine respective weights for each value vector, and combining the value vectors using the respective weights to determine the attention layer output for each element of the input sequence.
11 . The video processing system of claim 1 , wherein the video processing system output defines a probability distribution over a set of possible responses to the query.
12 . The video processing system of claim 1 , wherein the video processing system further comprises a prediction neural network that is configured to receive and process, for each video frame, the transformed input elements for the representations of each object of the plurality of objects in the video frame, and to generate from the transformed input elements a predicted representation for the representation of each object in each video frame, and wherein the video processing system has been trained by performing operations comprising: suppressing, in the object segmentation output, the representation of one or more the objects for one or more of the video frames; determining an infill loss dependent upon a measure of difference between the suppressed representation and the predicted representation, and backpropagating gradients of the infill loss to adjust parameters of the transformer neural network subsystem.
13 . The video processing system of claim 1 , wherein the video processing system has been trained by performing operations comprising: determining an auxiliary loss dependent upon a measure of a distance between the response to the query and a correct response to the query, and backpropagating gradients of the auxiliary loss to adjust parameters of the transformer neural network subsystem.
14 . One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to implement:
an object segmentation subsystem that is configured to receive an input comprising a sequence of video frames and to process each video frame of the sequence of video frames to generate an object segmentation output wherein, for each video frame, the object segmentation output comprises a representation of each object of a plurality of objects in the video frame;
a transformer neural network subsystem that is configured to receive an input sequence comprising i) an object input sequence including the object segmentation output for each video frame combined with an encoding of a position of the video frame in the sequence of video frames and ii) a query input sequence comprising a sequence of vectors representing a query, and to process the input sequence in parallel for all video frames in the sequence of video frames by applying a succession of attention neural network layers to generate an output sequence comprising a transformed input element for each element of the input sequence, wherein one or more of the attention neural network layers has an attention layer input for each element of the input sequence, the elements of the input sequence comprising a respective element for each representation of each of the plurality of objects of each of the video frames and for each of the sequence of vectors representing the query, and is configured to attend over representations of each of the plurality of objects in each of the video frames and each of the sequence of vectors representing the query by applying an attention mechanism over the attention layer inputs for the elements of the input sequence that include a respective element for each of the representations for each of the video frames and for each of the sequence of vectors representing the query to generate an attention layer output for each element of the input sequence; and
an output subsystem that is configured to receive one or more elements of the transformed input sequence and to process the one or more elements of the transformed input sequence by applying one or more output neural network layers to generate a video processing system output, wherein the video processing system output represents a response to the query on the objects in the sequence of video frames.
15 . A method performed by one or more computers, the method comprising:
receiving an input comprising a sequence of video frames; and
processing the input comprising the sequence of video frames using a video processing system to generate a video processing system output, wherein the video processing system comprises:
an object segmentation subsystem that is configured to receive the input comprising the sequence of video frames and to process each video frame of the sequence of video frames to generate an object segmentation output wherein, for each video frame, the object segmentation output comprises a representation of each object of a plurality of objects in the video frame;
a transformer neural network subsystem that is configured to receive an input sequence comprising i) an object input sequence including the object segmentation output for each video frame combined with an encoding of a position of the video frame in the sequence of video frames and ii) a query input sequence comprising a sequence of vectors representing a query, and to process the input sequence in parallel for all video frames in the sequence of video frames by applying a succession of attention neural network layers to generate an output sequence comprising a transformed input element for each element of the input sequence, wherein one or more of the attention neural network layers has an attention layer input for each element of the input sequence, the elements of the input sequence comprising a respective element for each representation of each of the plurality of objects of each of the video frames and for each of the sequence of vectors representing the query, and is configured to attend over representations of each of the plurality of objects in each of the video frames and each of the sequence of vectors representing the query by applying an attention mechanism over the attention layer inputs for the elements of the input sequence that include a respective element for each of the representations for each of the video frames and for each of the sequence of vectors representing the query to generate an attention layer output for each element of the input sequence; and
an output subsystem that is configured to receive one or more elements of the transformed input sequence and to process the one or more elements of the transformed input sequence by applying one or more output neural network layers to generate a video processing system output, wherein the video processing system output represents a response to the query on the objects in the sequence of video frames.
16 . The method of claim 15 wherein the input sequence further comprises a summary token element, wherein the output sequence comprises a transformed summary token element representing a summary of the input sequence, and wherein the output subsystem is configured to receive the transformed summary token element and to process the transformed summary token element to generate the video processing system output.
17 . The method of claim 15 wherein the object input sequence comprises a sequence of elements for each video frame, each comprising one of the representations combined with the encoding of the position of the video frame.
18 . The method of claim 15 wherein the object segmentation subsystem is configured to determine, for each of the objects, a set of parameters parameterizing a distribution of a latent variable vector representing the object, and the object segmentation output comprises a mean value vector of the distribution for each object, where the mean value vector is the representation of the object.
19 . The video processing system of claim 1 wherein the query comprises a request for a prediction of a future event or state relating to one or more of the objects in the sequence of video frames.
20 . The method of claim 15 further comprising a query embedding neural network subsystem that is configured to receive a sequence of tokens representing the query and to process the sequence of tokens to generate the sequence of vectors representing the query, wherein the vectors representing the query comprise embeddings of the tokens.