Multimodal multitask alternating gradient descent

Abstract

The technology involves integrating a set of multimodal tasks into a single training pipeline with a single shared model, which can employ alternating gradient descent and mixture of experts processing. A plurality of datasets is used to train the model, where each dataset has a different modality. A set of loss objectives is associated with the plurality of datasets. In one scenario, a process includes performing, by a neural network, embedding of the different modalities into a shared modality - agnostic space to obtain a set of embedded tokens. The process includes determining, by the neural netw ork according to a mixture of experts encoder module, semantic contextual embeddings from the embedded tokens. Then the process includes generating, by the neural network according to a set of encoder heads, a set of final predictions by re-projecting the semantic contextual embeddings into a modality -specific space according to the set of loss objectives.

Description

MULTIMODAL MULTITASK ALTERNATING GRADIENT DESCENTCROSS-REFERENCE TO RELATED APPLICATIONS

[0001] This application claims the benefit of and priority to the filing date of U.S. Provisional Application No. 63 / 465,002, filed May 9, 2023, the entire disclosure of which is hereby incorporated by reference herein.BACKGROUND

[0002] Machine learning techniques can be used to address many different ty pes of applications, such as image recognition and natural language processing. Certain machine learning approaches employ large-scale multitask frameworks and models, which can provide foundations for integrating certain capabilities that may help to unify disparate tasks under one model. However, given the vast quantity of independent variables involved in designing such a system, achieving an integrated multimodal machine learning model can be highly complex.

[0003] More specifically, designing a multi-task model that integrates many multimodal signals is challenging due to various reasons. For instance, modalities such as text, vision, and audio may require structurally different I / O signatures to properly train. Standard input configurations for different modalities and output loss objectives for different tasks may be incompatible. For instance, audio classification may require 2D spectrogram inputs with a linear head while video-text retrieval may require 3D color (e.g., RGB) input with one or more contrastive heads. Also, when training across multiple datasets, some modalities or objectives may not exist or cannot be applied, depending on the input data and the task to perform. For example, a video-text dataset may be used for captioning, but another dataset might only have unlabeled audio for self-supervised training. In addition, the presence of multiple input modalities calls for careful considerations on the architectural design and allocation of parameters to certain modalities, often requiring extensive hyperparameter tuning to find the best use of computational resources.BRIEF SUMMARY

[0004] The technology relates to integrating any number of heterogeneous unimodal and multimodal tasks into a single training pipeline with a single shared model. This may be accomplished without slowing down the model training or compromising the model’s accuracy, which are significant technical advantages for such a system. As a model is scaled, it can become increasingly expensive to redesign the architecture or search for a better training objective. The issue is exacerbated in multimodal multi-task modeling, where there is a need to consider the combination of input modalities or datasets, loss functions, and tasks at large scales. Therefore, one aspect of the technology involves a training approach that can be scaled incrementally. In particular, for any new task or objective, regardless of its input shape or output loss, the system may be able to add it to the existing pretraining without compromising the previous tasks.

[0005] As discussed herein, integrated multimodal perception (IMP) provides a simple and scalable multimodal multi-task training and modeling approach. This approach integrates multimodal inputs including image, video, text, and / or audio into a single encoder, such as a transformer encoder, with minimal modality-specific components. IMP makes use of a technical solution that combines alternating multi-task training and Mixture-of-Experts (MoE) for efficient model scaling. Alternating gradient descent (AGD) can be applied to efficiently improve multimodal understanding. In addition, model sparsification with MoE on a single modality -agnostic encoder can substantially improve the performance of the system.

[0006] According to one aspect of the technology, a computer-implemented method is provided for multimodal multi-task training of an integrated multimodal perception model. The method comprises: obtaining, by one or more processors of a computing system, a plurality’ of datasets, each dataset having a different modality, each dataset also having an input resolution; obtaining, by the one or more processors, a set of loss objectives associated with the plurality of datasets; performing, by a neural network implemented via the one or more processors, embedding of the different modalities into a shared modality -agnostic space to obtain a set of embedded tokens; determining, by the neural network according to a mixture of experts encoder module, semantic contextual embeddings from the embedded tokens; and generating, by the neural netw ork according to a set of encoder heads, a set of final predictions by re-projecting the semantic contextual embeddings into a modality -specific space according to the set of loss objectives.

[0007] The different modalities of the plurality’ of datasets may be selected from the set consisting of imagery, video, text, and audio modalities in any combination thereof. The input resolution of each dataset may be distinct from the resolution of the other datasets. The set of loss objectives may be selected from the set consisting of image-text contrastive, video-text contrastive, video-audio contrastive, image classification, video classification, and audio classification.

[0008] In one scenario, the generating is performed according to an alternating gradient descent process. At each gradient step, a different loss objective may be optimized given one or both of different sets of model w eights or input modalities. Here, the set of final predictions may include at least one of softmax probabilities, aggregation heads for noise contrastive estimation, or per-sequence heads for language modeling. A set of tasks to be performed by the alternating gradient descent process may be performable in any order. In an example here, the set of tasks to be performed includes one or more of supervised classification, masked auto encoding, image-text contrastive, and video-audio-text contrastive tasks.

[0009] The mixture of experts encoder module may be a sparse mixture of experts encoder module.Alternatively’ or additional to any’ of the above, the method may further comprise defining multiple variants of each dataset, in which each variant of a given dataset corresponds to a different number of spatiotemporal patches or a different batch size. Alternatively’ or additional to any of the above, themethod may further comprise performing inference using trained model to generate one or more results. Moreover, performing the embedding may further include performing flatting of the different modalities.

[0010] According to another aspect of the technology, a system for multimodal multi-task training of an integrated multimodal perception model is provided. The system comprises memory configured to store a plurality of datasets, in which each dataset has a different modality and also has an input resolution, and one or more processors operatively coupled to the memory. The one or more processors are configured to: obtain the plurality of datasets from the memory: obtain a set of loss objectives associated with the plurality of datasets; perform, by a neural network implemented via the one or more processors, embedding of the different modalities into a shared modality -agnostic space to obtain a set of embedded tokens; determine, according to a mixture of experts encoder module, semantic contextual embeddings from the embedded tokens; and generate, according to a set of encoder heads, a set of final predictions by re-projecting the semantic contextual embeddings into a modality-specific space according to the set of loss objectives.

[0011] Generation of the set of final predictions may be performed according to an alternating gradient descent process. Here, at each gradient step, a different loss objective may be optimized given one or both of different sets of model weights or input modalities. The set of final predictions may include at least one of softmax probabilities, aggregation heads for noise contrastive estimation, or persequence heads for language modeling. A set of tasks to be performed by the alternating gradient descent process may be performable in any order.

[0012] The mixture of experts encoder module may be a sparse mixture of experts encoder module. The one or more processors may be further configured to define multiple variants of each dataset, each variant of a given dataset corresponding to a different number of spatiotemporal patches or a different batch size. The one or more processors may be further configmed to perform inference using trained model to generate one or more results.BRIEF DESCRIPTION OF THE DRAWINGS

[0013] Figs. 1A-B illustrate example IMP training and architecture configurations in accordance with aspects of the technology.

[0014] Fig. 2 illustrates a general transformer architecture for use with aspects of the technology.

[0015] Figs. 3A-B illustrates a comparison of mixed batch training and alternating gradient descent, and an AGD example, in accordance with aspects of the technology.

[0016] Fig. 4 illustrates an alternating gradient descent algorithm in accordance with aspects of the technology.

[0017] Fig. 5 illustrates a table of results according to aspects of the technology.

[0018] Fig. 6 illustrates a table of results according to aspects of the technology.

[0019] Fig. 7 illustrates a table of results according to aspects of the technology.

[0020] Fig. 8 illustrates a table of results according to aspects of the technology.

[0021] Fig. 9 illustrates a table of results according to aspects of the technology.

[0022] Fig. 10 illustrates a table of results according to aspects of the technology.

[0023] Fig. 11 presents a chart of results according to aspects of the technology.

[0024] Fig. 12 presents a chart of results according to aspects of the technology.

[0025] Fig. 13 presents a chart of results according to aspects of the technology.

[0026] Fig. 14 presents a chart of results according to aspects of the technology.

[0027] Fig. 15 presents a chart of results according to aspects of the technology.

[0028] Fig. 16 presents a chart of results according to aspects of the technology.

[0029] Fig. 17 illustrates a table of results according to aspects of the technology.

[0030] Fig. 18 illustrates a table of results according to aspects of the technology.

[0031] Fig. 19 illustrates a table of results according to aspects of the technology.

[0032] Fig. 20 illustrates an example implementation for integrated multimodal perception in accordance with aspects of the technology.

[0033] Figs. 21(a)-(b) illustrate a system for use with aspects of the technology.

[0034] Fig. 22 illustrates an example method in accordance with aspects of the technology .

[0035] The patent or application fdc contains at least one drawing executed in color. Copies of this patent or patent application publication with color drawing(s) will be provided by the Office upon request and payment of the necessary' fee.DETAILED DESCRIPTION

