Tracking query points in videos using point tracking neural networks

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for tracking query points in videos using generation of tokens and point tracking neural networks. In one aspect, a method includes obtaining a video sequence including a plurality of video frames; obtaining a set of one or more query points that each specify a spatial position in a corresponding video frame; generating a plurality of video tokens representing the video sequence; generating a query token for each query point for the corresponding video frame that characterizes the spatial position of the query point; and processing the video tokens and the query tokens for the query points using a neural network to generate a network output that includes a predicted spatial position of each query point in video frames other than the corresponding video frame for the query point.

Description

[0001] TRACKING QUERY POINTS IN VIDEOS USING POINT TRACKING NEURAL NETWORKS

[0002] BACKGROUND

[0003] This specification relates to processing inputs that include video frames using neural networks.

[0004] Neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to the next layer in the network, e.g., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current value inputs of a respective set of parameters.

[0005] SUMMARY

[0006] This specification describes a system implemented as computer programs on one or more computers in one or more locations that processes an input that includes (i) a video sequence that includes a plurality of video frames and (ii) a set of one or more query points.

[0007] Each query point is a point in a corresponding one of the video frames. More specifically, each query point can specify a respective spatial position for a respective pixel in a corresponding one of the plurality of video frames. The input data representing the set of one or more query points and the video sequence can also include intensify values of the pixels of the video frames in the video sequence.

[0008] The system processes the set of one or more query points and the video sequence using a point tracking neural network to generate a network output. Processing the set of one or more query points and the video sequence can include processing the intensify values of the pixels of the video frames in the video sequence. The neural network output includes, for each query point, a respective predicted spatial position of the query point in frames of the video sequence other than the video frame to which the query point corresponds. In other words, for a given query point in a video frame, the point tracking neural network can predict a spatial position of that query point in other video frames.

[0009] The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following advantages.

[0010] This specification describes a neural network that performs tracking of arbitrary query points across a video. This processes low level features (such as pixel attributes) within the frames of the video. This can provide information about how object surfaces of objects depicted in the video deform and move during the video. Such information can be useful for a variety of downstream tasks, e.g., video understanding or robotics tasks. Moreover, in some implementations, the neural network can predict when points become occluded, providing for more fine-grained occlusion tracking.

[0011] The techniques described in this specification can limit the computational complexity involved in tracking query points across a video, while still generating highly accurate tracking outputs, by using learned mask tokens for query points. For example, the system can generate uery tokens for uery points for video frames included in a sequence of video frames. The query tokens can characterize spatial positions of the query' points in corresponding video frames. For video frames other than the corresponding video frame for the uery point, the system can generate learned mask tokens for the query point. The system can process the learned mask tokens, e.g., by updating them using a token updating neural network, to generate a predicted spatial position for the query point in the video frame for the learned mask token. Performing the task of predicting spatial positions for query points as a masked decoding task can simplify the computational operations involved. Performing the task in this way can reduce a likelihood that inductive biases introduce error into predictions. Reducing inductive bias can improve the performance of the system by allowing it to leam and utilize a more diverse range of tracking strategies to predict spatial positions of querypoints.

[0012] The techniques described in this specification can allow a system to process videos of arbitrary length while still limiting computational complexify, computational resources, or both, involved in processing the videos. For example, in some implementations, the querytokens do not encode a temporal position of the corresponding video frame in the sequence. Omitting a temporal encoding from the query tokens can allow the system to generalize to processing videos of different lengths, e.g., because temporal encodings for query tokens for video frames included in videos could depend on the lengths of the videos. Therefore, the system need not account for lengths of videos in generating query- tokens, reducing the computational complexify of the operation of generating query tokens.

[0013] For example, the system can use a neural network that includes a recurrent block to perform temporal processing on video and query tokens. The temporal processing can include updating query tokens for query points based on other query- tokens corresponding to the same query points in adjacent video frames (e.g., only the video frames immediately before and / or after the video frame corresponding to the query token in the video sequence). Updating the query tokens in this way can facilitate making predictions for a given frame using only information from the given frame and the adjacent frame. In this way. the system is able to process videos of any length without increasing the number of computer resources used to process each video frame. The system can achieve this processing efficiency because predictions for a given frame need not incorporate information from larger numbers of frames for longer videos. Updating query tokens in this way can also increase the speed with which the system processes videos, e.g.. because the system can process each frame of the video more quickly by making predictions for the frame based only on information from adjacent frames. This increase in speed can enhance the performance of the system in online, real-time tracking applications.

[0014] The techniques described herein can increase the accuracy with which a system tracks query points in videos by facilitating separate spatial and temporal processing of the videos. For example, the system can process video tokens and query tokens using a neural network that includes a plurality' of layers that include a recurrent block and a self-attention block. The recurrent block can perform temporal processing on the tokens that includes updating query tokens, each for a query point in a given video frame, based on query tokens for the same query point in other video frames. The self-attention block can perform spatial processing on the tokens that includes updating query tokens, each for a query' point in a given video frame, based on query tokens for other query points in the same video frame and video tokens for the same video frame. This separation of the spatial and temporal processing of the tokens can decrease a likelihood that predictions of spatial positions for a given video frame depend on predictions from other video frames, in turn decreasing a likelihood that errors in such predictions propagate across multiple video frames.

[0015] The techniques described herein can improve the performance of a neural network in tracking query points in videos by training the neural network using a loss function that includes a respective auxiliary' loss term for each of a plurality of layers of the neural network. Including such terms in the loss function can improve the training of the neural network, for example, by resulting in respective updates to neural network parameters for each layer that are more likely to be effective for the layer (e.g., because the auxiliary loss term for the layer provides the layer with more direct feedback). Including such terms in the loss function can also increase the stability', speed, or both, of the training of the neural network. Such enhancements to the training of the neural network can improve the performance of the trained neural network. The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.

[0016] BRIEF DESCRIPTION OF THE DRAWINGS

[0017] Fig. l is a diagram of an example video processing system.

[0018] Fig. 2 shows an example of processing query points and video frames using neural networks to generate predicted positions for the query' points.

[0019] Fig. 3 is a flow diagram of an example process for predicting spatial positions of query points in video frames.

[0020] Fig. 4 is a flow diagram of an example process for training a neural network that is configured to generate predicted spatial positions for query points in video frames.

[0021] Like reference numbers and designations in the various drawings indicate like elements.

[0022] DETAILED DESCRIPTION

[0023] Fig. 1 is a diagram of an example video processing system 100. The video processing system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented. The system 100 is configured to process an input including query' points 102 and a video sequence 104 to generate a network output 118 that includes predicted spatial positions 120 and, optionally, visibility scores 122 for the query points. The network output 118 can be used to track the query’ points 102 in the video sequence 104.

[0024] The system 100 receives an input that includes (i) a video sequence 104 that includes a plurality' of video frames and (ii) a set of one or more query points 102.

[0025] Each query point 102 is a point in a corresponding one of the video frames, e.g., that specifies a respective spatial position, e.g., a respective pixel, in a corresponding one of the plurality of video frames. Thus, each query point 102 can be represented as a tuple of coordinates (x, y, t), where x, y are the spatial coordinates of the query' point 102 and t is the index of the corresponding video frame in the video sequence 104. Each video frame in the video sequence 104 can be represented as a multi-dimensional array or tensor of pixel intensity values, such as RGB color values for each pixel. In some examples, the video frames in the video sequence 104 can depict a real -world environment at different respective times, e.g., a real -world environment that a robot or other agent interacts with to perform a task.