[0036] The technology involves combined multisensory learning in the machine learning context. Scalability is important, because as a model is scaled, it can become increasingly expensive to redesign the architecture or search for a better training objective. The issue can be exacerbated in multimodal multi-task modeling, where one needs to consider the combination of input modalities or datasets, loss functions, and tasks at large scales. Therefore, an aspect of the technology includes a training approach that can be scaled incrementally: for any new task or objective, regardless of its input shape or output loss, the system should be able to add it to the existing pretraining without compromising the previous tasks.

[0037] The scaling issue can be addressed via a solution that trains one multimodal model such that it (1) leverages as many existing datasets as possible, (2) can train on any combination of tasks or loss functions, and (3) does not slow down with the addition of any new dataset, task, or loss function. By solving all of these points simultaneously, the multimodal model can scale with an increasingly diverse and rich set of training data without needing to redesign the training framework when new tasks are integrated.

[0038] For instance, a transformer encoder may be tasked with integrating two different objectives: image classification and image-text noise contrastive estimation. Either task can be suitable for buildinga strong vision understanding model, but combining them could possibly be even better. However, one question is how one should combine them. Two approaches are: (1) either sum the two objectives in a single forward-backward pass, or (2) alternate steps between both objectives.

[0039] It has been discovered that the combination of diverse, heterogeneous tasks that have been previously established as strong objectives individually (e.g.. supervised classification, image-text contrastive learning) across multiple modalities are not only complementary, but can offer better convergence than training on individual tasks. It has been observed that naively summing loss objectives across such a wide range of tasks and modalities often results in poor convergence and inefficient training. Thus, according to one aspect of the technology an efficient and scalable solution implements AGD via selected primitives, such as according to a domain -specific tracing just-in-time (JIT) compiler, such as JAX.

[0040] This arrangement enables an IMP model to use a fraction of the computational cost and memory required by large-scale perception models (see, e.g.. “Learning transferable visual models from natural language supervision” by Radford et al., 2021; “Scaling up visual and vision-language representation learning with noisy text supervision” by Jia et al.. 2021; and “Contrastive captioners are image-text foundation models” by Yu et al., 2022), despite the addition of multiple modalities which could otherwise require 2-8 x compute at similar batch sizes.

[0041] The following introduces the certain aspects and benefits of the technology, which are detailed further below. First, a single, modality -agnostic encoder model is defined, which is able to leverage a strong combination of, e.g., image-text contrastive, video-text contrastive, video-audio contrastive, and image / video / audio classification losses during pretraining to create an image, video, audio, and text understanding model. An example 100 of the approach is illustrated Fig. 1A. As seen at (1), there is a mixture of datasets with varying modalities, resolutions, and loss objectives. The datasets may include video, imagery, text and / or audio datasets. As seen by the arrows (2), these datasets are randomly sampled with weighted probability at each optimization step and fed into the model. Then, as shown at (3). the system runs a forward-backward pass on the input batch and objective(s).

[0042] The arrangement as shown includes an embedding and flattening section, an MoE encoder section, and an encoder head section. The embedding and flattening section may apply ID text projection, 3D vision patch projection for imagery, ID waveform patch projection and / or 2D spectrogram patch projection for audio, as well as 3D positional encoding. The system linearly projects each modality into the same representation space, which may involve applying an MoE Transformer encoder in the same manner for all modalities from the prior stage. Then, simple projection heads are applied on the encoder to compute the respective objective losses that are applicable to the encoded features. In one scenario, the system can use jax.jit to compile and cache computation graphs to keepeach step efficient while also allowing input / output shapes to change for every optimization step without requiring any costly padding or masking strategies.

[0043] Fig. IB illustrates a high-level view 120 showing an overall approach for IMP. This figure illustrates the use of alternating gradient descent, and how an MoE Transformer encoder can be employed for training and inference. Note that the system can also apply a decoder for generative tasks like language modeling. In addition, MoE can also be used in the decoder to increase the parameters of the network and provide better generalization.

[0044] Contrasting the approach of summing the losses of multiple objectives, it is shown that alternating between objectives results in a design that allows seamless integration of virtually any number of tasks and datasets without significant memory overhead and results in better downstream evaluations. It is also shown that optimization between multiple heterogeneous multimodal tasks is complementary and results in a higher quality model than trained on any individual task. To train on large batches of video and audio modalities without reducing training efficiency or loss of accuracy, a dynamic mixture of various resolutions, sequence lengths, and batch sizes is employed throughout pretraining, and alternate training on input variations. Sparse mixture-of-experts (MoE) can be integrated in a single modality -agnostic encoder model, showing strong performance gains compared to an existing multi-tower contrastive model, even when applying MoE to both towers. Also, the resulting sparse IMP model can be scaled to a vast number of parameters (e.g., millions or billions of parameters) with similar compute to ViT-L (which has 3 OOM parameters), resulting in state-of-the-art evaluations on several large-scale image and video understanding datasets.General Transformer Approach

[0045] The techniques discussed herein may employ a self-attention architecture, e.g., the Transfonner neural network encoder-decoder architecture. An exemplary general Transformer-type architecture is shown in Fig. 2. which is based on the arrangement shown in U.S. Patent No. 10,452,978, entitled “Attention-based sequence transduction neural networks”, the entire disclosure of which is incorporated herein by reference.

[0046] System 200 of Fig. 2 is implementable as computer programs by processors of one or more computers in one or more locations. The system 200 receives an input sequence 202 and processes the input sequence 202 to transduce the input sequence 202 into an output sequence 204. The input sequence 202 has a respective network input at each of multiple input positions in an input order and the output sequence 204 has a respective network output at each of multiple output positions in an output order.

[0047] System 200 can perform any of a variety of tasks that require processing sequential inputs to generate sequential outputs. System 200 includes an attention-based sequence transduction neural network 206, which in turn includes an encoder neural network 208 and a decoder neural network 210. The encoder neural network 208 is configmed to receive the input sequence 202 and generate arespective encoded representation of each of the network inputs in the input sequence. An encoded representation is a vector or other ordered collection of numeric values. The decoder neural network 210 is then configured to use the encoded representations of the network inputs to generate the output sequence 204. Generally, both the encoder 208 and the decoder 210 are attention-based. In some cases, neither the encoder nor the decoder includes any convolutional layers or any recurrent layers. The encoder neural network 208 includes an embedding layer (input embedding) 212 and a sequence of one or more encoder subnetworks 214. The encoder neural 208 network may N encoder subnetworks 214.

[0048] The embedding layer 212 is configured, for each network input in the input sequence, to map the network input to a numeric representation of the network input in an embedding space, e g., into a vector in the embedding space. The embedding layer 212 then provides the numeric representations of the network inputs to the first subnetwork in the sequence of encoder subnetworks 214. The embedding layer 212 may be configured to map each network input to an embedded representation of the network input and then combine, e.g., sum or average, the embedded representation of the network input with a positional embedding of the input position of the network input in the input order to generate a combined embedded representation of the network input. In some cases, the positional embeddings are learned. As used herein, “learned” means that an operation or a value has been adjusted during the training of the sequence transduction neural netw ork 206. In other cases, the positional embeddings may be fixed and are different for each position.

[0049] The combined embedded representation is then used as the numeric representation of the network input. Each of the encoder subnetworks 214 is configured to receive a respective encoder subnetwork input for each of the plurality of input positions and to generate a respective subnetwork output for each of the plurality of input positions. The encoder subnetwork outputs generated by the last encoder subnetwork in the sequence are then used as the encoded representations of the network inputs. For the first encoder subnetwork in the sequence, the encoder subnetwork input is the numeric representations generated by the embedding layer 212, and, for each encoder subnetwork other than the first encoder subnetwork in the sequence, the encoder subnetwork input is the encoder subnetwork output of the preceding encoder subnetwork in the sequence.

[0050] Each encoder subnetwork 214 includes an encoder self-attention sub-layer 216. The encoder self-attention sub-layer 216 is configured to receive the subnetwork input for each of the plurality of input positions and, for each particular input position in the input order, apply an attention mechanism over the encoder subnetwork inputs at the input positions using one or more queries derived from the encoder subnetwork input at the particular input position to generate a respective output for the particular input position. In some cases, the attention mechanism is a multi-head attention mechanism as shown. In some implementations, each of the encoder subnetworks 214 may also include a residual connection layer that combines the outputs of the encoder self-attention sub-layer with the inputs to the encoder self-attention sub-layer to generate an encoder self-attention residual output and alayer normalization layer that applies layer normalization to the encoder self-attention residual output. These two layers are collectively referred to as an "Add & Norm” operation in Fig. 2.

[0051] Some or all of the encoder subnetworks can also include a position-wise feed-forward layer 218 that is configured to operate on each position in the input sequence separately. In particular, for each input position, the feed-forward layer 218 is configured receive an input at the input position and apply a sequence of transformations to the input at the input position to generate an output for the input position. The inputs received by the position-wise feed-forward layer 218 can be the outputs of the layer normalization layer when the residual and layer normalization layers are included or the outputs of the encoder self-attention sub-layer 216 when the residual and layer normalization layers are not included. The transformations applied by the layer 218 will generally be the same for each input position (but different feed-forward layers in different subnetworks may apply different transformations).

[0052] In cases where an encoder subnetwork 214 includes a position-wise feed-forward layer 218 as shown, the encoder subnetwork can also include a residual connection layer that combines the outputs of the position-wise feed-forward layer with the inputs to the position-wise feed-forward layer to generate an encoder position-wise residual output and a layer normalization layer that applies layer normalization to the encoder position-wise residual output. As noted above, these two layers are also collectively referred to as an "Add & Norm" operation. The outputs of this layer normalization layer can then be used as the outputs of the encoder subnetwork 214.

[0053] Once the encoder neural network 208 has generated the encoded representations, the decoder neural netw ork 210 is configured to generate the output sequence in an auto-regressive manner. That is, the decoder neural network 210 generates the output sequence, by at each of a plurality of generation time steps, generating a network output for a corresponding output position conditioned on (i) the encoded representations and (ii) network outputs at output positions preceding the output position in the output order. In particular, for a given output position, the decoder neural network generates an output that defines a probability distribution over possible network outputs at the given output position. The decoder neural network can then select a network output for the output position by sampling from the probability distribution or by selecting the network output with the highest probability.

[0054] Because the decoder neural network 210 is auto-regressive, at each generation time step, the decoder network 210 operates on the network outputs that have already been generated before the generation time step, i.e., the network outputs at output positions preceding the corresponding output position in the output order. In some implementations, to ensure this is the case during both inference and training, at each generation time step the decoder neural netw ork 210 shifts the already generated network outputs right by one output order position (i.e., introduces a one position offset into the already generated network output sequence) and (as will be described in more detail below) masks certain operations so that positions can only attend to positions up to and including that position in the output sequence (and not subsequent positions). While the remainder of the description below' describes that.when generating a given output at a given output position, various components of the decoder 210 operate on data at output positions preceding the given output positions (and not on data at any other output positions), it will be understood that this type of conditioning can be effectively implemented using shifting.

[0055] The decoder neural network 210 includes an embedding layer (output embedding) 220, a sequence of decoder subnetworks 222, a linear layer 224, and a softmax layer 226. In particular, the decoder neural network can include N decoder subnetworks 222. However, while the example of Fig. 2 shows the encoder 208 and the decoder 210 including the same number of subnetworks, in some cases the encoder 208 and the decoder 210 include different numbers of subnetworks. The embedding layer 220 is configured to, at each generation time step, for each network output at an output position that precedes the current output position in the output order, map the network output to a numeric representation of the network output in the embedding space. The embedding layer 220 then provides the numeric representations of the network outputs to the first subnetwork 222 in the sequence of decoder subnetworks.

[0056] In some implementations, the embedding layer 220 is configured to map each network output to an embedded representation of the network output and combine the embedded representation of the netw ork output w ith a positional embedding of the output position of the netw ork output in the output order to generate a combined embedded representation of the network output. The combined embedded representation is then used as the numeric representation of the network output. The embedding layer 220 generates the combined embedded representation in the same manner as described above with reference to the embedding layer 212.

[0057] Each decoder subnetwork 222 is configured to, at each generation time step, receive a respective decoder subnetwork input for each of the plurality of output positions preceding the corresponding output position and to generate a respective decoder subnetwork output for each of the plurality of output positions preceding the corresponding output position (or equivalently, when the output sequence has been shifted right, each network output at a position up to and including the current output position). In particular, each decoder subnetwork 222 includes two different attention sub-layers: a decoder self-attention sub-layer 228 and an encoder-decoder attention sub-layer 230. Each decoder self-attention sub-layer 228 is configured to, at each generation time step, receive an input for each output position preceding the corresponding output position and, for each of the particular output positions, apply an attention mechanism over the inputs at the output positions preceding the corresponding position using one or more queries derived from the input at the particular output position to generate a updated representation for the particular output position. That is, the decoder self-attention sub-layer 228 applies an attention mechanism that is masked so that it does not attend over or otherwise process any data that is not at a position preceding the current output position in the output sequence.

[0058] Each encoder-decoder atention sub-layer 230, on the other hand, is configured to, at each generation time step, receive an input for each output position preceding the corresponding output position and, for each of the output positions, apply an atention mechanism over the encoded representations at the input positions using one or more queries derived from the input for the output position to generate an updated representation for the output position. Thus, the encoder-decoder atention sub-layer 230 applies atention over encoded representations while the decoder self-atention sub-layer 228 applies atention over inputs at output positions.

[0059] In the example of Fig. 2. the decoder self-atention sub-layer 228 is shown as being before the encoder-decoder atention sub-layer in the processing order within the decoder subnetwork 222. In other examples, however, the decoder self-attention sub-layer 228 may be after the encoder-decoder atention sub-layer 230 in the processing order within the decoder subnetwork 222 or different subnetworks may have different processing orders. In some implementations, each decoder subnetwork 222 includes, after the decoder self-atention sub-layer 228, after the encoder-decoder atention sublayer 230, or after each of the two sub-layers, a residual connection layer that combines the outputs of the atention sub-layer with the inputs to the atention sub-layer to generate a residual output and a layer normalization layer that applies layer normalization to the residual output. These two layers being inserted after each of the two sub-layers, both referred to as an "Add & Norm" operation.

[0060] Some or all of the decoder subnetw ork 222 also include a position-wise feed-forward layer 232 that is configured to operate in a similar maimer as the position-wise feed-forward layer 218 from tire encoder 208. In particular, the layer 232 is configured to, at each generation time step: for each output position preceding the corresponding output position: receive an input at the output position, and apply a sequence of transformations to the input at the output position to generate an output for the output position. The inputs received by the position-wise feed-forward layer 232 can be the outputs of the layer nonnalization layer (following the last atention sub-layer in the subnetwork 222) when the residual and layer normalization layers are included or the outputs of the last atention sub-layer in the subnetwork 222 when the residual and layer normalization layers are not included. In cases where a decoder subnetwork 222 includes a position-wise feed-forward layer 232, the decoder subnetwork can also include a residual connection layer that combines the outputs of the position-wise feed-forward layer wdth the inputs to the position-wise feed-forward layer to generate a decoder position-wise residual output and a layer normalization layer that applies layer normalization to the decoder position-wise residual output. These two layers are also collectively referred to as an "Add & Norm" operation. The outputs of this layer normalization layer can then be used as the outputs of the decoder subnetwork 222.

[0061] At each generation time step, the linear layer 224 applies a learned linear transformation to the output of the last decoder subnetwork 222 in order to project the output of the last decoder subnetwork 222 into the appropriate space for processing by the softmax layer 226. The softmax layer 226 then applies a softmax function over the outputs of the linear layer 224 to generate the probabilitydistribution (output probabilities) 234 over the possible network outputs at the generation time step. The decoder 210 can then select a network output from the possible network outputs using the probability distribution.Exemplary Methods and Implementations

[0062] As noted above, one aspect of the technology involves task scalability. In particular, different combinations of data and loss objectives should be interchangeable throughout training, while the addition of any new data or objective should not cause memory or computation overhead. It has been found that a common issue in training large-scale foundation models in a distributed setting is that input signatures and loss objectives need to be static to avoid major inefficiencies. Using transforms like torch.compile or jax.jit allow for low-level graph optimizations that maximize hardware FLOPs utilization on distributed devices such as GPUs and TPUs. but can come at a cost of requiring static I / O signatures. However, in the context of multimodal foundation models, it is desirable to leverage multiple datasets with different modalities and input signatures and with different objective functions.

[0063] One approach to handle the issue with the static input signature would be to use mixed batching, where all possible inputs are constructed, and inapplicable inputs for a given dataset are padded and outputs are masked accordingly at each training step. This would require minimal changes to the underlying optimization code, but comes at a great efficiency cost. The more tasks that arc added, the more time is spent computing on padded inputs, in other words calculating operations that do not influence the optimization process.

[0064] An issue with having multiple objective functions can be resolved by mixed mini-batching, where in addition to padding / masking the inputs / outputs, a batch is divided to multiple mini-batches with their corresponding objective functions. The gradients for each mini-batch and objective function pair may be calculated and accumulated across multiple mini-batches and the model weights are updated once using an aggregated gradient. However, this approach can also be difficult to scale since the gradients across multiple mini-batches are accumulated in memory and per-task batch size naturally reduces as more tasks are added.

[0065] To address the problem more generally, the system uses the AGD optimization technique (see. e.g., “Non-convex optimization for machine learning” by Jain et al., 2017). Fig. 3A compares mixed batch training (302) to the AGD approach (304). Note that AGD can be seen as a superset of stochastic gradient descent (SGD), where at each gradient step, a different loss objective may be optimized given different sets of model weights and / or input modalities.