[0026] The system 100 processes the set of one or more query points 102 and the video sequence 104, e.g., the intensity values of the pixels of the video frames in the video sequence 104, using a token generation engine 106 to generate a set of video tokens 108 and a set of query tokens 110. The token generation engine 106 can include one or more neural networks. In some implementations, at least one of the one or more neural networks can be an embedding neural network that is configured to process input data (such as, e.g., a patch of a video frame) to generate an embedding characterizing the input data, where the embedding can be a vector or other ordered collection of numerical values, e.g., a vector of floating point or other numerical values. For example, the embedding can be a multi-dimensional vector that is located in an embedding space. The one or more neural networks included in the engine 106 can be trained using any of a variety of training techniques. The one or more neural networks included in the engine 106 can be jointly trained with other neural networks included in the system 100, e.g.. the neural network included in the point tracking engine 112 described below. For example, training the one or more neural networks can include determining gradients of an objective function and updating parameters of the neural networks using the determined gradients. The embedding neural network can have any appropriate architecture. For example, an embedding neural network can include a linear projection layer that transforms input data into a feature vector, e.g., an embedding.

[0027] The token generation engine 106 can generate the set of video tokens 108 by first dividing each video frame in the video sequence 104 into multiple non-overlapping image patches. The token generation engine 106 can then use an embedding neural network to apply a linear projection to each image patch of each video frame, transforming the pixel data within the patch into an embedding. Each video token of the set of video tokens 108 can be an embedding resulting from application of the linear projection to a respective one of the image patches. In some implementations, each video token can also include a positional encoding, which captures (characterizes) the spatial location of the respective patch within its frame.

[0028] The token generation engine 106 can generate the set of query' tokens 110 by processing each query point from the set of query points 102 using a neural network included in the engine 106. For each query point, the engine can process the query point to generate a respective query token for each of the video frames in the video sequence 104. For each query point, the token generation engine 106 can generate the query' token for the corresponding video frame of the query point (e.g., the frame specified by the 't' coordinate of the query point) to characterize the spatial position of the query point within the corresponding video frame. For example, the token generation engine 106 can generate the query token for the corresponding video frame by computing a spatial positional encoding of the (x, y) coordinates of the query point in the corresponding video frame. The spatial positional encoding can be an embedding that is located in the same embedding space in which the video tokens 108 are located. In some implementations, the query token for the corresponding video frame does not encode the temporal position of the corresponding video frame within the video sequence 104 (e.g., does not encode the value of ‘f of the tuple of coordinates of the query token). In some other implementations, the query tokens can each encode the temporal position of the corresponding video frame in the video sequence 104.

[0029] For each query point, the token generation engine 106 can generate the query7token for every other video frame in the video sequence other than the corresponding frame as a learned mask token. For each of the other video frames, the learned mask token can, for example, be an embedding, e.g.. in the form of a vector (a vector representation), that indicates that the spatial position of the query point in the video frame is unknown and is to be predicted by the system. The embedding or vector representation of the learned mask token can, for example, be learned by the neural network that is used to generate the query tokens, i.e., during training of the neural network. That is. the neural network can be trained to embed the learned mask token representing the unknown spatial location of a query point in the embedding space. The set of query tokens 110 can include, for each query point, the query token for the corresponding video frame and the respective query tokens for each of the other video frames other than the corresponding video frame.

[0030] The point tracking engine 112 is configured to process the video tokens 108 and the query7tokens 110 to predict the positions of each of the query points 102 in each of the video frames of the video sequence 104 other than the corresponding frame for the query point. The point tracking engine 112 can include a neural network. The neural network included in the point tracking engine 112 can include a token updating neural network 114 and a position prediction head 116.

[0031] The token updating neural network 114 can include one or more neural network layers. For example, the token updating neural network 114 can include any appropriate types of neural network layers (e.g., fully connected layers, convolutional layers, attention layers, and / or recurrent layers) in any appropriate number (e.g., 5 layers, or 10 layers, or 50 layers) and connected in any appropriate configuration (e.g., as a directed graph of layers). One or more of the layers included in the token updating neural network 114 can include a selfattention block, such as a self-attention block. One or more of the layers included in the token updating neural network 114 can include a recurrent block. The token updating neural network can be trained using any of a variety of training techniques, such as the techniques described below with reference to Fig. 4.

[0032] The token updating neural network 114 can process the video tokens 108 and the query tokens 110 to update the query tokens 110. For example, the token updating neural network 114 can update the learned mask tokens included in the set of query tokens 110 to infuse them with (e.g., introduce) contextual information from both the video content and the known query point locations. For example, for each query’ point, for each video frame in the video sequence 104 other than the corresponding video frame, the token updating neural network 114 can use a self-attention block to perform attention operations over the video tokens of the frame and the query’ tokens for the frame (e.g., including the learned mask token for the query point for the frame and the query’ tokens for the other query points for the frame) to update the learned mask token for the query point for the frame. For example, updating the learned mask tokens in this way can allow the token updating neural network 114 to infuse the learned mask tokens with spatial information about the video frame for the learned mask token.

[0033] For example, the token updating neural network 114 can update the query’ tokens 110 by using a recurrent block to scan over the video frames for each of a set of temporal “tubes”. There can be a temporal tube for each query’ point, and a temporal tube for each image patch. The temporal tube for a given query point can include query' tokens corresponding to the query point across time. For example, the temporal tube for a given query’ point can be a set of query tokens that includes the query token for the query point for each of the video frames in the video sequence 104. The temporal tube for a given image patch can include video tokens corresponding to the same spatial region of the video frames (e.g., the spatial region included in the image patch) across time. For example, the temporal tube for a given image patch can be a set of video tokens that includes the video token for the image patch for each of the video frames of the video sequence 104. The set of temporal tubes processed by the recurrent block can include the temporal tubes corresponding to each of the query’ points and to each of the image patches.

[0034] The recurrent block can process each temporal tube to update the tokens included in the temporal tube. For example, updating the tokens in this way can allow the token updating neural network 114 to infuse the tokens with information related to motion of the query' points, changes in video frames over time, or both. For example, using a recurrent block to update the tokens in this way can enable the processing of videos of arbitrary length with consistent per-frame computational complexity.

[0035] In some implementations, the token updating neural network 114 can include multiple self-attention blocks, multiple recurrent blocks, or both. In some implementations, the selfattention blocks and recunent blocks included in the token updating neural network 114 can be interleaved with one another. In such implementations, the video tokens 108 and the query tokens 110 can be processed by each of the self-attention blocks and recurrent blocks as described above. In this way, the system 100 can perform multiple rounds of updating the query tokens, e.g., which can improve the quality of the generated updated query tokens.

[0036] The token updating neural network 114 can generate a set of updated query tokens. The set of updated query' tokens can include, for each learned mask token of the set of query tokens 110, an updated query' token that encodes information characterizing a predicted location of the corresponding query point, a predicted visibility of the corresponding query point, or both.

[0037] The position prediction head 116 can be a neural network component configured to process the updated query' tokens to generate a respective predicted spatial position for each of the query points 102. The position prediction head 116 can be implemented, for example, as a multi-layer perceptron (MLP) network. For each query point and for each video frame other than its corresponding frame, the position prediction head 1 16 processes the updated query' token to generate a predicted spatial position of the query' point in the video frame. For example, the position prediction head 116 can generate the predicted spatial positions using the techniques described below with reference to Fig. 3.