[0066] In mixed batch settings, the model is compiled with all available inputs and outputs. For a given step (t). certain inputs may be unavailable, so padding and masking may be required, reducing training efficiency linearly with the number of tasks. In contrast, in AGD, the training loop is assigned a pair of inputs and loss objectives and performs backpropagation on each pair. In combination withjax.jit, each unique inpul / output signature causes the graph to be re-compiled at the time of introduction and is cached for fast execution. Tasks can be executed in any order using a scheduling function f(t).

[0067] Fig. 3A provides an example of how AGD might play out in multimodal multitask optimization over a series of steps. Inputs can consist of images, video, audio, or text. Fig. 3B illustrates one example of AGD training for supervised classification, masked auto encoding (MAE), image-text contrastive, and video-audio-text contrastive. The system can define a loop over all of the tasks and execute each task one by one by calling the jit-ed model function on the task. Each task has a dataset / objective pair that are all pre-defined before beginning the experiment. The ordering can depend on the choice of scheduling algorithm. Fig. 3B illustrates the simple case of round-robin scheduling.

[0068] Outputs from the AGD process can be, e g., softmax probabilities, aggregation heads for noise contrastive estimation (NCE), per-sequence heads for language modeling, etc. Correspondingly, the loss function could be cross-entropy, NCE, mean squared error, or any related function that is desired.

[0069] A substantial benefit of the AGD approach is that each task can also be scheduled in any order, enabling stochastic sampling methods to influence the execution of tasks, as seen in curriculum learning methods. Furthermore, unlike in AGD, in the mixed batch setting it is shown that summing the gradients results in an optimization problem that is less convex. This is shown empirically by various test results, discussed below.

[0070] The optimization problem may be approached in a general setting, allowing any changes to the three elements of the optimization system: inputs, model, and objective. More specifically, any input modality with arbitrary shape could consume any subset of tire model while focusing on minimizing any specific combination of objective functions. Furthermore, the limits of scaling may be pushed in these directions in order to study the effect of each on the overall downstream results. Testing has been done to show this, with empirical results to assess the efficacy and validity of the IMP approach.

[0071] From a teclmical point of view, the approach involves compiling multiple computation graphs, one for each unique task structure. This is because instead of defining the structure of the optimization in one large static computation graph to handle all tasks as in mixed batch training, AGD instead factors them into individual smaller graphs and executes a forward-backward pass on each individual graph. This allows mixing and matching of I / O, model, and objective in their own respective graphs.

[0072] To enable efficient execution of multiple computation graphs, according to one implementation JAX can be used, as it offers native Just-in-Time (JIT) compilation with the jax.jit API. which compiles the graph-of-interest at runtime and compiles a new graph if a change in structure is seen in any of the next optimization steps. Graphs themselves can be cached by JAX in an in-memory lookup table so that tasks only need to be compiled once. As long as the number of new I / O structures are not too frequent, compilation only constitutes a small fraction of total training time, and the cachedgraphs themselves take up only a small amount of memory. Further reductions can also be achieved in compilation time using compiler support for some types of dynamic input shapes. In experiments discussed below, up to 20 unique I / O structures were tested with no significant reduction in training speed, with compilation taking only a few minutes of the total training time, up to the largest scales.

[0073] An example AGD algorithm for the multimodal framework is provided in Algorithm 1 shown in Fig. 4. Dataset-objective pairs are sampled at each training step using a sampling function and the model optimization is applied with the dataset input and objective function. This algorithm does not apply any modifications to the optimizer, but uses it as-is to share optimization states across tasks. Note that despite its simplicity, the while loop has been carefully designed to incorporate only the minimum knowledge necessary about the data, loss function, and model. This allows for greater flexibility in defining logic inside the model to handle the processing of individual modalities and input shapes. An underlying goal is to allow the input data to provide the necessary information to dynamically control how the model should process any piece of data. This could be as simple as routing a specific modality to a different encoder, or more complex such as altering local hyperparameters on each step like input resolution, dropout, parameter-wise learning rate. MoE capacity factor, etc.

[0074] The sampling function can also incorporate state from the optimization process itself. For instance, the loss for each task in the data mixture can be used to affect different sampling behavior depending on the loss value, as seen in dynamic pretraining. In a default setup, each unique task may be sampled on a given step from a (single trial) multinomial distribution with probabilities directly proportional to the number of examples in each task.

[0075] To improve AGD memory and training time efficiency even further, some additional teclmiques can be applied for further optimization. For instance, after compiling each graph, it may be necessary' to store all of the activations to perfonn a single gradient descent step, which can consume large amounts of memory. To reduce memory overhead, one implementation uses JAX’s native rematerialization API, jax.checkpoint, to save memory by not checkpointing the activations within each transformer layer during the forward call and only store them between layers. In certain experiments, an average reduction of 70-80% TPU HBM usage was observed while resulting in only 18-20% longer step times.

[0076] Also, large models with many different objectives may still incur long compilation times. Therefore, the system may apply scan-over-layers with jax.lax. scan, a method which compresses the construction of all transformer layers into tensors with a function that scans over the first axis. This computes a single transformer layer regardless of depth, which can result in 20-40x faster compilation time, with relative time savings increasing with larger model sizes.

[0077] These improvements, can be accomplished across distributed accelerators through the jax.pjit API3. which distributes JIT compilation across multiple accelerators.

[0078] In one scenario, in view of a goal to reuse objectives that have been shown to be robust for learning each modality, established supervised and unsupervised learning objectives can be used. These include (1) supervised classification, in which labels are provided for image recognition, video action recognition, and audio recognition objectives, and (2) Noise Contrastive Estimation (NCE). NCE is calculated betw een different pairs of modalities given the available data: image-text, video-text, audiotext, video-audio. Unless otherwise specified herein, the above losses are not summed or gradients accumulated. Instead, backprop can be applied on each objective individually with AGD.

[0079] Returning to Fig. 1 A, it can be seen that the IMP architecture has three main modules: the embedder module 102, the MoE encoder module 104, and the encoder heads module 106. The embedder module 102 accepts specific modalities and embeds them in a shared modality -agnostic space. The MoE encoder module 104 computes semantic contextual embeddings from the embedded tokens. The encoder heads module 106 produces all the final predictions from the encoder by reprojecting its embeddings back into a modality-specific space.Input Embeddings

[0080] For vision modalities, the VATT scheme (see, e.g., “Vatt: Transformers for multimodal self-supervised learning from raw video, audio and text” by Akbari et al., 2021) may be used to patchily each 3D video tensor. A video tensor of size F x H x W is defined with F frames and H x W resolutionF H VK using a patch size of x h x w, producing - x - x — x 3 voxels that are flattened into a single sequence.This is followed by linearly projecting the sequence into the model’s hidden size. To allow for more robust generalization, images may be treated as a special case of a video, assuming sequences are of shape fxH x W and tiling frames f times to fit in a single patch. For a base model, a patch kernel size of 4x16x16 may be used, up to 16 frames, and resolutions up to 512x512. Following VATT, to allow' the model to adapt to different resolution scales, learnable positional encodings can be applied to each patch position, which consist of the sum of 3 embedding tables along each separate axis: one for the temporal dimension, and two for tire spatial dimensions of the video patches.

[0081] For text, the system may apply T5 encoding (see, e.g., “Exploring the limits of transfer learning with a unified text-to-text transformer” by Raffel et al., 2020). using a default English vocabulary with 32k SentencePiece tokens, which are embedded into the same hidden space as image patches.