[0038] In some implementations, the point tracking engine can also include a visibility prediction head 117. The visibility prediction head 117 can be implemented, for example, as a MLP network. For each query point and for each video frame (“other video frame’") other than its corresponding frame, the visibility prediction head 117 processes the updated query token to generate a visibility score for the query point in the other video frame. The visibility’ score can indicate whether the query' point is visible in the other video frame For example, the visibility prediction head 117 can generate the visibility’ scores using the techniques described below with reference to Fig. 3.

[0039] The network output 118 includes the predicted spatial positions 120 and, in some implementations, the visibility scores 122. The predicted spatial positions 120 can include, for each query point and for each video frame other than the corresponding video frame for the query point, a predicted spatial position of the query point in the video frame, e.g., as generated by the position prediction head 116. Each predicted spatial position can be a set of (x, y) coordinates for the video frame that indicates a prediction that the query point is located at the position in the video frame indicated by the set of coordinates.

[0040] The visibility scores 122 can include, for each query point and for each video frame other than the corresponding video frame for the query point, a visibility score for the query point in the video frame, e g., as generated by the visibility prediction head 1 17. Each visibility score can indicate whether the query point is visible in the video frame. For example, the visibility score can be a number, e g., between 0 and 1, representing the likelihood that the query point is visible in the video frame.

[0041] The predicted spatial positions 120 (and optionally the visibility scores 122) can be used for any of a variety of purposes.

[0042] For example, the predicted spatial positions 120 can be used to generate a reward signal for training a robot or other agent through reinforcement learning, e.g., if the task being performed by the agent requires moving a point in the scene (e.g., a point associated with an object in the scene) from one location to another, the distance between the predicted position of the point in the last video frame in the sequence and the target location can be used to generate a reward.

[0043] As another example, the predicted spatial positions 120 (and optionally the visibility scores 122) can be provided as an additional input to a policy neural network for controlling a robot or other agent interacting with an environment, e.g., a real -world environment. In this example, the query points can be points of interest in the last frame of the video sequence, and the predictions for the earlier frames in the sequence can be provided as input to the policy neural network, e.g., to provide a signal as to the recent motion of the agent or of other objects in the environment.

[0044] As another example, the predicted spatial positions 120 (and optionally the visibility scores 122) can be provided, along with the video sequence 102, as input to a video understanding neural network, e.g., an action classification neural network or a topic classification neural network, to provide additional information to the video understanding neural network regarding motion in the scene.

[0045] As another example, the predicted spatial positions 120 (and optionally the visibility scores 122) can be used for imitation learning, i.e., to enable the imitation of motion rather than appearance. As another example, the predicted spatial positions 120 (and optionally the visibility scores 122) can be used for training a video generation neural network. For example, the video generation neural network can process the predicted spatial positions 120 (and optionally the visibility scores 122) to generate a video that depicts motion of an object that is reflected by the predicted spatial positions 120 (and optionally the visibility scores 122). For example, training the video generation neural network can include conditioning the video generation neural network on the predicted spatial positions 120 (and optionally the visibility scores 122). As another example, the predicted spatial positions 120 (and optionally the visibility scores 122) can be used for training a point trajectory generation neural network to process an input image to predict future (or past) spatial positions and / or visibility scores for one or more query points in the input image. The predicted spatial positions and / or visibility scores can then, for example, be used to animate the input image, e.g., by processing the predicted spatial positions and / or visibility' scores using a trained video generation neural network. Videos generated in this way can be used by devices, such as a robot or other agent, to plan movements needed to perform a task.

[0046] Fig. 2 shows an example 200 of processing query points and video frames using neural networks to generate predicted positions for the query points. The process 200 can be implemented by a system, e.g., the video processing system 100 of Fig. 1, implemented as computer programs on one or more computers in one or more locations. For example, the system can include a token generation engine, e.g., the token generation engine 106 of Fig. 1, and a point tracking engine, e.g., the point tracking engine 1 12 of Fig. 1. The process 200 facilitates the transformation of raw video and coordinate data into a structured tensor format suitable for processing by a neural network architecture that can include both spatial and temporal processing modules. By processing the video and coordinate data in this way, the system can generate predicted spatial positions for query points associated with the video and coordinate data.

[0047] The process 200 includes obtaining a set of query' points 202 and an input video 206. Each of the query’ points 202 can be a point in a corresponding one of the video frames of the input video 206. Each query point 202 can be represented as a tuple, for example (x, y. t specifying spatial coordinates of the point in the corresponding frame (e g., x and ) and a temporal index (e.g., t) that indicates the index of the corresponding video frame within the input video 206. The set of query' points 202 can include Q query points, where Q is an integer greater than or equal to 1. The input video 206 can be a sequence of video frames, where each frame is a two- dimensional grid of pixels, and each pixel can have associated color values (e.g., RGB values). The input video 206 can include T frames, where T is an integer greater than or equal to 1.

[0048] The system processes the input video 206 to generate video tokens. The system can partition each frame of the input video 206 into a grid of non-overlapping image patches. For example, an input frame of size H x W pixels can be divided into h x w image patches, where each patch has a size of (H / h) x (W / w) pixels. The system can process each of the image patches using a linear projection layer 208. The linear projection layer 208 can be a neural network layer, such as a fully connected layer with a weight matrix and bias vector, configured to process pixel data of an image patch (e.g., a flattened vector of RGB values for the image patch) to generate an embedding for the image patch. The embedding can have a predetermined dimensionality, C. The embedding can represent feature data characterizing the image patch. For example, by applying the linear projection layer 208 to each image patch, the system can map the pixel-space representation of each patch into a higherdimensional feature space, e.g., the embedding space in which the embeddings are located. The embeddings generated by the linear projection layer 208 for all patches across all frames form a set of video tokens 212b.

[0049] Concurrently or sequentially, the system processes the query points 202 to generate a set of query tokens 212a. The set of query tokens 212a can include, for each of the query points 202, a respective query token for each video frame of the input video 206. For example, the set of query' tokens 212a can include, for each of the query' points 202, a query' token for the video frame corresponding to the query' point and query tokens for each of the video frames other than the corresponding video frame. For example, the video frame corresponding to the query point can be the video frame with the index t in the input video 206, where the query' point is represented by the tuple (x, v, I).

[0050] For each query point 202, the system can generate the query token for the corresponding video frame using a positional encoding operation 204. The positional encoding 204 can be a function that processes the query point to generate an embedding for the query’ point. For example, the positional encoding operation 204 can convert the continuous (x, y) spatial coordinates of the query' point into an embedding. The embedding can be a high-dimensional vector representation that is located in an embedding space with dimensionality C, e.g.. the same dimensionality as the embedding space in which the video tokens 212b are located. The positional encoding operation 204 can use various techniques. such as sinusoidal functions of different frequencies, to transform the spatial coordinates of the query point into the embedding representation.

[0051] For each query point 202, the system can generate a respective query token for each of the video frames other than the corresponding video frame such that the respective query token is a learned mask token for the query point for the video frame. The learned mask token can be an embedding, e.g.. also of dimensionality C, that indicates that the spatial position of the query point in the video frame is not provided as input, e.g.. to a neural network that is to process the query tokens, as described below. In this way, for each of the Q query points 202, the system can generate a sequence of T query' tokens 212a, one query' token for each frame in the input video 206.

[0052] In some implementations, the system can generate the video tokens 212b by adding to the embedding for each image patch (e.g., as generated by the linear projection layer 208) a positional encoding. The positional encoding can be generated using the positional encoding operation 204. The positional encoding that is added to each image patch can be a vector that represents information about the spatial location of the image patch within its video frame. For example, the system can perform a concatenation operation 210, which includes combining the embedding for each image patch from the linear projection layer 208 with a corresponding positional encoding for the image patch from the positional encoding operation 204.

[0053] The system can then generate an input tensor 212 by combining the video tokens 212b and the query' tokens 212a. The input tensor can have dimensions of [T, h x w + Q, C], where T is the number of frames in the video sequence, h x w is the number of image patches per frame, Q is the number of query points, and C is the dimensionality of the token embeddings. For example, the input tensor 212 can include, for each time index t of a video frame of the input video 206, a spatial vector of C-dimensional tokens that includes the video tokens for all h x w image patches of the video frame and the query tokens for all Q query' points for the video frame. The vector 212c of Fig. 2 is an example of such a spatial vector. As shown in the example of the vector 212c, the query tokens included in the spatial vector can include one or more query tokens that characterize spatial positions of query points in the video frame (e.g., represented by' the query token 212f of the vector 212c that includes a dot indicating the spatial position of the query' point); one or more query' tokens that are learned mask tokens (e.g., represented by the query tokens of the vector 212c that do not include dots); or both. For example, the input tensor 212 can include a number of tubes equal to h x w + Q. Each temporal tube can include a number of tokens equal to the number of video frames in the input video 206 (e.g., T). The tokens included in a temporal tube can be either: i) the set of query tokens 212a for a given query point that includes the query token for the query point for each of the video frames (e.g., the temporal tube 212d); or ii) the set of video tokens 212b for a given image patch that includes the video token for the image patch for each of the video frames (e.g.. the temporal tube 212e). Generating the input tensor 212 in this way can allow subsequent neural network layers to jointly process visual information from the video frames and the positional information from the uery points, as described below.

[0054] The system can process the input tokens tensor 212 using two or more neural network blocks, which can, for example, each include a self-attention block 216 and a recurrent block 214. These blocks can be arranged in an alternating or stacked fashion to iteratively refine the token representations. The neural network blocks can be included in a neural network of the system, e.g., the token updating neural network 114 of Fig. 1, that is configured to process an input tensor of video and query tokens to generate updated query tokens.

[0055] A self-attention block 216 can be configured to perform spatial processing of the input tensor 212 in order to update the query tokens 212a. The self-attention block 216 can operate on each spatial vector of the input tokens tensor 212 independently. For example, the self-attention block 216 can use the time dimension of the input tensor 212 as a batch dimension and apply a self-attention mechanism across the h x w video tokens and the Q uery tokens of each spatial vector. For example, the self-attention mechanism can cause each token to be updated based on a weighted sum of all other tokens in the same frame. For example, the self-attention block 216 can apply the self-attention mechanism using techniques similar to those described in Vaswani, et. al., ‘"Attention is All You Need'; arXiv: 1706.03762 (https: / / arxiv.org / abs / 1706.03762). For example, the self-attention block 216 can apply the self-attention mechanism using techniques similar to those described in Dosovitskiy, et. al., “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale”.

[0056] In this way, the updated tokens can capture spatial relationships and context. For example, a uery token can attend to relevant video tokens that depict an object or feature associated with the query point for the query token, thereby infusing the uery token with relevant visual information from the frame. For example, each of the learned mask tokens can be updated with spatial context from the video frame. A recurrent block 214 can be configured to perform temporal processing of the input tensor 212 in order to update the query’ tokens 212a. The recurrent block 214 can process each temporal tube of the input tensor 212 independently. For example, the recurrent block 214 can use the spatial dimension of the input tensor 212 as a batch dimension and scan over the time dimension T. The recurrent block 214 can be a recurrent neural network or a similar architecture that maintains a hidden state. The hidden state can be updated at each time step and can carry information from past frames to future frames. For example, the recurrent block 214 can perform the temporal processing of the temporal tubes using techniques similar to those described in Kuznetsov, et. al., “Discriminative State Space Models”, 31st Conference on Neural Information Processing Systems (NIPS 2017), Long Beach, CA, USA (https: / / proceedings.neurips.cc / paper_files / paper / 2017 / file / 6fel31632103526e3a6e8114c78e blel-Paper.pdl)- For example, the recurrent block 214 can perform the temporal processing of the temporal tubes using techniques similar to those described in Gu, et. al., “Efficiently Modeling Long Sequences with Structured State Spaces” arXiv:2111.00396 (https: / / arxiv.org / pdf / 2111.00396).

[0057] In this way, the recurrent block 214 can model motion, temporal dynamics, and long- range dependencies in the video without a significant increase in computational complexity for longer videos. By processing temporal tubes, the recurrent block 214 can update each query token by propagating information along the trajectory of the query point over time. For example, the recurrent block 214 can update the learned mask tokens for the query point in this way. For example, propagation of information along the trajectory of a query' point associated with a learned mask token over time can allow the known position of the query’ point at time t to be incorporated into the learned mask token for the query point for the video frame with index t+1. In this way. the known position can influence a prediction of the position of the query point at time t+1. Processing the tokens using the recurrent block 214 in this way can be beneficial for online or real-time tracking applications, as predictions for a given frame can be made using only information from that frame and preceding frames.

[0058] By processing the input tokens tensor 212 using the two or more neural network blocks, the system can generate an updated set of query tokens. The updated set of query tokens can include, for each query token of the set of query tokens 202, an updated query’ token that encodes both spatial and temporal context. As described with reference to Fig. 1, the updated set of query tokens can be processed by a position prediction head (e.g., the position prediction head 116 of Fig. 1) and optionally a visibility prediction head (e.g.. the visibility prediction head 117 of Fig. 1) to generate predicted spatial positions and visibility scores for each query’ point in each video frame.

[0059] FIG. 3 is a flow diagram of an example process 300 for predicting spatial positions of query’ points in video frames. For convenience, the process 300 will be described as being performed by a system of one or more computers located in one or more locations . For example, a video processing system, e.g., the video processing system 100 of FIG. 1. appropriately programmed in accordance with this specification, can perform the process 300. The system can include a neural network that is configured to process video tokens representing video frames and query’ tokens representing query points in the video frames, to generate predicted spatial positions for the query' points in the video frames. For example, the neural network included in the system can include one or more blocks, such as a selfattention block and a recurrent block. For example, the neural network can include a token updating neural network similar to the token updating neural network 114 of Fig. 1. For example, the neural network can include a position prediction head similar to the position prediction head 116 of Fig. 1, which can be a multi-layer perceptron (MLP). For example, the neural network can include a visibility prediction head similar to the visibility prediction head 1 17 of Fig. 1, which can be a MLP.

[0060] The system obtains a video sequence including multiple video frames (302). The video sequence can include a temporally ordered sequence of video frames. Each video frame can be represented as a multi-dimensional array or tensor of pixel data, such as a three- dimensional tensor of dimensions H x W x 3 for an HxW pixel frame with RGB color channels.

[0061] The system obtains a set of one or more query points, where each query point specifies a respective spatial position in a corresponding one of the video frames (304). The set of query points can be received from a user input device, a data storage system, or another automated process. Each query' point can be represented as a data structure, such as a tuple (x, v, / ). where x and v are floating-point or integer values representing the spatial coordinates within a specific video frame, and t is an integer identifying the index of the corresponding frame within the obtained video sequence. For example, each query point in the set of query points can correspond to a respective object or feature depicted in the video sequence. Each respective object or feature can be an object or feature to be tracked over the frames of the video sequence, e.g., for which a spatial position in each video frame other than the corresponding video frame is to be predicted. The system generates a plurality of video tokens representing the video sequence (306). The system can generate the plurality of video tokens by first partitioning each video frame into a grid of non-overlapping image patches. The system can, for example, partition each video frame into the same grid, e.g., so that each image patch of a given video frame has a corresponding image patch in each of the other video frames of the video sequence, where the corresponding image patches define the same area of each of the other video frames as the area of the given video frame defined by the image patch. Each patch can, for example, be represented as a vector of pixel data.

[0062] For each patch, the system can, for example, apply a linear projection operation to the vector representing the pixel data for the patch to generate an embedding for the patch that is located in a multi-dimensional embedding space. The system can apply the linear projection operation using a linear projection layer, such as the linear projection layer 208 of Fig. 2. Each video token of the plurality of video tokens can comprise the embedding for a respective image patch that is generated via the linear projection operation. The plurality of video tokens can include a respective video token for each image patch in each of the video frames.

[0063] In some implementations, the system can generate the plurality of video tokens by, for each image patch, after generating the embedding for the image patch, adding a spatial positional encoding to the embedding. The spatial positional encoding can be a second embedding that represents the spatial location of the image patch within its frame. The video token for the image patch can then be the combination of the embedding for the image patch generated via the linear projection operation and the spatial positional encoding for the image patch.

[0064] For each query point of the set of query points, the system generates a respective query token for each of the plurality of video frames (308). For example, for each query point, the system can generate both: i) a uery token for the corresponding video frame for the query point; and ii) a respective query token for each of the video frames other than the corresponding video frame for the query point.