[0082] For audio, the system may apply both wavefonn and audio spectrogram as input. For spectrogram, follow ing AudioMAE (see. e.g., "Masked autoencoders that listen” by Huang et al.. 2022), after downsampling the audio waveform to 16000 kHz. the system can extract Melspectrograms with a duration of, e.g., 8 seconds, producing 128 feature vectors with 128 dimensions each (or more or fewer feature vectors and / or dimensions. Applying a patch kernel size of 16x16 would produce 64 total patches as input, although a different size could be employed. For wavefonn, a kernel size of 256samples and an embedding of up to 256 tokens may be employed, although more or fewer samples and tokens can be used.

[0083] We use separate learned positional embeddings for the linear sequence of text tokens and audio patches similar to the video encoding scheme above.Unified Encoder with Sparse Mixture-of-Experts (MoE) Lavers

[0084] One design decision important to multimodal modeling is how to allocate parameters to each modality. An asymmetric modality-specific design can be more optimal than using a similar-sized model for each modality. However, can come at the cost of requiring additional hyperparameter tuning to find the optimal parameterization. In experimentation, it has been observed that through the use of model sparsification with MoE, a unified encoder design coupled with certain modality-specific pre- and post-encoder layers is more optimal than a traditional multi-encoder setup as seen in CLIP model variants.

[0085] A standard Transformer architecture can be used as a shared encoder as in the VATT model.Once the modality -specific embeddings are extracted, all embedded inputs can be passed, regardless of modality, as-is through the shared encoder. One important finding as shown in experimental results is that MoE can bridge the gap between modality -agnostic and modality-specific designs. Instead of using separate encoders or even separate experts for each modality, MoE routing functions can be applied on all input tokens. Hence, any token can be routed to any expert regardless of their modality. This can be seen as an inductive bias, allowing each expert to be allocated to multiple modalities if the optimization benefits. One immediate benefit is that the addition of new modalities for fine-tuning does not need any specific changes to the encoder, unlike modality-specific experts which require additional modifications and input handling.

[0086] Applying a fixed set of feedforward networks as the experts in the encoder can be done, e.g., according to V-MoE (see, e.g., “Scaling vision with sparse mixture of experts” by Riquelme et al., 2021) or LIMoE (see, e.g., “Multimodal contrastive learning with limoe: the language-image mixture of experts” by Mustafa et al., 2022). For all MoE encoders, expert-choice routing can be used (see. e.g., “Mixture-of-experts with expert choice routing” by Zhou et al., 2022), which provides a strong baseline for all of the four modalities. Using expert-choice (top-c) routing, a much higher accuracy has been observed compared to the standard tokens-choose (top- / c) routing. This is because experts-choose routing guarantees even load balancing, which has been found to be an important factor for using an encoder shared across modalities. Moreover, it has been found that only applying MoE to the last 50% of layers provided similar accuracy to applying them for all layers.

[0087] Table 1 in Fig. 5 provides a description of model sizes for various IMP architectures. Results are presented for three main variants, IMP-S, IMP-B. and IMP-L corresponding to encoder sizes of ViT-S. ViT-B, and ViT-L respectively. Note that here MoE was only applied to the last half of thelayers in the encoder. Three additional sparse MoE sub-variants were also evaluated, which are indicated as IMP-MoE in subsequent tables noted herein.Output Heads

[0088] For the encoder heads module 106, modality-specific heads can be applied on the encoder representations to produce the final outputs for loss and prediction calculations. Here, a global average pooling operation can also be applied across the entire output sequence of the encoder, in which the resulting vector is used as the global features for classification and noise contrastive estimation objectives. For classification objectives, a dataset-specific linear classifier can be applied to the average- pooled outputs. For noise-contrastive estimation (NCE), the CLIP architecture can be followed, applying separate feedforward heads for each modality -to-common-space projection. Each feedforward head may comprise a two-layer linear projection with GeLU activation in between. The projection dimension size may be the same as the model’s hidden size.Multi-Resolution Training

[0089] One concern when training vision transformers on video data is that computation and memory efficiency are usually bottlenecked because of the Transformer’s quadratic complexity as a function of the input length. For example, having a 1 -second video sampled at 32 fps with a temporal patch size of 8 would result in 4x longer sequence which inherently results in a 16x increase in step time and memory consumption compared to a single-frame still image. To counteract this, the system may adjust batch size or resolution to compensate the additional temporal tokens, hence achieving a similar total number of input tokens compared to a single-frame still image.

[0090] According to another aspect of the technology, a multi-resolution encoding technique can be used to train on video data as efficiently as image data without any changes to the model architecture. As noted above, one can define multiple variants of each dataset with different I / O signatures. In the case of multi-resolution encoding, each dataset can provide a different number of spatiotemporal patches or batch size, in effect creating a different augmentation or view of each dataset. Under AGD, each unique input can be sampled uniformly from its corresponding variants, which may result in compiling a unique graph per I / O variant.

[0091] To accomplish this, the system may first fix a set tokens per batch T = B xFXTHXTW, which is factorized by the batch size B. frame tokens TP, height tokens TH, and width tokens Tw representing each patchified video. It is observed that one can further factorize each batch by trading off different dimensions such that the total number of input tokens per step are roughly equal so that peak memory usage is preserved. For example, the system can halve the spatial resolution while quadrupling the number of frames. This can increase convergence especially at the start of training, and provide a more memory efficient encoding of each objective.

[0092] Furthermore, vision patches, especially regarding video, contain large amounts of information redundancy and therefore can be dropped without loss of accuracy. Therefore, the systemcan leverage the DropToken approach (see, e.g., “Vatt: Transformers for multimodal self-supervised learning from raw video, audio and text” by Akbari et al., 2021) as an additional method to reduce tokens per batch by randomly dropping a fixed ratio of tokens per example. DropToken enables the system to randomly drop a portion of visual or audio tokens. In practice, this can cut training time in half with minimal accuracy loss. For TF temporal frame tokens, the system can randomly drop a ratio of 1 — - tokens per example to match the same tokens per batch as images.

[0093] In general, for certain objectives it has been found that a different mix of trade-offs may be more optimal. For example, contrastive objectives favor large batch sizes (e.g., 65536), so the system can reduce the resolution or apply DropToken to be more memory efficient. On the other hand, classification objectives do not need as large batch sizes for optimal convergence (e.g., 8192). so the system reduce the batch size while increasing the spatiotemporal tokens.

[0094] Furthermore, to be able to handle different numbers of patches across different dimensions, the positional encoding of vision modalities may need to be addressed carefully. Unlike the 1 - dimensional sequences of text and audio which can be truncated to a given length, the presence of 2D spatial dimensions mean that images with double the patches along a dimension could be subdivided into quadrants so that adjacent positions are close to each other in the embedding space. This can be accomplished using a dilated positional encoding. For a given dimension a spatial positional encoding of B buckets, if we encode a resolution with P patches, we dilate the positional encoding with a stride of- In the case of the temporal dimension in video, this can be treated the same as a 1 -dimensional truncation independent of the spatial dimensions, which is applied in the same way for text.Experiments and Test Results

[0095] This section describes the details of an experimental setup. The experiments may be organized into two parts, the first part indicating the downstream evaluation results of a final model, and the second part detailing the ablations on a smaller model variant which can motivate the design decisions of the model.

[0096] The datasets evaluated in the testing encompass a diverse set of learnable signals across multiple modalities: (1) images with text captions for image-text contrastive learning, (2) images with labels for supervised learning, (3) videos with audio and text captions for video-audio-text contrastive learning, and (4) videos with labels for supervised learning. Combmed, the data represents a set of 10 billion image-text pairs and 1 billion video-text-audio triplets. For large-scale pretraining, the following datasets were used:1. WebLI, having on the order of 4 billion English-only image-text pairs. This dataset was used for both image-text contrastive loss as well as image captioning loss.2. JFT-3B, which contains a large collection of multi-class labels per image. Here, BASIC was followed for encoding multiclass indices as text and use the dataset for image-text contrastive loss as well as supervised classification loss.3. LA10N-400M is a public dataset of 400 million image-text pairs for image-text contrastive loss.4. Wikipedia Image Text (WIT) has 37 millin image-text pairs sourced from Wikipedia for imagetext contrastive loss.5. Conceptual Captions (CC) 12M was used for image-text contrastive loss.6. ImageNet21K (I21K) has on the order of 11 million labeled images for image-text contrastive loss and supervised classification loss.7. VideoCC (VCC). a video dataset with a variant expanded to about 1 billion English video-text pairs for video-audio-text triplet contrastive loss and video captioning loss.8. HowTolOOM (HT100M) was used for video-audio-text triplet contrastive loss.9. Weak Text Supervision (WTS-70M) is a dataset on the order of 70 million video-text pairs for video-audio-text triplet contrastive loss.10. AudioSet was used for video-audio-text triplet contrastive loss focusing on audio understanding.

[0097] According to one aspect of the testing, a proportionally weighted sampling algorithm was used, executing each task in succession. To ensure that datasets were evenly sampled, each task was weighted by the number of examples, normalized to a probability distribution. For each dataset variant with different resolution sizes, the same weight was applied. For a fair evaluation on downstream tasks, all near-domain examples were filtered from the pretraining datasets (about 5 million examples total).

[0098] In the experiments, the input parameters were reconfigured so that the number of frame tokens was always equal to 4. This was done to result in the base tokens per video batch being exactly 4x of images. For video datasets, three variants were constructed as follows:1. Reduce the resolution by half in each dimension:TH TryB x TFx — x — = B x THx Tw, subject to TF— 42. Reduce the batch size by 4x so that:3. Apply DropToken d = 1 — — = 0.75.TpWhere in (1) and (2) Tvtlieois on the left side of the equation, and Tlmageis on the right side.

[0099] During training, the system uniformly sampled from each variant. It was observed that this provides significantly faster training times on video data (see Table 2 of Fig. 6) while simultaneously avoiding excessive memory consumption, as tokens per batch are roughly equivalent to those of the image datasets. Table 2 illustrates zero-shot results of IMP on Video Action Recognition (top-1). IMPis shown to achieve a new state-of-the-art by a wide margin with relatively low training cost. Train compute is measured in TPU-days

[0100] For image datasets, a similar strategy was applied but for the purpose of high-resolution learning. The three variants here are:1. The base resolution with Timage= B x THx Tw2. Reduce the batch size by 4x and double each spatial dimension so that:3. Apply DropToken d = 1 — = 0.75.

[0100] For certain experiments, a patch size of 4x16x16 was used on base input resolutions of 16x256x256 and 4x256x256 on video and image modalities respectively, resulting in a total of 1024 and 256 patches per sample. The text inputs in ImageNet21K and JFT were truncated to 16 tokens to improve step efficiency with no loss of information, while keeping the text length of the rest of the datasets to a maximum of 256 tokens. A base batch size of 65536 was used for training using the Adam optimizer, a peak learning rate of le-3 with a cosine schedule, and applying no weight decay. For MoE parameters, experts-choose routing was applied with a top-c capacity factor of 1.0 and without applying any jittering to the routing or other auxiliary losses.

[0101] During inference, the system evaluated on the largest available resolution that the model was trained on, here 16x512x512, using a total of 8 clips per video at approximately 12.5 fps on all evaluated datasets. Training resulted in roughly 16 billion examples seen, or about 5 trillion tokens. Taken together, these datasets represent about 11 billion unique image and video examples.

[0102] For certain testing, the model was scaled up and IMP was evaluated on a variety of downstream datasets. For zero-shot evaluation, denote text (T), vision (V). and audio (A) modalities with their intended mapping (e.g., V — > T for video-text retrieval). For linear evaluation, all parameters were frozen and a linear layer was applied on top of the pooled features, similar to classification tasks in pretraining.

[0103] In Tables 2 (Fig. 6) and 3 (Fig. 7), the zero-shot capabilities of the model were compared on several classification datasets. As noted above. Table 2 illustrates zero-shot results of IMP on Video Action Recognition (top-1). Table 3 illustrates retrieval results of IMP on image datasets (top-1). Compared to VideoCoCa, IMP -MoE -L was trained on 256 TPU v4 chips for 6 days, representing 15% of the total training cost of VideoCoCa. Strong performance is noted on video datasets especially, with a substantial lead on Kinetics (see columns K400, K600 and K700), UCF101, and HMDB-51 top-1 accuracy (Table 2). On ImageNet (see Table 3), IMP shows higher zero-shot metrics with respect to CoCa, assuming similar model parameters as IMP-MoE-B. Similar accuracy was observed when scaling to IMP-MoE-L, however most gains are observed on video benchmarks.

[0104] In Table 4 (Fig. 8), results were compared on the fine-tuned model on UCF-101. And in Table 5 (Fig. 9) the results of IMP on audio were evaluated with the ESC-50 dataset. It is shown that IMP-MoE-L achieves state-of-the-art when fine-tuned on UCF-101. Note that in terms of high-quality audio-text examples, AudioSet was used for testing, but despite this IMP remains competitive on ESC- 50 with AudioCLIP, an audio-only baseline.

[0105] The testing also included various ablation studies. Here, the experiments used IMP-S trained for 250k steps with a base batch size of 8192. A fixed video / image resolution of 16x224x224 / 4x224x224 using a patch size of 4x16x16 was employed. Unless otherwise specified, multi-scale resolution was not applied.

[0106] This testing confirmed that combined objectives are mutually beneficial. The ablation process included evaluating whether a contrastive objective (NCE) or classification objective (softmax) is more important for downstream evaluation on vision tasks. Using ImageNet21K as the pretraining dataset, the following was done: train on the objectives separately, combine the objectives by summing them, or alternating (AGD) between the objectives on each step. In the case of alternating, for a fair comparison so that training time was equivalent, the same number of steps (250k) was fixed so that each objective only optimizes 50% of the total steps (125k). The system was evaluated on ImageNet and CIFAR-100 by training a linear classifier and on imagc-to-tcxt retrieval.

[0107] The results in Table 6 (Fig. 10) indicate several important insights. It is not a surprise that classification objective benefits fine-tuning evals the most, while contrastive objective benefits open vocabulary classification. However, it is observed that combining both objectives is better than optimizing on them individually. And alternating between the objectives is better than non-AGD objective mixing. Thus, according to one aspect of the technology, the system may fix one objective per training step and alternate optimization between them.

[0108] Multi-task multi-dataset AGD is also mutually beneficial. The previous results of AGD may be extended by adding additional datasets to the training mixture. Chart 1100 of Fig. 11 compares the result of adding additional datasets to the pretraining mixture. This chart additionally compares results across Flickr30k and COCO datasets. The evaluation started with Conceptual Captions (CC) 12M dataset and gradually added new datasets and objectives. Most notably, the addition of ImageNet21K (i21k) dataset is also compared, showing complementary improvement when combining NCE and Softmax objectives from classification-based pretraining objectives. Similar to i21k isolated experiments, adding softmax objective benefits the entire pretraining mixture. Adding softmax objectives can be better for zero-shot evaluation than adding more contrastive data. The reverse can also be true: adding contrastive objectives can be better for linear evaluation in the case of CIFAR-100 and Flickr30k evaluation. Certain dataset combinations (CC+i21k, CC+LAION) can be seen to cause instability at the beginning of training. Adding a classification objective had a stabilizing effect.significantly reducing the chance of slow convergence. Optimizing on LAION directly was difficult, but benefited training a lot more when mixed in with other datasets.

[0109] In some situations, the addition of multiple datasets, even under the same objective, can be detrimental to the optimization process. For instance, at the beginning of training CC and LAION datasets, a loss plateau was observed, where the loss remained relatively constant from the start of training, and the modeled fails to start converging for a long period of time. On the other hand, when the same dataset was applied and a softmax objective was added from the ImageNet21K dataset, a loss plateau was no longer observed, as softmax tends to be more stable for optimization processes in general. And integrating both objectives on ImageNet21 K. the plateau was also missing. This highlights the importance of selecting the right dataset mixture, especially at the start of training where the inherent nature of the random parameters can make it difficult for some task gradients to solidify a good direction in the optimization process. Thus, according to one aspect of the technology, the system may further integrate a larger set of diverse datasets and objectives.

[0110] It was also found that multi-scale resolution provides universal improvement. Chart 1200 in Fig. 12 shows a comparison of using different combinations of resolution, batch size, and DropToken as input. In particular, this chart presents a comparison of dynamic input sizes on IMP-S. For each addition, die baseline configuration was copied and the input sizes modified accordingly, and each new dataset configuration was sampled equally. In all settings, input sizes were fixed so that the total tokens per batch was roughly equivalent, and it was ensured that all training runs used the same number of total steps (more dataset variants mean a smaller number of steps executed per dataset). It can be seen that certain ty pes of datasets respond well to droptoken while others may not. Conceptual Captions with double the batch size and droptoken 0.5 improved ImageNet / CIFAR zero-shot. Droptoken + 320x320 image on i21k softmax pretrain was better for linear evals and Flickr30k. Adding multiple versions of smaller batch size + higher res, droptoken + higher res. larger batch size + lower res. can significantly improve downstream evaluations It was formd that dynamic mixtures of resolution, batch size, and droptoken were always helpful.

[0111] Testing also showed that adding more modalities may impact single tower (dense) encoder accuracy. Chart 1300 of Fig. 13 presents a comparison of the addition of video and audio datasets. The addition of video data (here, WTS) can be seen to substantially improve video accuracy (Kinetics400, UCF) at the cost of slightly reducing the model’s accuracy on image tasks (ImageNet, Flickr30k, COCO). The introduction of audio in the contrastive loss (here. AudioSet) also reduced both image and video accuracy slightly, but enabled fine-tuning on audio data (here, ESC). Finally, the introduction of a dedicated audio class contrastive objective reduced image and video accuracy the most, but enabled zero-shot audio classification.

[0112] In particular, this testing compared the addition of more modalities via video datasets. Adding a video dataset (WTS), to pretraining boosted the Kinetics classification significantly, allowingthe model to more easily discriminate between action classes. However, the addition of video data impacted image classification performance slightly, especially when parameters were constrained. Likewise, the addition of audio data with video had a slight negative impact, and the addition of dedicated audio classification dataset (AudioSet) had an even larger negative impact. This may be due to the additional audio-text contrastive signal which required the network to allocate dedicated processing for audio-to-text understanding. Therefore, according to one aspect of the technology, a standard single tower encoder may not be sufficient for optimal multimodal learning due to parameter bottlenecks.

[0113] The ablation testing also confirms that MoE provides universal improvement across modalities, and that it also resolves the single tower encoder parameter bottleneck. Compared to a two- tower contrastive model, the encoder of a unified image-text model contains half the parameters, while keeping training efficiency the same. From the MoE comparisons in chart 1400 of Fig. 14, it is observed that simply replacing a model with an equivalent with just four experts, the system can provide a large gain in accuracy, especially for zero-shot metrics. Even larger improvements were made when scaling up the model size and increase the experts to 16. This indicates that MoEs can be used to bridge the multimodal gap.

[0114] Moreover, it is observed that with the addition of MoE, the system can significantly close the gap between multiple modalities as seen in chart 1500 of Fig. 15, which provides an MoE comparison with multimodal inputs. These results show reduced image accuracy on the dense model when joint-training on images and video. However, switching to MoE not only closes the evaluation gap, but improves both image and video metrics simultaneously, showing mutual generalization on both vision modalities. The addition of audio somewhat reduces accuracy on image and video metrics across the board, but is much less prominent when using MoE.

[0115] Since experts are free to choose which tokens are allocated to different experts, strong alignment is observed between experts and modalities. Note that experts-choose routing is particularly important for strong single tower performance. Of all the variants tested, a unified MoE encoder using experts-choose routing provided the most parameter and compute efficient design. This is opposed to a multi-encoder modality -specific model, as seen in chart 1600 of Fig. 16, which compares single tower and multi-tower architectures on IMP-S.

[0116] Here, two variants of two tower models were applied, one with similar parameters to single tower, and another with similar compute. It is observed that multi-tower models are more compute efficient than single tower models, but less parameter efficient. However, a single tower MoE model with experts-choose routing is both more compute and parameter efficient than all variants, showing improved generalization and using fewer parameters with the same compute budget as a multi-tower dense model. When comparing two-tower models, the system can either split the parameters to be roughly equal in size to a single tower, or duplicate the towers to double the parameter count whileproviding equivalent computation. In either case, results show superior parameter efficiency and higher accuracy when using just 4 experts.

[0117] The effect of experts-choose routing vs. tokens-choose are compared in Table 7 of Fig. 17. It is observed that separating experts by modality in the case of tokens-choose routing is useful for improving accuracy. However, when switching to experts-choose routing, it is found that performance increases further, and a multi-tower model is similar enough in accuracy that separating them per modality is no longer necessary. This allows for a much simpler model design, and the system can finetune new modalities in the encoder without any additional setup or new experts required.

[0118] Additional testing was also conducted on the IMP technology as follows. Table 8 of Fig. 18 is a prompt comparison of IMP, in which three settings of train-time prompts were compared with increasing diversity. Results show that randomized prompts in training tend to significantly improve metrics on classification tasks. In particular, this testing evaluated different prompt settings during pretraining. IMP was contrasted with prior works which only apply prompts during evaluation. It is observed that there were strong gains when randomizing the prompts during training. A trend can be seen that simpler prompts may be more useful for datasets with a smaller number of classes. Thus, in contrast, when dealing with large-scale datasets, prompt diversification may be used as a default setting.

[0119] And Tabic 9 of Fig. 19 provides a comparison of datasets & objectives with AGD. Here, Conceptual Captions (CC) was integrated with contrastive (NCE) loss, and ImageNet21K (I21K) with NCE and softmax loss. The addition of both NCE and Softmax objectives from classification-based pretraining datasets were mutually beneficial with the retrieval-based pretraining datasets, observing best performance with tire combination of both objectives. Further optimality was provided by adding larger, more diverse dataset like LAION-400M. However, it was found that LAION caused optimization on contrastive objectives to become unstable, so softmax loss can be used to stabilize this type of noisier dataset. In general, the following situations were observed where optimization can become more unstable: (1) Increase in dataset diversity; (2) Increase in model size; (3) Increase in batch size. In Table 9, it is shown that training on NCE alone can cause instability issues during training, especially for noisy text datasets. But with the addition of more clean data sources and softmax objectives, one can greatly reduce the instability.IMP Framework

[0120] As explained herein, IMP is a multimodal sequence-to-sequence approach that may utilize a Transformer-type encoder. It is able to take images, videos, audio and text as inputs to the encoder and generate their feature embeddings as outputs. By way of example, the inputs may comprise RGB images. RGB video frames, audio waveform or spectrogram, as well as text in one or more languages. The IMP model can be trained from scratch with random initialization. In one scenario, it may be a static model trained on offline datasets. IMP can be used for image classification, video classification, audio classification, image to text or text to image retrieval, as well as other applications.

[0121] IMP provides AGD training. Individual tasks may be confined toa single training step. New tasks do not affect the efficiency of existing tasks. This training also results in guaranteed convergence with marginal convexity with respect to model parameters and tasks. IMP also provides a unified encoder-decoder model for any combination of images, video, audio and / or text. The unified model is able to perform classification and zero-shot retrieval. According to one aspect of the technology, text and / or video may be generated from the model. As shown via the test results, integration and performance may improve substantially when multimodal MoE is utilized.

[0122] The system also supports dynamic input sizes, which provides a tradeoff of large batch size versus small sequence length (resolution) across training steps. This allows scaling to large batch sizes (e.g., on the order of 65000 or larger) or resolution (e.g., 10242or larger) from the start of pretraining. Dynamic input sizes also allow the system to easily support curriculum training, such as multigrid training. Moreover, the system can vary input resolution during inference.

[0123] In view of the above discussion, an IMP framework has been developed that employs AGC, which is referred to herein as MAX. MAX is a multimodal framework providing an integrated multimodal training pipeline supporting large-scale image, video, text, and audio pretraining, which can be implemented in JAX or another compiler application that is able provide just-in-time operation. MAX provides an cnd-to-cnd framework for running arbitrary multimodal data on models efficiently. It supports full horizontal scaling of multimodal tasks-multiple tasks and objectives at no additional cost. It also supports dynamic input sizes that allow arbitrary' batch sizes, token lengths, and resolutions without padding. An overview of modules that may be used in MAX is shown in group 2000 of Fig. 20, when JAX is utilized.

[0124] In this particular configuration, the data pipeline defines data using TensorFlow Datasets (TFDS) and SeqIO registries for vision and language tasks. Preprocessing of text is provided by SeqIO, while image, video, and audio preprocessors are provided by DeepMind Video Readers (DMVR). Datasets are emitted from a tf.data.Dataset object provide a key-value signature that can be tightly- integrated with models. This signature should be consistent with the model’s expected input structure. For IMP, named keys can be defined for each modality and emit the applicable modalities from each dataset. Each modality can further provide optional metadata, information that specify how to properly execute or route the input to various modules. Metadata can include the names of the component each modality is routed to. For example, for image captioning, “text” could be routed to a decoder, while “image” could be routed to an encoder. For supervised classification, metadata can include a text field that identifies which label set to apply when computing the classification head and loss, e.g.. “imagenet classification”. Metadata can also include how to combine modalities. For instance, metadata can contain a reference to a pair of modalities which are then concatenated before passing them through the encoder.

[0125] Models can be built as native Flax modules, partitioned jax.pjit, and optimized by transforms defined in Optax. Jax.pjit is a method to compile and distribute a computation graph, and it is used with partitioning annotations (see SPMD) to create high-performance device meshes. SPMD: jax spmd is a technique jax uses to run the same computation on different pieces of data. With pjit, it allows jax to send annotation hints (e.g.. jax.pjit. with sharding constraint) on how to distribute tensors across devices for efficient parallel computation and model parallelism.

[0126] The JAX framework provides a core selection of primitives that interface with XLA, a library that compiles and optimizes computation graphs across different devices. TensorStore may be used to efficiently checkpoint and restore partitioned model parameters using async parallel dispatch. In this example, configmation may be specified according to Python dataclasses which can be overridden. This allows the creation of many variants of datasets, models, and experiments without excessive code duplication.

[0127] On each training step, the training loop samples a dataset-objective pair, passing inputs from the dataset directly into the model. Note that the routing of inputs across different model components can be specifically avoided in the training loop logic to prevent tire training process from being tied to a specific way to handle different input ty pes. Instead, the model itself may handle the interpretation of any combination inputs provided from the dataset and produces a named collection of outputs. Loss functions are applied in the training loop after sampling a dataset-objective pair and executing the model's forward pass. Together, this provides a modular way to interchange datasets, models, and loss functions.

[0128] Training and inference step partitioning may be leveraged from jax.pjit, with further model and data parallelism abstractions provided by the T5x framework to partition model weights and activations across devices. On a high level. PJIT enables the use of dynamic graph compilation at runtime across many distributed devices. For each unique dataset-objective pair, PJIT will compile a new computation graph. These graphs may all be cached so on subsequent iterations re-compilation overhead is minimized. This is used in conjunction with MoE to efficiently dispatch sparse weights across multiple devices while minimizing communication overhead. In conjunction with the partitioner, we initialize states by defining a set of specifications of shapes that the model should accept as input, using jax.eval shape. To efficiently nm each training step, the system may pre-initialize the PRNG states of all training steps before any training takes place. In one scenario, the system can use t5x as a reference. The difference is that instead of applying a single type of input for training, a set of tasks (dataset-objective pairs) can be defined before the main training loop, and then iterate over all the tasks, emitting a potentially different task at each step, jax.jit would handle the rest, re-compiling the computation graph if encountering a new input / output shape it has not yet seen.Example Computing Architecture

[0129] The integrated multimodal perception technology discussed herein may be trained on one or more tensor processing units (TPUs). graphics processing units (GPUs), CPUs or other computing devices in accordance with the features disclosed herein. One example computing architecture is shown in Figs. 21(a) and 21(b). In particular, Figs. 21(a) and (b) are pictorial and functional diagrams, respectively, of an example system 2100 that includes a plurality of computing devices and databases connected via a network. For instance, computing device(s) 2102 may be implemented as a cloudbased server system. Databases 2104, 2106 and 2108 may store, e.g., training inputs (e.g.. video segments or clips, or full videos, a text corpus and / or an audio corpus, etc.), classification and / or retrieval results other output according to the model(s), and / or trained models, respectively. The server system may access the databases via network 2110. Client devices may include one or more of a desktop computer 2112. a laptop or tablet PC 2114, a mobile phone or PDA 2116, a wearable device 2118 such as a smartwatch or head-mounted display (e.g., a virtual reality headset), in-home devices such as smart display 2120a and / or a smart home device 2120b, etc.

[0130] As shown in Fig. 21(b), each of the computing devices 2102 and 2112-2120 may include one or more processors, memory , data and instructions. The memory’ stores information accessible by the one or more processors, including instructions and data (e.g., models) that may be executed or otherwise used by the proccssor(s). The memory’ may be of any ty pe capable of storing information accessible by the processor(s), including a computing device-readable medium. The memory is a non- transitory medium such as a hard-drive, memory card, optical disk, solid-state, etc. Systems may include different combinations of the foregoing, whereby different portions of the instructions and data are stored on different types of media. The instructions may be any set of instructions to be executed directly (such as machine code) or indirectly (such as scripts) by the processor(s). For example, the instructions may be stored as computing device code on the computing device-readable medium. In that regard, the terms “instructions”, “modules” and “programs” may be used interchangeably herein. The instructions may be stored in object code format for direct processing by the processor, or in any other computing device language including scripts or collections of independent source code modules that are interpreted on demand or compiled in advance.

[0131] The processors may be any conventional processors, such as commercially available CPUs, TPUs. GPUs, etc. Alternatively, each processor may be a dedicated device such as an ASIC or other hardware-based processor. Although Fig. 21 (b) functionally illustrates the processors, memory, and other elements of a given computing device as being within the same block, such devices may actually include multiple processors, computing devices, or memories that may or may not be stored within the same physical housing. Similarly, the memory may be a hard drive or other storage media located in a housing different from that of the processor(s), for instance in a cloud computing system of server 2102. Accordingly, references to a processor or computing device will be understood to include references to a collection of processors or computing devices or memories that may or may not operate in parallel.

[0132] The computing devices may include all of the components normally used in connection with a computing device such as the processor and memory described above as well as a user interface subsystem for receiving input from a user and presenting information to the user (e.g., text, imagery, videos and / or other graphical elements). The user interface subsystem may include one or more user inputs (e.g., at least one front (user) facing camera, a mouse, keyboard, touch screen and / or microphone) and one or more display devices (e.g., a monitor having a screen or any other electrical device that is operable to display information (e.g., text, imagery and / or other graphical elements). Other output devices, such as speaker(s) may also provide information to users.

[0133] The user-related computing devices (e.g.. 21 12-2120) may communicate with a back-end computing system (e.g.. server 2102) via one or more networks, such as network 2110. The network 2110, and intervening nodes, may include various configurations and protocols including short range communication protocols such as Bluetooth™, Bluetooth LE™, the Internet, World Wide Web, intranets, virtual private networks, wide area networks, local networks, private networks using communication protocols proprietary to one or more companies, Ethernet, WiFi and HTTP, and various combinations of the foregoing. Such communication may be facilitated by any device capable of transmitting data to and from other computing devices, such as modems and wireless interfaces.

[0134] In one example, computing device 2102 may include one or more server computing devices having a plurality of computing devices, e.g., a load balanced server farm or cloud computing system, that exchange information with different nodes of a netw ork for the purpose of receiving, processing and transmitting the data to and from other computing devices. For instance, computing device 2102 may include one or more server computing devices that are capable of communicating with any of the computing devices 2112-2120 via the network 2110.Example process

[0135] Fig. 22 illustrates an example flow diagram 2200 for multimodal multi-task training of an integrated multimodal perception model in accordance with aspects of the technology. This includes, at block 2202, obtaining, by one or more processors of a computing system, a plurality of datasets, each dataset having a different modality, each dataset also having an input resolution. At block 2204, the process also includes obtaining, by the one or more processors, a set of loss objectives associated with the plurality of datasets. Then at block 2206 the process includes performing, by a neural network implemented via the one or more processors, embedding of the different modalities into a shared modality-agnostic space to obtain a set of embedded tokens, and at block 2208 determining, by the neural network according to a mixture of experts encoder module, semantic contextual embeddings from the embedded tokens. Then at block 2210 the process includes generating, by the neural network according to a set of encoder heads, a set of final predictions by re-projecting the semantic contextual embeddings into a modality -specific space according to the set of loss objectives.

[0136] Finally, as shown and discussed herein. IMP is able to achieve competitive performance on a wide range of downstream tasks including image classification, video classification, image-text, and video-text retrieval. Most notably, a sparse IMP-MoE-L has been trained to focus on video tasks that achieves new state-of-the-art in zero-shot video classification. Testing has shown that the model achieved 77.0% on Kinetics-400, 76.8% on Kinetics-600, and 76.8% on Kinetics-700 zero-shot classification accuracy, improving upon known several existing approaches by +5%, +6.7%, and +5.8%. respectively, while using only 15% of their total training computational cost. Thus, there are substantial technical benefits to the IMP approach.

[0137] Although the technology herein has been described with reference to particular embodiments, it is to be understood that these embodiments are merely illustrative of the principles and applications of the present technology. It is therefore to be understood that numerous modifications may be made to the illustrative embodiments and that other arrangements may be devised without departing from the spirit and scope of the present technology as defined by the appended claims.

Claims

CLAIMS1. A computer-implemented method for multimodal multi-task training of an integrated multimodal perception model, the method comprising: obtaining, by one or more processors of a computing system, a plurality of datasets, each dataset having a different modality, each dataset also having an input resolution; obtaining, by the one or more processors, a set of loss objectives associated with the plurality’ of datasets; performing, by a neural network implemented via the one or more processors, embedding of the different modalities into a shared modality-agnostic space to obtain a set of embedded tokens; determining, by the neural network according to a mixture of experts encoder module, semantic contextual embeddings from the embedded tokens; and generating, by the neural network according to a set of encoder heads, a set of final predictions by re-projecting the semantic contextual embeddings into a modality-specific space according to the set of loss objectives.

2. The method of claim 1, wherein the different modalities of the plurality of datasets are selected from the set consisting of imagery, video, text, and audio modalities.

3. The method of claim 1, wherein the input resolution of each dataset is distinct from the resolution of the other datasets.

4. The method of claim 1, wherein the set of loss objectives is selected from the set consisting of image-text contrastive, video-text contrastive, video-audio contrastive, image classification, video classification, and audio classification.

5. The method of claim 1, wherein the generating is performed according to an alternating gradient descent process.

6. The method of claim 5. wherein at each gradient step, a different loss objective is optimized given one or both of different sets of model weights or input modalities.

7. The method of claim 5, wherein the set of final predictions includes at least one of softmax probabilities, aggregation heads for noise contrastive estimation, or per-sequence heads for language modeling.

8. The method of claim 5. wherein a set of tasks to be performed by the alternating gradientdescent process is performable in any order.

9. The method of claim 8, wherein the set of tasks to be performed includes one or more of supervised classification, masked auto encoding, image-text contrastive, and video-audio-text contrastive tasks.

10. The method of claim 1. wherein the mixture of experts encoder module is a sparse mixture of experts encoder module.

11. The method of claim 1, further comprising defining multiple variants of each dataset, each variant of a given dataset corresponding to a different number of spatiotemporal patches or a different batch size.

12. The method of claim 1, further comprising performing inference using trained model to generate one or more results.

13. The method of claim 1, wherein performing the embedding further includes performing flatting of the different modalities.

14. A system for multimodal multi-task training of an integrated multimodal perception model, tire system comprising: memory configured to store a plurality of datasets, each dataset having a different modality, each dataset also having an input resolution; and one or more processors operatively coupled to the memory, the one or more processors being configured to: obtain the plurality of datasets from the memory; obtain a set of loss objectives associated with the plurality of datasets; perform, by a neural network implemented via the one or more processors, embedding of the different modalities into a shared modality -agnostic space to obtain a set of embedded tokens; determine, according to a mixture of experts encoder module, semantic contextual embeddings from the embedded tokens; and generate, according to a set of encoder heads, a set of final predictions by re-projecting the semantic contextual embeddings into a modality-specific space according to the set of loss objectives.

15. The system of claim 14, wherein generation of the set of final predictions is performed according to an alternating gradient descent process.

16. The system of claim 15, wherein at each gradient step, a different loss objective is optimized given one or both of different sets of model weights or input modalities.

17. The system of claim 15, wherein the set of final predictions includes at least one of softmax probabilities, aggregation heads for noise contrastive estimation, or per-sequence heads for language modeling.

18. The system of claim 1 , wherein a set of tasks to be performed by the alternating gradient descent process is performable in any order.

19. The system of claim 14, wherein the mixture of experts encoder module is a sparse mixture of experts encoder module.

20. The system of claim 14, wherein the one or more processors are further configured to define multiple variants of each dataset, each variant of a given dataset corresponding to a different number of spatiotemporal patches or a different batch size.

21. The system of claim 14, wherein the one or more processors are further configmed to perform inference using trained model to generate one or more results.

Images