[0065] For each query point, the query token for the corresponding video frame can characterize the respective spatial position of the query point in the corresponding video frame. In some implementations, for each query point, the respective query token for the corresponding video frame for the query point includes a spatial positional encoding of the respective spatial position of the query point in the corresponding video frame. For example, the system can perform a positional encoding operation, e.g., the positional encoding operation 204 of Fig. 2, to convert coordinates of the query point in the corresponding video frame into an embedding that represents the spatial position of the query point in the corresponding video frame in a multi-dimensional embedding space, e.g., the same embedding space in which the embeddings for the video tokens are located.

[0066] In some implementations, for each query point, the respective query token for the corresponding video frame does not encode a temporal position of the corresponding video frame within the video sequence. For example, skipping encoding a temporal position in the uery tokens in this way can help to facilitate separate processing of the video and query tokens in each of a spatial and temporal dimension, as described below. Without encoding temporal positions, query' tokens can be simplified and therefore the system can use fewer computational resources to generate and process the query tokens. Omitting a temporal encoding from the query tokens can also allow the system to more easily generalize to processing video sequences of different lengths, e.g., because temporal encodings for query tokens could depend on the lengths of the corresponding video sequences.

[0067] For each query point, the respective query token for each video frame other than the corresponding video frame can, for example, be a learned mask token. The learned mask token can be an embedding that indicates that the spatial position of the query point in the other video frame is not provided as input to the neural network included in the system. In this way, the learned mask token for each query point for each of the other video frames can represent a “placeholder’7that indicates that the position of the query point in the frame is unknown and is to be predicted by the system. For example, generating query tokens in this way can allow the system to predict spatial positions of query' points using a masked decoding approach, e.g., an approach in which the system is tasked with "filling in" the information for the learned mask tokens for each of the query points for each of the other video frames.

[0068] The system processes the video tokens and the respective query' tokens for the query points to generate, for each query' point, a respective predicted spatial position for the query' point in one or more other video frames other than the corresponding video frame (310). The system can use the neural network to process the video tokens and the respective querytokens to generate a network output that includes the respective predicted spatial positions for the query' points in the one or more other video frames.

[0069] In some implementations, the system first processes the video tokens and the respective query tokens using the token updating neural network included in the neural network to update the respective query' tokens for the query point for the one or more other video frames. In such implementations, the token updating neural network can generate an input tensor that includes both the video tokens and the query tokens. For example, the input tensor can be similar to the input tensor 212 of Fig. 2. For example, the input tensor can include, for each video frame of the plurality of video frames, a spatial vector of tokens. The spatial vector of tokens for a video frame can include video tokens for each of the image patches of the video frame and query’ tokens for each of the query points for the video frame. In this way, each spatial vector can include a number of tokens equal to the sum of the total number of image patches and the total number of query points. For example, each spatial vector can represent a “slice’’ of the input tensor at a given point along a time dimension of the input tensor, e.g.. as illustrated by the spatial vector 212c of Fig. 2.

[0070] The input tensor can include, for each query point and for each set of corresponding image patches (e.g., that define the same area in each of the video frames), a temporal tube that is a vector of tokens. For each query' point, the temporal tube can include the respective query token for the query' point for each of the video frames. For each set of corresponding image patches, the temporal tube can include the respective video token for each of the corresponding image patches in the set. In this way. each temporal tube can include a number of tokens equal to the number of video frames in the video sequence. For example, each temporal tube can represent a “slice” of the input tensor at a given point along a spatial dimension of the input tensor, e.g., as illustrated by the temporal tubes 212d and 212e of Fig. 2.

[0071] For example, the input tensor can have dimensions of T, h X w + Q and C, where T represents a number of video frames in the video sequence; h represents a number of rows of non-overlapping image patches; w represents a number of columns of non-overlapping image patches; Q represents a number of query points in the set of one or more query points; and C represents a dimension of a space into which the non-overlapping image patches are embedded, as described above with reference to the example input tensor 212 of Fig. 2.

[0072] The system processes the input tensor using the token updating neural network. In some implementations, the token updating neural network can include a plurality of layers that each update the video tokens and the query tokens (e.g., included in the input tensor). The token updating neural network can, for example, include an alternating sequence of: i) self-attention blocks, e.g., for spatial processing of the input tensor; and ii) recurrent, e.g., for temporal processing of the input tensor. Each self-attention block can perform spatial processing on the video tokens and the respective query tokens. The spatial processing can include, for each query point, updating the respective query tokens for the query point for each of the one or more other video frames based on video tokens for the video frame and respective query tokens for the video frame for query points in the set of one or more query points other than the uery point.

[0073] For example, each self-attention block can apply a self-attention operation across each of the spatial vectors of the input tensor. For example, the self-attention operation across a spatial vector can include updating query tokens included in the spatial vector using information encoded in other tokens included in the spatial vector. For example, each selfattention block can process each spatial vector using techniques described above with reference to Fig. 2. By applying the self-attention operation in this way, the system can update, for each query point and each query token for the query point for the one or more other video frames, the query token (which is a learned mask token) included in the spatial vector with spatial context from the video frame (e.g., context that is encoded in the video tokens for the video frame and in the query7tokens for other query7points for the video frame).

[0074] Each recurrent block can perform temporal processing on the video tokens and the respective query tokens. For example, the temporal processing can include, for each query point, updating each of the respective query tokens for the query point for the one or more other video frames based on a respective query7token for the query7point in an adjacent video frame. For example, each recurrent block can maintain a hidden state, which can be a vector or a tensor that represents information about

[0075] Each recurrent block can scan across each of the temporal tubes to update query tokens included in the temporal tube. For each temporal tube, the recurrent block can scan across the temporal tube by processing each token included in the temporal tube in sequential order (e.g., such that the first token processed corresponds to the first video frame in the temporally ordered sequence of video frames, the second token processed corresponds to the second video frame in the temporally ordered sequence of video frames, and so on). As the recurrent block scans across the temporal tube, it can maintain a respective hidden state. The hidden state can be a set of one or more variables that is dynamically updated that represents information from tokens that were previously processed by the recurrent block. In response to processing each token, the recurrent block can update the hidden state, e.g., to incorporate information from the most recently processed token into the hidden state. In this way, the hidden state can maintain information from all tokens that w ere previously processed by the recurrent block. For each token in the sequence of tokens included in the temporal tube, the recurrent block can process the token by processing both the token and the hidden state. As a result of processing both the token and the hidden state, the recurrent block can update the hidden state based on the token. As a result of processing both the token and the hidden state, the recurrent block can update the token based on the hidden state. Updating the token based on the hidden state can allow the recurrent block to infuse the token with information from all of the previously processed query tokens that is maintained in the hidden state. For example, the recurrent block can update each token corresponding to a respective video frame that is included in the temporal tube with information maintained in the hidden state, which includes information encoded in tokens included in the temporal tube that correspond to earlier video frames (e.g., earlier than the respective video frame in the temporally ordered sequence of video frames). For example, each recurrent block can process each temporal tube using techniques described above with reference to Fig. 2.

[0076] In this way, the recurrent block can propagate information across all frames despite only processing, for each token, the preceding token from the preceding frame. By scanning the temporal tubes of the input tensor in this way. the system can update query tokens for query points so as to incorporate information about known position of query points at earlier times. For example, scanning the temporal tubes in this way can enable the system to capture motion and temporal dynamics over the sequence of video frames.

[0077] After updating the respective query tokens for the query point for the one or more other video frames using the token updating neural network in this way, the system can process, for each query point, each updated query token for each of the other video frames to generate a respective predicted spatial position of the query point in the other video frame.

[0078] In some implementations, the system can use the position prediction head to generate the respective predicted spatial position of each query point in each of the other video frames. For example, the position prediction head can generate, for each query point and for each of the other video frames, a respective probability7distribution over each of an x-coordinate and a y-coordinate of the other video frame. The position prediction head can generate the probability distributions based on the updated query token for the query point for the other video frame.

[0079] For example, each video frame can be a two-dimensional array of pixels, with one of the two dimensions corresponding to an x-coordinate axis and the other of the two dimensions corresponding to a y-coordinate axis. Each coordinate axis of each video frame can be discretized into a fixed number of x- and y-coordinates, respectively. Each spatial location in each video frame can be defined by an x-coordinate (e.g., that defines a location along the x-coordinate axis of the video frame) and a y-coordinate (e.g., that defines a location along the y-coordinate axis of the video frame).

[0080] Each of the probability distributions over the x- and y-coordinates can assign a probability to each x-coordinate and each y-coordinate, respectively. The probability for each coordinate can be a floating-point value that is greater than or equal to 0 and less than or equal to 1 that represents a likelihood that the query point is located at the coordinate in the other video frame.

[0081] The position prediction head can then identity7a spatial location with a maximum probability in the other video frame. For example, the position prediction head can identity7the spatial location with maximum probability by identifying a maximum x-coordinate of the other video frame as the x-coordinate that is assigned the highest probability according to the generated probability7distribution over the x-coordinate for the frame; and a maximum y- coordinate of the other video frame as the y-coordinate that is assigned the highest probability according to the generated probability distribution over the y-coordinate for the frame. The spatial location with maximum probability7can be the spatial location in the other video frame defined by the maximum x-coordinate and the maximum y-coordinate.

[0082] The position prediction head can use the spatial location with the maximum probability to generate the final predicted spatial position for the query point. For example, the position prediction head can identify a set of other spatial locations in the other video frame within a fixed size window of the spatial location with the maximum probability'. For example, the fixed size window can define an area of the other video frame that includes all spatial locations from which the distance to the spatial location with maximum probability is less than or equal to a pre-defined distance.

[0083] The position prediction head can generate the respective predicted spatial position of the query point in the other video frame by computing a weighted sum of the spatial locations within the fixed size window of the spatial location with the maximum probability. The weight for each spatial location can be determined based on a probability assigned to the spatial location by the respective probability distributions.

[0084] For example, the position prediction head can determine a predicted x-coordinate and a predicted y-coordinate. The predicted x-coordinate can be a weighted sum of the x- coordinates of the spatial locations within the fixed-size window of the spatial location with maximum probability. For example, the weights of the x-coordinates in the weighted sum can be the probabilities assigned to the x-coordinates by the generated probability' distribution over the x-coordinate. Similarly, the predicted y-coordinate can be a weighted sum of the y- coordinates of the spatial locations within the fixed-size window of the spatial location with maximum probability, where the weights can be probabilities assigned to the y-coordinates by the generated probability distribution over the y-coordinate. The position prediction head can determine the predicted spatial position of the query7point in the other video frame to be the spatial location in the other video frame defined by the predicted x-coordinate and the predicted y-coordinate.

[0085] In some implementations, the position prediction head can include two position prediction heads. For example, one of the two heads can generate the probability7distribution over the x-coordinate of each video frame. The other of the two heads can generate the probability distribution over the y-coordinate of each video frame. The distributions generated by the two prediction heads can then be used to generate the predicted spatial position of the query7point in the video frame as described above.

[0086] In this way, the system can generate a network output that includes a predicted spatial position of each query point in each of the frames other than the corresponding frame for the query point.

[0087] In some implementations, the system also processes the video tokens and the query7tokens to generate, for each query point, a respective visibility7score for each of the one or more other video frames of the video frames in the video sequence other than the corresponding video frame for the query point (312). The respective visibility scores can be included in the network output along with the respective predicted spatial positions, and each respective visibility7score can indicate whether the query7point is visible in the other video frame. The system can generate the respective visibility scores by processing the updated query tokens using a visibility prediction head included in the system. The visibility prediction head can generate the visibility scores in parallel with the generation of the predicted spatial positions. The visibility7prediction head can generate the visibility7scores before or after the generation of the predicted spatial positions.

[0088] The visibility prediction head can process, for each query point, each updated query7token for each of the one or more other video frames to generate a respective visibility score for the query point for the other video frame. The visibility score can, for example, be a value between 0 and 1 (inclusive), representing a likelihood that the query7point is visible (e.g., not occluded) in the video frame. For example, the inclusion of visibility scores can enhance the network output generated by the system, e.g., by providing valuable contextual information for downstream tasks, indicating a reliability of the predicted spatial positions, or both. FIG. 4 is a flow diagram of an example process 400 for training a neural network that is configured to generate predicted spatial positions for query points in video frames. For convenience, the process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, a video processing system, e.g., the video processing system 100 depicted in Fig. 1, appropriately programmed in accordance with this specification, can perform the process 400. For example, the process 400 can be used to train the neural network included in the video processing system 100 of Fig. 1.

[0089] The system obtains a batch of one or more training video sequences (402). Each of the one or more training video sequences includes multiple video frames. For example, each training video sequence can include a temporally ordered sequence of video frames. Each video frame in the training video sequence can be represented as a multi-dimensional array of pixel data, as described with reference to the operation 302 of Fig. 3.

[0090] The system obtains a set of one or more query points for each of the one or more training video sequences (404). Each query point specifies a respective spatial position in a corresponding one of the video frames for the training video sequence. For example, the set of one or more query points for each training video sequence can be similar to the set of one or more query points described with reference to the operation 304 of Fig. 3.

[0091] The system obtains, for each training video sequence and each query point in the set of query points for the training video sequence, a ground truth spatial position of the query point in each of the video frames in the training video sequence (406). For example, for each query’ point in the set of one or more query' points for the training video sequence, the system can also obtain a complete ground truth trajectory. For example, the complete ground truth trajectory for a query point can include respective ground truth (x, y) coordinates for the query point in each frame of the training video sequence.

[0092] In some implementations, the system can obtain, for each training video sequence, a ground truth visibility7score for each query7point in each video frame. The ground truth visibility score can indicate whether the point is occluded in the video frame.

[0093] For example, the ground truth spatial positions and. optionally, the ground truth visibility scores, can be generated by manual annotation of the training sequences of video frames, e.g., by an operator of the system. For example, the ground truth spatial positions and, optionally, the ground truth visibility scores, can be synthetically generated. For example, the ground truth spatial positions and, optionally, the ground truth visibility scores, can be generated using a point tracking neural network that has already been trained, e.g., that is different than the neural network. The ground truth spatial positions and, optionally, the ground truth visibility scores, can be used for calculating a loss function that is used to train the neural network, as described below. One example training dataset is the TAP-Vid dataset, Doersch et al.. “TAP-vid: A benchmark for tracking any point in a video”, Advances in Neural Information Processing Systems, 35: 13610-13626, 2022.

[0094] For each training video sequence, the system generates a plurality of video tokens representing the training video sequence (408). For each training video sequence, the video tokens representing the training video sequence can include video tokens for image patches into which each video frame of the training video sequence is partitioned. For example, the system can generate the plurality of video tokens for each training video sequence using the techniques described above with reference to the operation 306 of Fig. 3.

[0095] For each training video sequence, for each query point in the set of query points for the training video sequence, the system generates a respective query token for the query point for each video frame in the training video sequence (410). The respective query token for the corresponding video frame for the query point characterizes the respective spatial position of the query point in the corresponding video frame. In some implementations, these query tokens do not encode a temporal position of the corresponding video frame in the training video sequence. The respective uery token for each video frame other than the corresponding video frame for the query point is a learned mask token that indicates that the spatial position of the query point in the video frame is not provided as input to the neural network. For example, the system can generate the query tokens using the techniques described above with reference to the operation 308 of Fig. 3.

[0096] The system processes, for each training video sequence, the plurality of video tokens and the respective query token for each query point for each video frame in the training video sequence using the neural network to generate, for each query point, for each learned mask token for each video frame other than the corresponding video frame for the query point, a respective training output that includes a respective predicted spatial position of the query7point in the video frame (412). For example, the system can process the video tokens and query tokens to generate the respective predicted spatial positions included in the training output using the techniques described above with reference to the operation 310 of Fig. 3. For example, the neural network can include a token updating neural network and a position prediction head that are used by the system to generate the respective predicted spatial positions, as described above with reference to Fig. 3. In implementations in which the system obtains ground truth visibility' scores, the neural network can also include a visibility prediction head, and the training output generated by the system can also include a respective visibility score for each learned mask token for each query point. The generated respective visibility scores can each indicate a prediction by the neural network of whether the query point is visible in the video frame associated with the learned mask token, e.g.. as described with reference to Fig. 3.

[0097] In some implementations, the neural network includes a plurality of layers and processing, for each training video sequence, the plurality’ of video tokens and the respective query’ token for each query’ point for each video frame in the training video sequence using the neural network can include generating, for each query point, for each learned mask token for each video frame other than the corresponding video frame for the query point, a respective auxiliary training output for each of the plurality of layers.

[0098] For example, in such implementations, each of the plurality’ of layers included in the neural network can receive video tokens and query' tokens as input (which can be, e.g., the video tokens and query’ tokens generated by the system at operations 408 and 410, or video tokens and query tokens received from a preceding layer). Each layer can then process the received video tokens and query’ tokens to update query’ tokens that are learned mask tokens, e.g., according to the techniques described above with reference to Fig. 3.

[0099] In such implementations, the neural network can include, for at least some of the plurality of layers, a respective auxiliary position prediction head. The respective auxiliary position prediction head for each layer can generate the auxiliary output for the layer by processing each updated learned mask token to generate a predicted spatial position for the query point associated with the learned mask token in the video frame associated w ith the learned mask token. For example, each auxiliary’ position prediction head can generate the predicted spatial positions for the learned mask tokens using techniques similar to those used by the position prediction head to generate predicted spatial positions (e.g., after the tokens have been updated by all layers of the neural network), as described with reference to Fig. 3. The auxiliary training output for each layer can include the predicted spatial positions for each of the learned mask tokens. The auxiliary training outputs for the layers can be used to generate auxiliary loss terms for training the neural network, as described below.

[0100] The system computes gradients with respect to the parameters of the neural netw ork of a loss function to train the neural network (414), e.g., using backpropagation of gradients of the loss function through layers of the neural network. The loss function includes a first loss term that measures, for each training video sequence, for each query point, for each video frame other than the corresponding frame for the query point, errors between the ground truth spatial position of the query point and the respective predicted spatial position of the query point in the video frame. For example, the respective error measured by the first loss term for each query point for each video frame can be based on a distance between the ground truth spatial position and the respective predicted spatial position. For example, the first loss term can be based on a mean-squared error (MSE) that is based on the distance. For example, the first loss term can be based on a mean-absolute error (MAE) that is based on the distance. In some implementations, the first loss term can be a Huber loss. For example, in such implementations, the first loss term can include, for each training video sequence, for each query point, for each video frame other than the corresponding frame for the query point, a term given by the expression: Where ypis the ground truth spatial position for the query point for the video frame; p is the predicted spatial position for the query point for the video frame; and 8 is a hyperparameter for the first loss term. In implementations in which the system generates visibility scores for the learned mask tokens, the loss function includes a second loss term that measures, for each training video sequence, for each query point, for each video frame other than the corresponding video frame for the query point, errors between the ground truth visibility score for the query point and the respective predicted visibility score for the query point in the video frame. In some implementations, the second loss term can be a cross-entropy loss. For example, in such implementations, the second loss term can include, for each training video sequence, for each query' point, for each video frame other than the corresponding frame for the query point, a term given by the expression:

[0101] LCEyv> yv) = = ~[yv*log log ( ty,) + (1 - yv) Xlog log (1 - ty.) ]

[0102] Where yvis the ground truth visibility score for the query point for the video frame; and yvis the predicted visibility score for the query point for the video frame.

[0103] In implementations in which the system also generates an auxiliary training output for each of a plurality of layers included in the neural network, the loss function can include, for each of the plurality of layers, one or more auxiliary loss terms based on the auxiliary training output for the layer. Each auxiliary' loss term can measure an error between the auxiliary' training output for the layer and the ground truth spatial position. Similar to the first loss term, each auxiliary loss term can measure the error using a MSE, a MAE, a Huber loss, or any of a variety of error functions. For example, in such implementations, the loss function can be given by the following equation:

[0104] Where B is the number of training video sequences in the batch; T is the total number of frames in the video; Q is the number of query points in the set of query points; LHuberis a Huber loss function used to compute the first loss term; ypis the ground truth spatial position for the query point for the video frame;pis the predicted spatial position for the query point for the video frame; LCEis a cross-entropy loss function used to compute the second loss term; yvis the ground truth visibility score for the query point for the video frame; yvis the predicted visibility score for the query point for the video frame; I is the number of layers included in the neural network; Lauxis a loss function used to compute the auxiliary loss terms; ytis the auxiliary training output for each layer; andp, and A, are weights for the first loss term, the second loss term, and each of the auxiliary loss terms, respectively.

[0105] The system can compute the gradients with respect to the loss function using backpropagation. The system can use the gradients to update parameter values of neural network parameters of the neural network. For example, the system can use the gradients to update the parameter values using an update rule of any appropriate gradient descent optimization algorithm, e.g., RMSprop or Adam.

[0106] This specification uses the term "‘configured7’ in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions. With reference to particular neural networks, a neural network may be configured to perform a particular action by being trained to perform that particular action.

[0107] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, e.g., one or more modules of computer program instructions encoded on a tangible non transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be a machine- readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

[0108] The term “data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g.. code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0109] A computer program, which may also be referred to or described as a program, software, a software application, an app. a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it can be deployed in any form, including as a stand alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, sub programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communication network.

[0110] In this specification, the term “database” is used broadly to refer to any collection of data: the data does not need to be structured in any particular way, or structured at all, and it can be stored on storage devices in one or more locations. Thus, for example, the index database can include multiple collections of data, each of which may be organized and accessed differently.

[0111] Similarly, in this specification the term “engine’’ is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers.

[0112] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0113] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read only memory or a random access memory or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

[0114] Computer readable media suitable for storing computer program instructions and data include all forms of non volatile memory, media and memory devices, including by way of example semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto optical disks; and CD ROM and DVD-ROM disks.

[0115] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditor}' feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.

[0116] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and compute-intensive parts of machine learning training or production, e.g., inference, workloads.

[0117] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework or a Jax framework.

[0118] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back end, middleware, or front end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.

[0119] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g.. a result of the user interaction, can be received at the server from the device.

[0120] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.

[0121] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

[0122] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

[0123] What is claimed is:

Claims

CLAIMS1. A method performed by one or more computers, the method comprising: obtaining a video sequence comprising a plurality of video frames; obtaining a set of one or more query points, each query point specifying a respective spatial position in a corresponding one of the plurality of video frames; generating a plurality of video tokens representing the video sequence; for each query point, generating a respective query token for the query point for each of the plurality of video frames, wherein the respective query token for the corresponding video frame for the query point characterizes the respective spatial position of the query point in the corresponding video frame; and processing the video tokens and the respective query tokens for the one or more querypoints using a neural network to generate a network output that comprises, for each query point, a respective predicted spatial position of the query point in one or more other video frames of the video frames in the sequence other than the corresponding video frame for the query point.

2. The method of claim 1, wherein the neural network comprises a token updating neural network and a position prediction head, and wherein processing the video tokens and the respective query tokens for the one or more query- points using the neural network to generate the network output comprises: processing the video tokens and the respective query tokens using the token updating neural network to, for each query point, update the respective query tokens for the query point for the one or more other video frames; and for each query- point and for each of the one or more other video frames for the querypoint, processing the updated query token for the other video frame for the query token using the position prediction head to generate the respective predicted spatial position of the query point in the other video frame.

3. The method of any preceding claim, wherein, for each query point, the respective query token for each video frame other than the corresponding video frame for the query point is a learned mask token that indicates that the spatial position of the query point in the video frame is not provided as input to the neural network.

4. The method of any preceding claim, wherein, for each query' point, the respective query token for the corresponding video frame for the query’ point comprises a spatial positional encoding of the respective spatial position of the query point in the corresponding video frame.

5. The method of claim 4, wherein, for each query point, the respective query token for the corresponding video frame for the query point does not encode a temporal position of the corresponding video frame yvithin the video sequence.

6. The method of any preceding claim, wherein: each video frame is partitioned into a number of non-overlapping image patches; and the plurality of video tokens comprises a respective token for each image patch in each of the video frames.

7. The method of any preceding claim, when dependent on claim 2, wherein: the token updating neural network comprises a plurality of layers that each update the video tokens and the respective query tokens.

8. The method of claim 7, wherein one or more layers of the plurality of layers comprise a recurrent block that performs temporal processing on the video tokens and the respective query tokens.

9. The method of claim 8, wherein the temporal processing comprises, for each query point, updating each of the respective query tokens for the query point for the one or more other video frames based on a respective query token for the query point in an adjacent video frame.

10. The method of any of claims 7-9, yvherein one or more layers of the plurality of layers comprise a self-attention block that performs spatial processing on the video tokens and the respective query tokens.

11. The method of claim 10, wherein the spatial processing comprises, for each query' point, updating each of the respective query tokens for the query point for the one or moreother video frames based on video tokens and respective query tokens for query points in the set of one or more query points other than the query point in the video frame.

12. The method of any preceding claim, wherein: the neural network is configured to process the video tokens and the respective querytokens to generate, for each of the query points, respective probability- distributions over each of an x-coordinate and ay-coordinate for each of the one or more other video frames; and processing the video tokens and the respective query tokens for the one or more query points using the neural network to generate the network output comprises generating the respective predicted spatial position of the query- point in the one or more other video frames based on the respective probability- distributions.

13. The method of claim 12, wherein generating the respective predicted spatial position of the query- point in the one or more other video frames based on the respective probabilitydistributions comprises: identifying a spatial location with a maximum probability in each of the one or more other video frames; identifying, for each spatial location with a maximum probability in each of the one or more other video frames, a set of other spatial locations within a fixed size window of the spatial location with the maximum probability; and generating the respective predicted spatial position of the query point in the one or more other video frames by computing a weighted sum of the spatial locations within the fixed size window of the spatial location w ith the maximum probability, wherein the weight for each spatial location is determined based on a probability assigned to the spatial location by the respective probability distributions.

14. The method of any preceding claim, wherein the network output comprises a visibility score for each query point, wherein the visibility- score indicates whether the query point is visible in the one or more other video frames of the video frames in the video sequence other than the corresponding video frame for the query point.

15. The method of claim 14, when dependent on claim 2, wherein the neural network comprises a visibility- prediction head, and wherein processing the video tokens and therespective query tokens for the one or more query points using the neural network to generate the network output comprises: for each query point and for each of the one or more other video frames for the query point, processing the updated query token for the other video frame for the query token using the visibility prediction head to generate the respective visibility' score for the query point in the other video frame.

16. The method of any preceding claim, when dependent on claim 6, wherein an input to the neural network comprises a tensor with dimensions of T, h X w + Q and C. wherein:T represents a number of video frames in the video sequence; h represents a number of rows of non-overlapping image patches; w represents a number of columns of non-overlapping image patches;Q represents a number of query' points in the set of one or more query points; andC represents a dimension of a space into which the non-overlapping image patches are embedded.

17. The method of any preceding claim, wherein the neural network is trained by: obtaining a batch of one or more training video sequences, yvherein each of the one or more training video sequences comprises video frames; obtaining a set of one or more query points for each of the one or more training video sequences, each query point specifying a respective spatial position in a corresponding one of the video frames for the training video sequence; obtaining, for each training video sequence and each query point in the set of query points for the training video sequence, a ground truth spatial position of the query point in each of the video frames in the training video sequence; for each training video sequence, generating a plurality' of video tokens representing the training video sequence; for each training video sequence, for each query point in the set of query points for the training video sequence, generating a respective query token for the query point for each video frame in the training video sequence, yvherein the respective query token for the corresponding video frame for the query point characterizes the respective spatial position of the query point in the corresponding video frame and the respective query token for each video frame other than the corresponding video frame for the query point is a learned masktoken that indicates that the spatial position of the query point in the video frame is not provided as input to the neural network; processing, for each training video sequence, the plurality of video tokens and the respective query token for each query point for each video frame in the training video sequence using the neural network to generate, for each query point, for each learned mask token for each video frame other than the corresponding video frame for the query point, a respective training output that includes a respective predicted spatial position of the query point in the video frame; and computing gradients with respect to the parameters of the neural network of a loss function that includes a first loss term that measures, for each training video sequence, for each query point, for each video frame other than the corresponding frame for the query point, errors between the ground truth spatial position of the query point and the respective predicted spatial position of the query point in the video frame.

18. The method of claim 17, further comprising: obtaining, for each training video sequence and each query point in the set of query points for the training video sequence, a respective ground truth visibility score for the query point for each of the video frames in the training video sequence, wherein: for each learned mask token for each video frame other than the corresponding video frame for the query point, the respective training output includes a respective predicted visibility score for the query point in the video frame; and the loss function includes a second loss term that measures, for each training video sequence, for each query point, for each video frame other than the corresponding video frame for the query point, errors between the ground truth visibility score for the query point and the respective predicted visibility score for the uery point in the video frame.

19. The method of claim 18, wherein the second loss term is a cross-entropy loss.

20. The method of any one of claims 17-19, wherein the first loss term is a Huber loss.

21. The method of any preceding claim, when dependent on claim 2, wherein the position prediction head is a multi-layer perceptron network.

22. The method of any one of claims 17-21, when dependent on claim 17, wherein the neural network comprises a plurality of layers, wherein processing, for each training video sequence, the plurality of video tokens and the respective query token for each query point for each video frame in the training video sequence using the neural network comprises: generating, for each query7point, for each learned mask token for each video frame other than the corresponding video frame for the query point, a respective auxiliary training output for each of the plurality of layers, and wherein the loss function further comprises, for each of the plurality of layers, one or more auxiliary7loss terms based on the auxiliary training output for the layer.

23. A system comprising: one or more computers; and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the respective method of any one of claims 1 -22.

24. One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations of the respective method of any one of claims 1-22.

Images