Inverted bottleneck architecture search and efficient attention mechanism for machine-learned models

Abstract

Generally, the present disclosure is directed to efficient neural network architectures, techniques for constructing new efficient neural networks, and approaches to low-latency execution of neural networks. In an example aspect, the present disclosure provides a more powerful configuration of an inverted bottleneck block that maintains an efficient execution profile. An example search system can parameterize a universal inverted bottleneck block with a first parameter that categorically activates a spatial mixing operation in the unexpanded state. In this manner, for instance, a large variety of different network architectures can be explored using a relatively compact search space that admits high levels of parameter sharing. Further, the present disclosure introduces an efficiency-optimized multi-query attention block.

Description

INVERTED BOTTLENECK ARCHITECTURE SEARCH AND EFFICIENT ATTENTION MECHANISM FOR MACHINE-LEARNED MODELS PRIORITY

[0001] This application claims priority to United States Provisional Patent Application No.63 / 550,970, filed February 7, 2024. United States Provisional Patent Application No.63 / 550,970 is hereby incorporated by reference herein in its entirety. FIELD

[0002] The present disclosure relates generally to machine learning processes and machine-learned devices and systems. More particularly, the present disclosure relates to efficient neural network architectures, techniques for constructing new efficient neural networks, and approaches to low-latency execution of neural networks. BACKGROUND

[0003] A computer can receive input(s). The computer can execute instructions to process the input(s) to generate output(s) using a parameterized model. The computer can obtain feedback on its performance in generating the outputs with the model. The computer can generate feedback by evaluating its performance. The computer can receive feedback from an external source. The computer can update parameters of the model based on the feedback to improve its performance. In this manner, the computer can iteratively “learn” to generate the desired outputs. The resulting model is often referred to as a machine-learned model. SUMMARY

[0004] Aspects and advantages of embodiments of the present disclosure will be set forth in part in the following description, or can be learned from the description, or can be learned through practice of the embodiments.

[0005] In an aspect, the present disclosure provides a first example computer- implemented method. The first example computer-implemented method can include receiving, at an input to the twice-mixed inverted bottleneck block, an input state of a neural network having a first channel depth. The first example computer-implemented method can include convolving, using a first depthwise convolutional layer of the twice-mixed invertedbottleneck block, the input state to obtain a spatially-mixed input state. The first example computer-implemented method can include expanding, using an up-projection pointwise convolutional layer of the twice-mixed inverted bottleneck block, the spatially-mixed input state to generate an expanded state having a second channel depth greater than the first channel depth. The first example computer-implemented method can include convolving, using a second depthwise convolutional layer of the twice-mixed inverted bottleneck block, the expanded state to obtain a spatially-mixed expanded state. The first example computer- implemented method can include projecting, using a down-projection pointwise convolutional layer of the twice-mixed inverted bottleneck block, the spatially-mixed expanded state to obtain a projected state having a third channel depth less than the second channel depth. The first example computer-implemented method can include using the projected state to perform a prediction task with the neural network. For example, in an implementation, the input state represents image data, and the prediction task comprises an image processing task.

[0006] In an aspect, the present disclosure provides a second example computer- implemented method. The second example computer-implemented method can include initializing a search space comprising a plurality of configurations of a universal inverted bottleneck block. The universal bottleneck block can include a first placeholder for a depthwise convolutional layer that is discretely activated during search using a first parameter. The universal bottleneck block can include an up-projection pointwise convolutional layer that follows the first placeholder and expands a channel depth of an internal state of the network to an expanded state. The universal bottleneck block can include a second placeholder for a depthwise convolutional layer that is discretely activated during search using a second parameter, the second placeholder following the expanding pointwise convolutional layer. The universal bottleneck block can include a down-projection pointwise convolutional layer that follows the second placeholder and decreases a channel depth of an internal state of the network from the expanded state to a projected state. In an example, the placeholder can perform a unit transform or passthrough if not activated by the first or second parameter. The second example computer-implemented method can include constructing a plurality of candidate models, wherein constructing the plurality of candidate models comprises, for each candidate model of the plurality of candidate models: selecting, for each respective network block of a plurality of network blocks of a neural network architecture, a respective value of the first parameter and a respective value of the second parameter to instantiate a respective inverted bottleneck block at the respective network block. The secondexample computer-implemented method can include training the plurality of candidate models over a training dataset. The second example computer-implemented method can include evaluating performance of the plurality of candidate models. The second example computer-implemented method can include outputting a selected candidate model based on the evaluating.

[0007] In an aspect, the present disclosure provides a third example computer- implemented method. The third example computer-implemented method can include obtaining a current input tensor representing an input sequence of one or more elements. The third example computer-implemented method can include computing attention values over the current input tensor using one or more tensor contractions, wherein at least a portion of respective tensor contractions of the one or more tensor contractions comprise an output tensor having non-contracting indices inherited from a first tensor and non-contracting indices inherited from a second tensor, wherein: the non-contracting indices inherited from the first tensor are contiguous and ordered as in the first tensor; and the non-contracting indices inherited from the second tensor are contiguous and ordered as in the second tensor. The third example computer-implemented method can include generating an output of the machine-learned sequence processing model based on the attention values.

[0008] In an aspect, the present disclosure provides a fourth example computer- implemented method. The fourth example computer-implemented method can include obtaining an input tensor representing at least one input sequence of one or more elements, the input tensor having a sequence length dimension and an embedding dimension. The fourth example computer-implemented method can include obtaining a query projection tensor configured to generate a projection of the input tensor using a plurality of learned weights, the query projection tensor having a query head dimension, an embedding dimension, and a projection dimension, wherein adjacent indices of the query projection tensor respectively correspond to the query head dimension and the projection dimension. The fourth example computer-implemented method can include contracting the input tensor and the query projection tensor in the embedding dimension. Contracting the input tensor and the query projection tensor in the embedding dimension can be performed by reshaping the query projection tensor into a query matrix by combining the query head dimension and the projection dimension into a single dimension. Contracting the input tensor and the query projection tensor in the embedding dimension can be performed by obtaining an input matrix comprising at least a portion of the input tensor. Contracting the input tensor and the query projection tensor in the embedding dimension can be performed by multiplying the inputmatrix and the query matrix to obtain an output. Contracting the input tensor and the query projection tensor in the embedding dimension can be performed by reshaping the output into a query tensor having a query head dimension, a sequence length dimension, and a projection dimension, wherein adjacent indices of the query tensor respectively correspond to the query head dimension and the projection dimension. The fourth example computer-implemented method can include using the query tensor to compute attention values over the at least one input sequence.

[0009] In an aspect, the present disclosure provides a fifth example computer- implemented method. The fifth example computer-implemented method can include obtaining a distillation dataset that comprises a plurality of training examples labeled by a teacher model, wherein the plurality of training examples comprises: a first plurality of training examples generated using a first set of augmentation routines; and a second plurality of training examples generated using a second set of augmentation routines different from the first set of augmentation routines. The fifth example computer-implemented method can include generating a training batch for training a student model by combining one or more of the first plurality of training examples and one or more of the second plurality of training examples. The fifth example computer-implemented method can include processing batchwise the training batch using the student model to obtain a batch of outputs. The fifth example computer-implemented method can include evaluating the batch of outputs using corresponding labels for the training batch generated by the teacher model. The fifth example computer-implemented method can include updating one or more learnable parameters of the student model based on the evaluating.

[0010] In an aspect, the present disclosure provides one or more non-transitory, computer-readable media storing instructions that are executable by one or more processors to cause a computing system to perform operations, the operations comprising one or more features of the first computer-implemented method, the second computer-implemented method, the third computer-implemented method, the fourth computer-implemented method, or the fifth computer-implemented method, or any combination of features thereof.

[0011] In an aspect, the present disclosure provides a computing system comprising one or more processors and one or more non-transitory, computer-readable media storing instructions that are executable by the one or more processors to cause the computing system to perform operations, the operations comprising one or more features of the first computer- implemented method, the second computer-implemented method, the third computer-implemented method, the fourth computer-implemented method, or the fifth computer- implemented method, or any combination of features thereof.

[0012] In an aspect, the present disclosure provides a neural network configured to implement one or more features of the first computer-implemented method, the second computer-implemented method, the third computer-implemented method, the fourth computer-implemented method, or the fifth computer-implemented method, or any combination of features thereof.

[0013] Other example aspects of the present disclosure are directed to other systems, methods, apparatuses, tangible non-transitory computer-readable media, and devices for performing functions described herein. These and other features, aspects, and advantages of various implementations will become better understood with reference to the following description and appended claims. The accompanying drawings, which are incorporated in and constitute a part of this specification, illustrate implementations of the present disclosure and, together with the description, help explain the related principles. BRIEF DESCRIPTION OF THE DRAWINGS

[0014] Figure 1 is a block diagram illustrating an example implementation of a bottleneck block according to example implementations of aspects of the present disclosure;

[0015] Figure 2A is a block diagram illustrating an example implementation of a bottleneck block search space according to example implementations of aspects of the present disclosure;

[0016] Figure 2B is a block diagram illustrating an example implementation of a bottleneck block search space according to example implementations of aspects of the present disclosure;

[0017] Figure 3 is a block diagram illustrating an example implementation of a bottleneck block search system according to example implementations of aspects of the present disclosure;

[0018] Figure 4 is a block diagram illustrating an example implementation of a transformer encoder according to example implementations of aspects of the present disclosure;

[0019] Figure 5 is a block diagram illustrating an example implementation of a transformer encoder according to example implementations of aspects of the present disclosure;

[0020] Figure 6 is a block diagram illustrating an example implementation of an attention block according to example implementations of aspects of the present disclosure;

[0021] Figure 7 is a block diagram illustrating generation of a distillation training dataset according to example implementations of aspects of the present disclosure;

[0022] Figure 8 is a block diagram illustrating generation input batches using a distillation training dataset according to example implementations of aspects of the present disclosure;

[0023] Figure 9 is a chart illustrating model results according to example implementations of aspects of the present disclosure;

[0024] Figure 10 is a flow chart diagram illustrating an example method for implementing a twice-mixed inverted bottleneck block according to example implementations of aspects of the present disclosure;

[0025] Figure 11 is a flow chart diagram illustrating an example method for performing efficient neural network architecture search according to example implementations of aspects of the present disclosure;

[0026] Figure 12 is a flow chart diagram illustrating an example method for efficient computation of multi-query attention in a machine-learned sequence processing model according to example implementations of aspects of the present disclosure;

[0027] Figure 13 is a flow chart diagram illustrating an example method for efficient computation of multi-query attention in a machine-learned sequence processing model according to example implementations of aspects of the present disclosure;

[0028] Figure 14 is a flow chart diagram illustrating an example method for improved training efficiency of machine-learned image processing models according to example implementations of aspects of the present disclosure;

[0029] Figure 15 is a flow chart diagram illustrating an example method for training a machine-learned model according to example implementations of aspects of the present disclosure;

[0030] Figure 16 is a block diagram of an example processing flow for using machine-learned model(s) to process input(s) to generate output(s) according to example implementations of aspects of the present disclosure;

[0031] Figure 17 is a block diagram of an example sequence processing model according to example implementations of aspects of the present disclosure;

[0032] Figure 18 is a block diagram of an example technique for populating an example input sequence for processing by a sequence processing model according to example implementations of aspects of the present disclosure;

[0033] Figure 19 is a block diagram of an example model development platform according to example implementations of aspects of the present disclosure;

[0034] Figure 20 is a block diagram of an example training workflow for training a machine-learned model according to example implementations of aspects of the present disclosure;

[0035] Figure 21 is a block diagram of an inference system for operating one or more machine-learned model(s) to perform inference according to example implementations of aspects of the present disclosure;

[0036] Figure 22 is a block diagram of an example networked computing system according to example implementations of aspects of the present disclosure;

[0037] Figure 23 is a block diagram of an example computing device according to example implementations of aspects of the present disclosure; and

[0038] Figure 24 is a block diagram of an example computing device according to example implementations of aspects of the present disclosure. DETAILED DESCRIPTION

[0039] Generally, the present disclosure is directed to efficient neural network architectures, techniques for constructing new efficient neural networks, and approaches to low-latency execution of neural networks.

[0040] In an example aspect, the present disclosure provides a more powerful configuration of an inverted bottleneck block that maintains an efficient execution profile. A bottleneck block can be a portion of a neural network (e.g., an image processing network) that is configured to condense rich feature information into a compressed, lower-dimensional representation. An inverted bottleneck block can use expansion and contraction layers internally within the block to perform nonlinear activations in an expanded state, with the inputs and outputs to the block—and any associated residuals that bypass the block internals—operating in the lower dimensional states. In some aspects, the condensed rich features in the lower dimensional states (the bottleneck tensors) can maintain the important information, with the expanded internal states operating to help increase the expressivenessof nonlinear transformations (e.g., such that visual patterns can be more likely to be linearly separable).

[0041] In general, the framework of an inverted bottleneck block can include an expansion pointwise layer that projects features to a higher dimension, followed by an activation layer that can act as a classifier. A projection pointwise layer can be applied to contract features back to the original dimension (or a different dimension). This compression can use the low-rank nature of the features and can facilitate the use of skip connections, which can enable very deep networks.

[0042] To facilitate information transfer in spatial dimensions, the framework of the inverted bottleneck block can include spatial mixing operators. For example, depthwise convolutions can be used to mix features over one or more spatial dimensions of the input state(s). The factorization of a standard 2D convolution into a depthwise convolution and a pointwise convolution can reduce the number of parameters and operations with relatively little accuracy loss, advantageously resulting in fast and efficient networks.

[0043] In example implementations, the present disclosure provides for a twice-mixed inverted bottleneck block. In an example, a first depthwise convolution can be performed in a lower-dimensional “bottleneck” input state to perform spatial mixing over the rich feature information. After the expansion to the higher-dimensional internal states, a second depthwise convolution can be performed over the expanded features to leverage the additional expressive capacity of a higher-dimensional mixing operation.

[0044] The framework leveraged to powerful effect in the twice-mixed inverted bottleneck block provides for a flexible lattice for building a wide variety of bottleneck blocks. The structure can be parameterized to construct a search space for performing a neural architecture search to identify particularly effective combinations of different parameterizations of the bottleneck blocks. For instance, using the twice-mixed inverted bottleneck block can provide for an inexpensive increase of the network depth and receptive field. Using only spatial mixing in the higher dimension or only spatial mixing in the lower dimension (e.g., a singly-mixed inverted bottleneck block) can provide for ad-hoc spatial and channel mixing tradeoffs. An inverted bottleneck without spatial mixing can boost computational utilization due to highly efficient, accelerator-friendly subroutines for computing pointwise convolutions.

[0045] In this manner, for instance, the inverted bottleneck framework provided herein can provide a universal inverted bottleneck block that can be replicated across network stages while allowing for flexibility at each stage to tune the performance of the model. Dueto this high level of structural repetition, example implementations can facilitate network architecture search paradigms with high levels of weight sharing across candidate networks, thereby unlocking new levels of both search efficiency and resultant model performance.

[0046] For example, an example search system can parameterize a universal inverted bottleneck block with a first parameter that categorically activates a spatial mixing operation in the unexpanded state. The search system can parameterize the universal inverted bottleneck block with a second parameter that categorically activates a spatial mixing operation in the expanded state. The search system can construct multiple different candidate sequences of blocks. Each block of a candidate sequence can be assigned values for the two parameters. In this manner, for instance, a large variety of different network architectures can be explored using a relatively compact search space that admits high levels of parameter sharing.

[0047] The network architecture search can use multiple different optimization targets. For instance, the search can prioritize absolute result quality. The search can prioritize result generation speed (e.g., latency). The search can leverage a balance of quality and speed. Latency targets can be used that target estimated execution latencies on target hardware, allowing different model architectures to be optimized for different systems.

[0048] The inverted bottleneck blocks provided herein can also provide for improved feedforward sublayers in transformer blocks. For instance, a transformer block can include an attention block and a feedforward network that processes the attention output of the attention block. The advantages of the inverted bottleneck blocks provided herein can further improve the expressive power of the attention block.

[0049] To further facilitate efficient execution of transformer blocks, the present disclosure introduces an efficiency-optimized multi-query attention block that can be combined with implementations of the universal inverted bottleneck (e.g., in a neural architecture search). The efficiency-optimized multi-query attention block, variously termed “Mobile MQA” or “M-MQA” herein, includes carefully selected execution sequencing to decrease memory costs and other computational overheads that can slow computation and expended additional resources. For instance, M-MQA can include restructured weight matrices that are specially adapted for performing a tensor product with inputs to the attention block without introducing spurious transpositions. Although transposes can contain no MACs, they can be very time-consuming because they can require reading the full tensor out of memory and then writing it back to memory again. By reducing the number of transpose operations, M-MQA can achieve significant latency improvements over standard multiqueryattention implementations, with example tests herein demonstrating up to about 3X speedup during training and inference on some hardware.

[0050] Example implementations of M-MQA can also include spatial downsampling to reduce a dimensionality of the tensor products computed during the attention operations. This can include downsampling at the inputs and upsampling at the outputs.

[0051] Example implementations of M-MQA can also include tuned tensor reshaping to decrease compiler burden. For example, 3D tensors used in standard MQA can be decomposed into 4D tensors to improve accelerator or processor utilization. Example tests described herein illustrate nearly 30% processor utilization improvements that can provide 3X speed gains.

[0052] To further facilitate efficient search for and training of new architectures that include the universal inverted bottleneck blocks, example implementations of the present disclosure also provide for data-efficient distillation techniques that leverage dynamic mixing of differently-augmented training datasets into each training batch. Further implementations can leverage additional mixing of class-balanced additional training examples from a third dataset. The training batches can be associated with labels predicted by a teacher model. Training selected student models or an architecture search supernet using distillation (e.g., in lieu of directly training on the raw data) can lead to models that are more performant and more robust to minor hyperparameter tweaks. This in turn can provide for more efficient search and improved final outcomes. For example, it is well known that the model quality of ImageNet training can be highly sensitive to data augmentations, regularization, and optimizer hyperparameters. This presents particular challenges in the context of neural architecture search because candidate architectures can vary dramatically and change dynamically during the search and training process. This can render it challenging to select optimal hyperparameters at each stage. Advantageously, distillation can be robust to hyperparameter selection. Distillation data sets may require little to no additional augmentation and can be robust to regularization and optimization configurations. In this manner, the search space for example implementations of the present disclosure may be simplified by avoiding undue search over hyperparameter value spaces.

[0053] Example implementations of the present disclosure can provide a number of technical effects and benefits.

[0054] In an aspect, the present disclosure provides a technical solution that addresses the challenge of efficiently performing neural network architecture search to identify performant configurations for various hardware environments while maintaining highprediction quality. The technical effects of the technology include the reduction of search space complexity, the enhancement of computational efficiency, and the improvement of latency performance across different hardware configurations.

[0055] One of the technical effects is the introduction of a universal inverted bottleneck block that can be parameterized and instantiated in multiple configurations within a neural network architecture. By parameterizing this flexible architecture into a search space that offers a wide range of complementary performance profiles, a relatively compact search space can allow for exploration of a wide range of model performance tradeoffs for a given implementation. By constraining the search space width, more compute can be allocated to deeper searches, resulting in higher performance models output from the search routine for a given amount of search compute.

[0056] Another technical effect is the ability to evaluate and optimize the neural network architectures for different hardware configurations. The technology allows for the determination of execution latencies on various processors (e.g., having different performance capacities), enabling the selection of candidate models that are not only accurate but also tailored for efficient execution on targeted hardware. This is particularly beneficial for edge computing devices where power consumption and heat dissipation can be tightly constrained. By estimating execution times and incorporating latency targets into the reward function used for learning the policy that guides the search, example implementations can provide for resulting neural network architectures that are optimized for real-world deployment scenarios.

[0057] Another technical effect is reduced memory costs and computational overheads in performing multi-query attention with a decreased number of transpose operations required during tensor contractions. This technical effect is achieved by carefully restructuring weight matrices to avoid spurious transpositions, which are typically time- consuming due to reading and writing the full tensor to memory. By decreasing the number of these operations, M-MQA can substantially improve latency and energy cost, leading to speedups during both training and inference on various hardware platforms.

[0058] Furthermore, example implementations provide data-efficient distillation techniques for training neural network models, which is a technical solution to the problem of limited training data and the energy costs to obtain, process, and store training data. By dynamically mixing differently-augmented training datasets and leveraging class-balanced additional training examples, example implementations can provide for trained models thatare more performant and generalize better to unseen data. This in turn can translate to more efficient architecture search processes and improved final model outcomes.

[0059] A technical effect of example implementations of the present disclosure is increased energy efficiency in performing operations using machine-learned models, thereby improving the functioning of computers implementing such models. For instance, example implementations can provide for more energy-efficient runtime execution or inference. In some scenarios, increased energy efficiency can provide for less energy to be used to perform a given task (e.g., less energy expended to maintain the model in memory, less energy expended to perform calculations within the model, etc.). In some scenarios, increased energy efficiency can provide for more task(s) to be completed for a given energy budget (e.g., a larger quantity of tasks, more complex tasks, the same task but with more accuracy or precision, etc.).

[0060] In another example aspect, example implementations can provide for more energy-efficient training operations or model updates. In some scenarios, increased energy efficiency can provide for less energy to be used to perform a given number of update iterations (e.g., less energy expended to maintain the model in memory, less energy expended to perform calculations within the model, such as computing gradients, backpropagating a loss, etc.). In some scenarios, increased energy efficiency can provide for more update iterations to be completed for a given energy budget (e.g., a larger quantity of iterations, etc.). In some scenarios, greater expressivity afforded by model architectures and training techniques of the present disclosure can provide for a given level of functionality to be obtained in fewer training iterations, thereby expending a smaller energy budget. In some scenarios, greater expressivity afforded by model architectures and training techniques of the present disclosure can provide for an extended level of functionality to be obtained in a given number of training iterations, thereby more efficiently using a given energy budget.

[0061] In this manner, for instance, the improved energy efficiency of example implementations of the present disclosure can reduce an amount of pollution or other waste associated with implementing machine-learned models and systems, thereby advancing the field of machine-learning and artificial intelligence as a whole. The amount of pollution can be reduced in toto (e.g., an absolute magnitude thereof) or on a normalized basis (e.g., energy per task, per model size, etc.). For example, an amount of CO2 released (e.g., by a power source) in association with training and execution of machine-learned models can be reduced by implementing more energy-efficient training or inference operations. An amount of heatpollution in an environment (e.g., by the processors / storage locations) can be reduced by implementing more energy-efficient training or inference operations.

[0062] Various example implementations are described herein with respect to the accompanying Figures.

[0063] Figure 1 is a block diagram of an example machine-learned model 100 having an architecture divided into a plurality of blocks. In general, the blocks can perform processing on an input stream and pass a result downstream. Although not shown in Figure 1, it is to be understood that other connection between the blocks can be used (e.g., residuals, skips, etc.).

[0064] An example block can be a twice-mixed inverted bottleneck block that is shown in exploded view in Figure 1. An input state 102 can receive an input to the block. A first spatial mixing operation (e.g., depthwise convolution) can generate, from input state 102, a spatially-mixed input state 104. An expansion or up-projection operation (e.g., pointwise convolution) can expand a dimensionality (e.g., channel depth) of spatially-mixed input state 104 to obtain expanded state 106. A second spatial mixing operation (e.g., depthwise convolution) can generate, from expanded state 106, a spatially-mixed expanded state 108. A dimension reducing or down-projection operation (e.g., a pointwise convolution) can generate, from spatially-mixed expanded state 108, an input state of a next block in a lower dimensional state than the expanded state.

[0065] In this manner, input state(s) of each block can be “bottlenecks” that condense rich feature information. A residual connection 112 can connect these richly featured bottleneck tensors.

[0066] Nonlinear activations can be applied during generation of one or more of the states of the block. Example nonlinear activations include Relu, Relu6, Gelu, swish, hard swish, etc.

[0067] One or more of the states can be generated using linear convolutions. In an example, the down-projection from the expanded state can be performed using a linear operation to form a linear bottleneck. In an example, nonlinear activations are applied for every operation except the down-projection.

[0068] Input state 102 can be an input tensor or multi-dimensional data array. Input state 102 can include multiple channels and spatial dimensions. For example, input state 102 can include a first spatial dimension (e.g., spatial width W) and a second spatial dimension (e.g., spatial height H). Input state 102 can also include a number of input channels. For instance, input tensors can be three-dimensional tensors having a first spatial dimension, asecond spatial dimension, and the number of input channels (e.g., W×H×C). Input tensors can be four-dimensional tensors having the first spatial dimension, the second spatial dimension, the number of input channels, and a third spatial dimension (e.g., W×H×C×S). Input state 102 can include other combinations of one or more spatial and / or channel dimensions. Input state 102 can have a first channel depth (e.g., number of channels).

[0069] Input state 102 can represent various types of data. In an example, input state 102 is pixel data of an image or feature map data generated from pixel data of an image. Input state 102 can represent sensor data or other data structured in a multi-dimensional array.

[0070] Input state 102 can be processed by a depthwise convolution operator to generate spatially-mixed input state 104. The depthwise convolution operator can perform spatial mixing by sliding a convolution window of filters (e.g., kernels) over (e.g., across dimensions of) the first spatial dimension and the second spatial dimension of input state 102 to generate spatially-mixed input state 104.

[0071] Spatially-mixed input state 104 can be an tensor or multi-dimensional data array. Spatially-mixed input state 104 can include multiple channels and spatial dimensions. For example, spatially-mixed input state 104 can include a first spatial dimension (e.g., spatial width W) and a second spatial dimension (e.g., spatial height H). Spatially-mixed input 104 can also include a number of input channels. For instance, input tensors can be three-dimensional tensors having the first spatial dimension, the second spatial dimension, and the number of input channels (e.g., W×H×C). Spatially-mixed input tensors can be four- dimensional tensors having the first spatial dimension, the second spatial dimension, the number of input channels, and a third spatial dimension (e.g., W×H×C×S). Spatially-mixed input state 104 can include other combinations of one or more spatial and / or channel dimensions. Spatially-mixed input state 104 can have a same channel depth (e.g., number of channels) as input state 102.

[0072] Spatially-mixed input state 104 can be processed by an up-projection operator (e.g., a pointwise convolution) to generate expanded state 106 as an output. A pointwise convolution can perform expansion by applying a separate 1x1 convolutional kernel at each spatial location individually for each channel to generate a number of output channels corresponding to a number of filters of the pointwise convolution operator. In this manner, for instance, a pointwise convolution operator can perform up-projection to project a state having a first channel depth to an expanded state having a second channel depth higher than the first channel depth.

[0073] Expanded state 106 can be a tensor or multi-dimensional data array. Expanded state 106 can include multiple channels and spatial dimensions. For example, expanded state 106 can include a first spatial dimension (e.g., spatial width W) and a second spatial dimension (e.g., spatial height H). Expanded state 106 can also include a number of output channels. For instance, output tensors can be three-dimensional tensors having the first spatial dimension, the second spatial dimension, and the number of output channels (e.g., W× H×C). Expanded state 106 can be four-dimensional tensors having the first spatial dimension, the second spatial dimension, the number of output channels, and a third spatial dimension (e.g., W×H×C×S). Expanded state 106 can include other combinations of one or more spatial and / or channel dimensions. Expanded state 106 can have a greater channel depth (e.g., number of channels) than input state 102. For instance, the number of channels of the expanded state 106 can be greater than the number of channels of input state 102 by a ratio corresponding to an expansion factor.

[0074] Expanded state 106 can be processed by a convolution operator (e.g., a depthwise convolution) to generate spatially-mixed expanded state 108 as an output. This can be a second spatial mixing in addition to a first spatial mixing performed in the lower- dimensional state (e.g., at the first channel depth). In an example, spatial mixing can be performed by a second depthwise convolution in addition to a depthwise convolution performed in the lower-dimensional state to input state 102.

[0075] Spatially-mixed expanded state 108 can be a tensor or multi-dimensional data array. Spatially-mixed expanded state 108 can include multiple channels and spatial dimensions. For example, spatially-mixed expanded state 108 can include a first spatial dimension (e.g., spatial width W) and a second spatial dimension (e.g., spatial height H). Spatially-mixed expanded state 108 can also include a number of output channels. For instance, output tensors can be three-dimensional tensors having the first spatial dimension, the second spatial dimension, and the number of output channels (e.g., W×H×C). Spatially- mixed expanded state 108 can be four-dimensional tensors having the first spatial dimension, the second spatial dimension, the number of output channels, and a third spatial dimension (e.g., W× H×C×S). Spatially-mixed expanded state 108 can include other combinations of one or more spatial and / or channel dimensions.

[0076] Spatially-mixed expanded state 108 can be processed by a down-projection operator (e.g., a pointwise convolution) to generate a projected state, such as an input state 110 of a next network block. A pointwise convolution can perform down-projection by applying a separate 1x1 convolutional kernel at each spatial location individually for eachchannel to generate a number of output channels corresponding to a number of filters of the pointwise convolution operator. In this manner, for instance, a pointwise convolution operator can perform down-projection to project a state having a second channel depth to an output state having a third channel depth lower than the second channel depth (which can be the same or different than the first channel depth). The third channel depth of the output state can correspond to the channel depth of the input state of a subsequent network block.

[0077] The output of the inverted bottleneck block can be used to perform a prediction task of machine-learned model 100 or some other task. Example prediction tasks are described herein with respect to machine-learned model 1. For instance, example prediction tasks can include image processing tasks, such as object detection, object classification, object segmentation, image captioning, image embedding, etc.

[0078] Figure 2a is a block diagram illustrating that the spatial mixing portions can be associated with and discretely activated by separate search parameters 202 and 204. Figure 2b is a diagram illustrating different configurations based on different boolean values of the search parameters.

[0079] For example, the leftmost configuration in Figure 2b corresponds to a configuration in which a spatial mixing operator is applied prior to the expansion (e.g., a value 202 of True) and a spatial mixing operator is applied in the expanded state (e.g., a value 204 of True). With reference to the next configuration to the right, the diagram corresponds to a configuration in which a spatial mixing operator is not applied prior to the expansion (e.g., a value 202 of False) and a spatial mixing operator is applied in the expanded state (e.g., a value 204 of True). With reference to the next configuration to the right, the diagram corresponds to a configuration in which a spatial mixing operator is applied prior to the expansion (e.g., a value 202 of True) and a spatial mixing operator is not applied in the expanded state (e.g., a value 204 of False). With reference to the next configuration to the right, the diagram corresponds to a configuration in which a spatial mixing operator is not applied prior to the expansion (e.g., a value 202 of False) and a spatial mixing operator is not applied in the expanded state (e.g., a value 204 of False).

[0080] In an example, an additional configuration can be triggered with another search parameter value 206 in which the depthwise and pointwise operations are fused in a conv2d operation.

[0081] The various configurations shown in Figures 2a and 2b illustrate a flexible framework for building a wide variety of bottleneck blocks using a high degree of sharedinfrastructure. This can simplify a neural network architecture search space to allow for efficient deep searches.

[0082] Furthermore, such an adaptable framework provides flexibility to strike an ad- hoc spatial and channel mixing tradeoff based on the number and arrangement of spatial mixing stages. Furthermore, such an adaptable framework provides flexibility to enlarge the receptive field as needed by expanding the hidden state for processing. Furthermore, such an adaptable framework provides flexibility to max out computational utilization with accelerator-friendly operations (e.g., such as pointwise convolution in a FFN block).

[0083] Figure 3 is a block diagram illustrating a search system 300 for performing a neural architecture search over the constrained search space based on the universal inverted bottleneck block introduced in the present disclosure. A selection policy 302 can guide selection of candidate neural network architectures (e.g., search parameters 202, 204, 206 for one or more blocks of a sequence of blocks) to form candidate model configurations 304. Training inputs 306 can be input (e.g., batchwise) to candidate model configurations 304 to generate training outputs 308.

[0084] System 300 can compute one or more evaluation metric(s) 310 by comparing training outputs 308 with training labels 312, which can be reference outputs associated with training inputs 306. Training labels 312 can be supervised labels or self- or semi-supervised (e.g., generated using a teacher model). Evaluation metric(s) 310 can also process runtime metadata 314 to extract other signals regarding execution of candidate model configuration(s) 304, such as a duration or latency of execution, memory usage, FLOPs, processor utilization rate, etc.

[0085] Evaluation metrics 310 can generate feedback signal(s) 316 that can be used by system 300 to update selection policy 302, learnable parameters of candidate model configuration(s) 304, or both (together, alternating in sequence, etc.). The feedback signal(s) 316 can characterize a performance of the candidate model configuration based on its outputs (e.g., a loss applied to its outputs). The feedback signal(s) 316 can characterize a performance of the candidate model configuration based on its runtime metadata attributes (e.g., a cost of execution, a latency or duration of execution, memory usage, FLOPs, processor utilization rate, etc.). Feedback signal(s) 316 can represent a joint or combined characterization of the performance of the candidate model configuration based on a combination of a measure of its outputs and a measure of its runtime attributes.

[0086] Evaluation metrics 310 can provide an objective with multiple targets, such as multiple latency targets or evaluations. For instance, a latency of the candidate can beevaluated on multiple different hardware platforms to provide a diverse representation of candidate model performance across different hardware. In some implementations, providing dual latency targets resulted in an improvement of MACs and parameter counts of the resulting model without a measured decrease in accuracy.

[0087] In an example, candidate model configuration(s) 304 are formed from components of a supernet. A supernet can include multiple different computation pathways corresponding to subnets that represent different candidate neural network architectures. For example, a supernet can include one or more types of blocks, including one or more types of bottleneck blocks (e.g., universal inverted bottleneck blocks). Different computation pathways of a supernet can be activated and deactivated based on respective search parameters of each block (e.g., search parameters 202, 204, 206 for one or more blocks of the pathway). During search, values for these parameters can be sampled and combined to form candidate model configurations 304.

[0088] In some implementations, weights in the supernet can be shared across candidate architectures. For example, a given component used repeatedly across some or all of the candidate architectures—or a subportion thereof—can be implemented with the same set of learnable parameters. In this manner, for instances, the number of learnable parameters in the supernet can be decoupled from the size and diversity of the architecture search space.

[0089] For instance, in an example, a supernet can be initialized with a shared set of parameters for one or both of the up-projection and down-projection operators (e.g., for the pointwise convolution operators) across all instances of the universal inverted bottleneck block within the supernet. During the search process, different subnetworks are sampled from the supernet by activating or deactivating[Please continue this paragraph] the depthwise convolutional layers in various configurations. The shared parameters for the up-projection and down-projection layers are learned during training, while the architecture of each sampled subnetwork is determined by the values of the first and second parameters 202 and 204. This weight-sharing strategy allows for efficient exploration of a large search space, as most of the parameters are shared across different candidate architectures.

[0090] In some implementations, the neural architecture search can train learnable parameters of the various network components while also learning how to sample better subnetworks from the supernet. For example, the system can provide updates in a two-stage training approach.

[0091] In a first training stage, the learnable parameters of the supernet can be updated using gradient descent based on the gradient of a loss function computed over abatch of training examples from the training dataset. For example, evaluation metrics 310 can compute a loss based on an accuracy or other quality measure of the training outputs 308. This loss can be backpropagated into the supernet to the active parameters (e.g., in the active blocks forming candidate model configuration(s) 304)

[0092] In a second training stage, the subnetwork sampling policy can be updated using reinforcement learning. For instance, a reward signal can be computed for each sampled subnetwork based on its performance (e.g., accuracy, inference cost on one or more target hardware platforms, or a combination thereof). This reward is then used to update the policy using a reinforcement learning algorithm, such as REINFORCE. The system can alternate between these two stages iteratively until a stopping criterion is met (e.g., a maximum number of iterations or a convergence threshold). Other objectives besides reinforcement learning could be used to update the sampling policy, such as evolutionary algorithms or other optimization techniques.

[0093] An example search can be divided into two search stages: a coarse first search stage and a fine second search stage. In a first search stage, search system 300 can search across filter sizes with other searchable parameters frozen (e.g., parameters 202, 204). For instance, filter sizes can be searched while expansion factor and kernel size are frozen (e.g., at 4 and 3x3, respectively). During the coarse search stage, a type of inverted bottleneck block can also be frozen (e.g., an inverted bottleneck with 202 False and 204 True, or another configuration).

[0094] In the fine second search stage, using the filter sizes selected in the coarse stage, search system 300 can search over one or more of the parameters that were frozen in the coarse stage. For example, in the second search stage, search system 300 can search for configurations of the universal inverted bottleneck block using different combinations of first and second parameters 202 and 204 for each of a plurality of universal inverted bottleneck blocks in the candidate network configurations. Further, in the second search stage, search system 300 can search for kernel sizes used within the universal inverted bottleneck blocks. For instance, kernel size can be selected from a set of available kernel sizes (e.g., 3x3, 5x5). In some implementations additional search parameters can be explored in these stages. In some implementations, some of the search parameters can be fixed throughout all or one or more search stages. For instance, an expansion factor can be searched or can be frozen in both the first and the second search stages.

[0095] Comparing the experimental results of a model resulting from a single-stage search with a model resulting from a two-stage search, the two-stage search resulted in amodel with higher accuracy (+0.22% in validation split, +3.6% in training split) and lower latency (-0.18 ms execution time on Pixel 6 EdgeTPU).

[0096] In an example, a search space of search system 300 can be represented as a set of categorical decisions which control different aspects of the network architecture. For example, a single categorical decision might control whether a 3×3, 5×5, or 7×7 convolution is used at a particular position in the network. An architecture can be formed by assignment of values to these categorical decisions for each block of a number of available network blocks (the number and connection of blocks can be frozen or also searched). During a search, a policy π can be learned. The policy can be a probability distribution from which values are sampled for determining the categorical decisions.

[0097] In an example, π can be defined as a collection of independent multinomial variables, one for each of the decisions in our search space. A set of shared weights W can be used to efficiently estimate the quality of candidate architectures in the search space. The system can alternate between learning the shared weights W using gradient descent and learning the policy π using REINFORCE. At each step, search system 300 can first sample a network architecture α ∼ π. Next, system 300 can use the shared weights to estimate the quality Q(α) of the sampled architecture using a single batch of examples from the validation set. Search system 300 can then estimate the inference time of the sampled architecture T(α). The accuracy Q(α) and inference time T(α) can be used to jointly determine the reward r(α) which is used to update the policy π using REINFORCE. Finally, the system 300 can update the shared model weights W by computing a gradient update with respect to the architecture α on a batch of examples from the training set. The above process can be repeated until the search completes based on satisfaction of a stoppage criterion. At the end of the search, the final architecture can be obtained by independently selecting the most probable value for each categorical decision in π.

[0098] In an example, a search space can be constructed with defined attributes to limit a scope of the search. For instance, in an example, one or more initial layers of the network can be fixed. For instance, an example search can start with a Conv2D layer (3x3 kernel, stride 2) followed by one or more Fused Inverted Bottleneck blocks (e.g., parameter 206 True) with stride 2. Similarly, one or more output head layers can be fixed, such as to match the output head configuration in MobileNetV3.

[0099] Further, different operation densities can be prioritized at different layers in the model. For example, a search space can prioritize operations with higher computational density in the initial layers to balance efficiency and accuracy.

[0100] Search system 300 was utilized to directly search for models of different target sizes. For MNv4-Conv-L, the NAS search recipe described above was followed. The search was conducted with 384px as input and dual targets of 2.3ms latency on a Pixel 6 EdgeTPU simulator and 2.0ms latency on a Pixel 7 EdgeTPU simulator. For MNv4-Conv-M, the fine- grained search was relaxed to include extra expansion factor search options of 2, 4, and 6. A target of 0.6ms latency on the Pixel 6 EdgeTPU simulator was set, and 256px inputs were used. MNv4-Conv-S was searched using 0.212 ms as the target with 224px input resolution.

[0101] To create MobileNet-V4 hybrid models, attention blocks were manually added to MobileNet-V4 convolution models. MNv4-Hybrid-L was built on top of MNv4-Conv-L. Four Mobile MQA blocks were added to the second-to-last stage. These blocks had 8 heads and 48 projected dimensions, with stride=2 asymmetric downsampling. Four more Mobile MQA blocks were added to the last stage. These attention blocks had 8 heads and 64 projected dimensions. MNv4-Hybrid-M was built on top of MNv4-Conv-M in a similar manner. The only difference was that 64 was used as the projected dimension for all the MQA blocks, and 4 heads instead of 8. Table 1. Architecture specification of MNv4-Conv-S. Input Block DW K_1 DW K_2 Expanded Dim Output Dim Stride 2242×3 Conv2D - 3×3 - 32 2 1122×32 FusedIB - 3×3 32 32 2 562×32 FusedIB - 3×3 96 64 2 282×64 ExtraDW 5×5 5×5 192 96 2 142×96 IB - 3×3 192 96 1 142×96 IB - 3×3 192 96 1 142×96 IB - 3×3 192 96 1 142×96 IB - 3×3 192 96 1 142×96 ConvNext 3×3 - 384 96 1 142×96 ExtraDW 3×3 3×3 576 128 2 72×128 ExtraDW 5×5 5×5 512 128 1 72×128 IB - 5×5 512 128 1 72×128 IB - 5×5 384 128 1 72×128 IB - 3×3 512 128 1 72×128 IB - 3×3 512 128 1 72×128 Conv2D - 1×1 - 960 1 72×960 AvgPool - 7×7 - 960 1 12×960 Conv2D - 1×1 - 1280 1 12×1280 Conv2D - 1×1 - 1000 1Table 2. Architecture specification of MNv4-Conv-M. Input Block DW K_1 DW K_2 Expanded Dim Output Dim Stride 2562×3 Conv2D - 3×3 - 32 2 1282×32 FusedIB - 3×3 128 48 2 642×48 ExtraDW 3×3 5×5 192 80 2 322×80 ExtraDW 3×3 3×3 160 80 1 322×80 ExtraDW 3×3 5×5 480 160 2 162×160 ExtraDW 3×3 3×3 640 160 1 162×160 ExtraDW 3×3 3×3 640 160 1 162×160 ExtraDW 3×3 5×5 640 160 1 162×160 ExtraDW 3×3 3×3 640 160 1 162×160 ConvNext 3×3 - 640 160 1 162×160 FFN - - 320 160 1 162×160 ConvNext 3×3 - 640 160 1 162×160 ExtraDW 5×5 5×5 960 256 2 82×256 ExtraDW 5×5 5×5 1024 256 1 82×256 ExtraDW 3×3 5×5 1024 256 1 82×256 ExtraDW 3×3 5×5 1024 256 1 82×256 FFN - - 1024 256 1 82×256 ConvNext 3×3 - 1024 256 1 82×256 ExtraDW 3×3 5×5 512 256 1 82×256 ExtraDW 5×5 5×5 1024 256 1 82×256 FFN - - 1024 256 1 82×256 FFN - - 1024 256 1 82×256 ConvNext 5×5 - 512 256 1 82×256 Conv2D - 1×1 - 960 1 82×960 AvgPool - 8×8 - 960 1 12×960 Conv2D - 1×1 - 1280 1 12×1280 Conv2D - 1×1 - 1000 1Table 3. Architecture specification of MNv4-Hybrid-M. Input Block DW K_1 DW K_2 Expanded Dim Output Dim Stride 2562×3 Conv2D - 3×3 - 32 2 1282×32 FusedIB - 3×3 128 48 2 642×48 ExtraDW 3×3 5×5 192 80 2 322×80 ExtraDW 3×3 3×3 160 80 1 322×80 ExtraDW 3×3 5×5 480 160 2 162×160 ExtraDW 3×3 3×3 640 160 1 162×160 ExtraDW 3×3 3×3 640 160 1 162×160 ExtraDW 3×3 5×5 640 160 1 162×160 M-MQA - - - 160 1 162×160 ExtraDW 3×3 3×3 640 160 1 162×160 M-MQA - - - 160 1 162×160 ConvNext 3×3 - 640 160 1 162×160 M-MQA - - - 160 1 162×160 FFN - - 640 160 1 162×160 M-MQA - - - 160 1 162×160 ConvNext 3×3 - 640 160 1 162×160 ExtraDW 5×5 5×5 960 256 2 82×256 ExtraDW 5×5 5×5 1024 256 1 82×256 ExtraDW 3×3 5×5 1024 256 1 82×256 ExtraDW 3×3 5×5 1024 256 1 82×256 FFN - - 1024 256 1 82×256 ConvNext 3×3 - 1024 256 1 82×256 ExtraDW 3×3 5×5 512 256 1 82×256 M-MQA - - - 256 1 82×256 ExtraDW 5×5 5×5 1024 256 1 82×256 M-MQA - - - 256 1 82×256 FFN - - 1024 256 1 82×256 M-MQA - - - 256 1 82×256 FFN - - 1024 256 1 82×256 M-MQA - - - 256 1 82×256 ConvNext 5×5 - 1024 256 1 82×256 Conv2D - 1×1 - 960 1 82×960 AvgPool - 8×8 - 960 1 12×960 Conv2D - 1×1 - 1280 1 12×1280 Conv2D - 1×1 - 1000 1Table 4. Architecture specification of MNv4-Conv-L. Input Block DW K_1 DW K_2 Expanded Dim Output Dim Stride 3842×3 Conv2D - 3×3 - 24 2 1922×24 FusedIB - 3×3 96 48 2 962×48 ExtraDW 3×3 5×5 192 96 2 482×96 ExtraDW 3×3 3×3 384 96 1 482×96 ExtraDW 3×3 5×5 384 192 2 242×192 ExtraDW 3×3 3×3 768 192 1 242×192 ExtraDW 3×3 3×3 768 192 1 242×192 ExtraDW 3×3 3×3 768 192 1 242×192 ExtraDW 3×3 5×5 768 192 1 242×192 ExtraDW 5×5 3×3 768 192 1 242×192 ExtraDW 5×5 3×3 768 192 1 242×192 M-MQA - - - 192 1 242×192 ExtraDW 5×5 3×3 768 192 1 242×192 M-MQA - - - 192 1 242×192 ExtraDW 5×5 3×3 768 192 1 242×192 M-MQA - - - 192 1 242×192 ExtraDW 5×5 3×3 768 192 1 242×192 M-MQA - - - 192 1 242×192 ConvNext 3×3 - 768 192 1 242×192 ExtraDW 5×5 5×5 768 512 2 122×512 ExtraDW 5×5 5×5 2048 512 1 122×512 ExtraDW 5×5 5×5 2048 512 1 122×512 ExtraDW 5×5 5×5 2048 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 ExtraDW 5×5 3×3 2048 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 ExtraDW 5×5 3×3 2048 512 1 122×512 ExtraDW 5×5 5×5 2048 512 1 122×512 M-MQA - - - 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 M-MQA - - - 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 M-MQA - - - 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 M-MQA - - - 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 Conv2D - 1×1 - 960 1 122×960 AvgPool - 12×12 - 960 1 12×960 Conv2D - 1×1 - 1280 1 12×1280 Conv2D - 1×1 - 1000 1Table 5. Architecture specification of MNv4-Hybrid-L. Input Block DW K_1 DW K_2 Expanded Dim Output Dim Stride 3842×3 Conv2D - 3×3 - 24 2 1922×24 FusedIB - 3×3 96 48 2 962×48 ExtraDW 3×3 5×5 192 96 2 482×96 ExtraDW 3×3 3×3 384 96 1 482×96 ExtraDW 3×3 5×5 384 192 2 242×192 ExtraDW 3×3 3×3 768 192 1 242×192 ExtraDW 3×3 3×3 768 192 1 242×192 ExtraDW 3×3 3×3 768 192 1 242×192 ExtraDW 3×3 5×5 768 192 1 242×192 ExtraDW 5×5 3×3 768 192 1 242×192 ExtraDW 5×5 3×3 768 192 1 242×192 M-MQA - - - 192 1 242×192 ExtraDW 5×5 3×3 768 192 1 242×192 M-MQA - - - 192 1 242×192 ExtraDW 5×5 3×3 768 192 1 242×192 M-MQA - - - 192 1 242×192 ExtraDW 5×5 3×3 768 192 1 242×192 M-MQA - - - 192 1 242×192 ConvNext 3×3 - 768 192 1 242×192 ExtraDW 5×5 5×5 768 512 2 122×512 ExtraDW 5×5 5×5 2048 512 1 122×512 ExtraDW 5×5 5×5 2048 512 1 122×512 ExtraDW 5×5 5×5 2048 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 ExtraDW 5×5 3×3 2048 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 ExtraDW 5×5 3×3 2048 512 1 122×512 ExtraDW 5×5 5×5 2048 512 1 122×512 M-MQA - - - 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 M-MQA - - - 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 M-MQA - - - 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 M-MQA - - - 512 1 122×512 ConvNext 5×5 - 2048 512 1 122×512 Conv2D - 1×1 - 960 1 122×960 AvgPool - 12×12 - 960 1 12×960 Conv2D - 1×1 - 1280 1 12×1280 Conv2D - 1×1 - 1000 1Table 6. Training hyper-parameters for ImageNet-1k classification. Conv-S Conv-M Hybrid-M Conv-L Hybrid-L Batch size 4096 4096 16384 16384 16384 Peak learning 0.002 0.004 0.016 0.004 0.01 rate Cosine decay 0.0 0.0 0.0 0.0 0.001 alpha Cosine decay 9600 500 500 500 500 epochs Warm-up epochs 5 5 20 20 20 Training epochs 9600 500 500 500 500 AdamW weight 0.01 0.1 0.1 0.2 0.2 decay AdamW β10.6 0.9 0.9 0.9 0.9 AdamW β20.999 0.999 0.999 0.999 0.999 AdamW ε 10-610-710-710-710-7EMA decay 0.9999 - - - - L2- 10-5- - - - regularization Gradient - - - - - clipping Label smoothing 0.1 0.1 0.1 0.1 0.1 Dropout 0.3 0.2 0.2 0.2 0.2 Peak Stochastic 0 0.075 0.075 0.35 0.35 Depth drop rate RandAugment 0.5 0.7 0.7 1.0 1.0 probability RandAugment 2 2 2 2 2 layers RandAugment 9 15 15 15 15 magnitude RandAugment Cutout Cutout Cutout Cutout Cutout excluded ops Mixup / Cutmix - - - 0.3 0.3 probability Mixup α - - - 0.8 0.8 Cutmix α - - - 1.0 1.0 Mixup / Cutmix - - - 0.5 0.5 switch probabilityTable 7. Training hyper-parameters for distillation. Conv-S Conv-M Hybrid-M Conv-L Hybrid-L Batch size 4096 24576 24576 24576 16384 Peak learning rate 0.004 0.024 0.024 0.024 0.0016 Warm-up epochs 5 30 60 30 12 Training epochs 9600 2000 2000 2000 2000 AdamW weight decay 0.01 0.05 0.05 0.05 0.05 AdamW β10.4 0.9 0.9 0.9 0.9 AdamW β20.999 AdamW ε 10-7Peak Stochastic Depth drop rate

[0102] Figure 4 is a block diagram of an example implementation of a transformer block that uses the universal inverted bottleneck block to process the attention output of a multi-headed attention block as an improved feedforward network.

[0103] In an example, multi-headed self-attention (MHSA) can include multiple attention layers (heads) in parallel with different linear transformations on the queries, keys, values and outputs.

[0104] Figure 5 is a block diagram of an example implementation of a transformer block that uses the universal inverted bottleneck block to process the attention output of a multi-query attention (MQA) block as an improved feedforward network.

[0105] In an example, MQA can be similar to MHSA except that the different heads share a single set of keys and values. Since the keys and values can be shared across all of the different attention “heads”, the size of these tensors and hence the memory bandwidth requirements of incremental decoding can be significantly reduced.

[0106] Figure 6 is a block diagram of an example implementation of M-MQA, or mobile multi-query attention as provided herein. In Figure 6, matrix projection operations in attention have been expressed as a 1x1 Conv2d and a tensor reshape. In some implementations, explicitly controlling the reshape operation (e.g., instead of letting a compiler handle implicitly) can allow for fine-tuned control over the low-level operations performed in the projection operation. For example, specific reshape operations can include decomposing a large dimension into multiple smaller dimensions to provide for increased parallelism, higher processor utilization, lower memory usage, lower latency, etc.

[0107] For example, the efficiency of deep learning models running on accelerators can be highly dependent on the shapes of their tensors. For instance, modern acceleratorarchitectures can be represented as a collection of parallel execution units and a multilayer memory hierarchy. Higher efficiency is generally achieved when the shape of a data tile is aligned with the hardware characteristics, including the memory bank, memory transaction length, and minimum schedulable unit (e.g., warp size in GPUs).

[0108] To address the efficiency problem at the source, expressing the model using 4D tensor shapes for MatMul can ease the burden of the compiler. As a concrete example, consider the top part of the M-MQA block, which corresponds to the matrix projection of the output tensor in attention. This operation can be implemented in three steps: a matrix multiply, a reshape, and a convolution. The reshape operation can have significant impact on the performance of the subsequent convolution.

[0109] For instance, as a baseline for comparison, consider the following execution graph:

[0110] OPERATOR 1: MATRIX MULTIPLY

[0111] Input Shape: [1][1]

[1152]

[0144]

[0112] Operation: Matrix multiplication (einsum_4 / Einsum2)

[0113] Kernel: 1x1

[0114] Stride: 1x1

[0115] Parameters:

[0116] num_macs: 10,618,632 (number of multiply-accumulate operations)

[0117] run_time: 2.41 microseconds

[0118] utilization: 76.94%

[0119] Output Shape: [1][1]

[1152]

[0064]

[0120] OPERATOR 2: RESHAPE

[0121] Input Shape: [1][1]

[1152]

[0064]

[0122] Operation: Reshaping tensor dimensions

[0123] Parameters:

[0124] num_macs: 0

[0125] run_time: 8.28 microseconds

[0126] utilization: 0.00%

[0127] Output Shape: [1][1]

[0144]

[0512]

[0128] OPERATOR 3: CONVOLUTION

[0129] Input Shape: [1][1]

[0144]

[0512]

[0130] Operation: Convolution (einsum_5 / Einsum1.conv)

[0131] Parameters:

[0132] param_size: 262,144

[0133] num_macs: 37,748,736

[0134] Energy Consumption:

[0135] io_milli_watts: 815.41436

[0136] run_time: 33.11 microseconds

[0137] utilization: 19.88%

[0138] memory_time: 7.00 microseconds

[0139] Output Shape: [1][1]

[0144]

[0512]

[0140] The RESHAPE operator can be adjusted to decompose one of the dimensions into multiple smaller dimensions to provide for increased parallelism across the additional dimension. For example, consider the following adjusted execution graph:

[0141] OPERATOR 1: MATRIX MULTIPLY

[0142] Input Shape: [1][1]

[1152]

[0144]

[0143] Operation: Matrix multiplication (einsum_4 / Einsum2)

[0144] Parameters:

[0145] num_macs: 10,618,632

[0146] run_time: 2.41 microseconds

[0147] utilization: 76.97%

[0148] Energy Consumption:

[0149] io_milli_watts: 0.00000

[0150] milli_watts: 0.00000

[0151] milli_joules: 0.00000

[0152] Output Shape: [1][1]

[1152]

[0064]

[0153] OPERATOR 2: RESHAPE

[0154] Input Shape: [1][1]

[1152]

[0064]

[0155] Operation: Reshaping tensor dimensions

[0156] Parameters:

[0157] num_macs: 0 (no computational operations)

[0158] run_time: 2.28 microseconds

[0159] utilization: 0.00%

[0160] Output Shape: [1]

[0012]

[0012]

[0512]

[0161] OPERATOR 3: CONVOLUTION

[0162] Input Shape: [1]

[0012]

[0012]

[0512]

[0163] Operation: Convolution (einsum_5 / Einsum1.conv)

[0164] Parameters:

[0165] param_size: 262,144 (number of parameters)

[0166] num_macs: 37,748,736

[0167] run_time: 11.19 microseconds

[0168] utilization: 58.83%

[0169] Energy Consumption:

[0170] io_milli_watts: 2412.86872

[0171] Memory Time: 7.00 microseconds

[0172] Output Shape: [1]

[0012]

[0012]

[0512]

[0173] In comparison against the baseline, reshaping the tensor by decomposing the size 144 dimension into two dimensions of size 12 enabled the convolution operator to achieve higher processing utilization (58.83% vs.19.88%), enabling faster runtime (11.19 µs vs.33.11 µs). For example, flattened tensors ([1][1]

[0144]

[0512] ) may have limited parallelization while, if the data is restructured spatially (

[0012]

[0012]

[0512] ), each tile (patch) can be processed in parallel, better leveraging the power of the accelerator.

[0174] In some implementations, spatial downsampling may be used. Spatial downsampling can be applied to one or more inputs. This can be implemented using depthwise convolutions with a stride parameter selected to control an amount of downsampling.

[0175] For example, a query input tensor may be downsampled using a depthwise convolution operator (e.g., 3x3 with stride 2). Similarly, a key input tensor and value input tensor can be downsampled using a depthwise convolution operator (e.g., 3x3 with stride 2). Such downsampling may be asymmetric – for instance, the key and value tensors may downsampled while the query tensor is preserved at full resolution. This asymmetric downsampling can reduce a number of tokens in stages with high spatial resolution while preserving token count in high impact processing paths (e.g., the query tensor). In an example, asymmetric downsampling as compared to no downsampling might only decrease top-1 accuracy by 0.06% while achieving a 3% decrease in MACs, a 23% decrease in CPU time, and a 25% decrease in GPU time (measured using MNv4-Hybrid-M base model on Samsung S23. Stride-2 down-sampling is applied at penultimate 16x16 stage.).

[0176] Furthermore, careful control of the ordering of the tensor product operations can yield additional improvements. For example, structuring the weight matrices such that tensor contractions can be performed without transposing the inputs can yield lower latencycomputation with fewer memory reads / writes. For example, a tensor product with two inputs can be translated to a matrix multiplication if both the contracted and non-contracted indices from each input are adjacent to each other. This is because adjacent indices can be very cheaply reshaped to a single index for the purpose of running the low-level matrix multiplications and then very cheaply reshaped back. If the indices are not cleanly split into a set of contracting indices and non-contracting indices, then transposition operations may be needed to bring them into this form for matrix multiplication. Exploiting this opportunity, example implementations of the M-MQA block are structured from the ground up to store the weight matrices in a format specially adapted for more efficient matrix reshaping, yielding significant efficiency gains. Further, the outputs of the tensor products can be requested / queried in a manner that avoid unnecessary transposes (e.g., all transposes), such as by preserving an order of indices between the inputs and outputs.

[0177] In the example below, for the slow implementation, two transposes must be introduced: one to transpose O, and one to transpose Po.

[0178] # Slow implementation: 2 transposes needed

[0179] Y = tf.einsum("bhnv,hdv->bnd", O, P_o)

[0180] # Fast implementation: no transpose needed

[0181] Y = tf.einsum("bnhv,hvd->bnd", O, P_o)

[0182] In the slow implementation of the following example, the non-contracted indices h and k from Pq is interleaved with index n from X. As a result, two transposes will be introduced, one before the matrix multiplication to transpose Pq, and one afterwards to swap the indices in the output tensor.

[0183] #Slow implementation: 2 transposes needed

[0184] Q = tf.einsum("bnd,hdk->bhnk", X, P_q)

[0185] # Fast implementation: no transpose needed

[0186] Q = tf.einsum("bnd,dhk->bnhk", X, P_q)

[0187] As an example, one implementation of multi-query attention can be expressed as in the following pseudocode:

[0188] def MQA(X, M, mask, P_q, P_k, P_v, P_o):

[0189] Q = tf.einsum("bnd,hdk->bhnk", X, P_q)

[0190] K = tf.einsum("bmd,dk->bmk", M, P_k)

[0191] V = tf.einsum("bmd,dv->bmv", M, P_v)

[0192] logits = tf.einsum("bhnk,bmk->bhnm", Q, K)

[0193] weights = tf.softmax(logits + mask)

[0194] O = tf.einsum("bhnm,bmv->bhnv", weights, V)

[0195] Y = tf.einsum("bhnv,hdv->bnd", O, P_o)

[0196] An example implementation of M-MQA can be expressed as in the following pseudocode:

[0197] def Mobile_MQA(X, M, mask, P_q, P_k, P_v, P_o):

[0198] Q = tf.einsum("bnd,dhk->bnhk", X, P_q)

[0199] K = tf.einsum("bmd,dk->bmk", M, P_k)

[0200] V = tf.einsum("bmd,dv->bmv", M, P_v)

[0201] logits = tf.einsum("bnhk,bmk->bnhm", Q, K)

[0202] weights = tf.softmax(logits + mask)

[0203] O = tf.einsum("bnhm,bmv->bnhv", weights, V)

[0204] Y = tf.einsum("bnhv,dhv->bnd", O, P_o)

[0205] In combination, these improvements can provide marked computational gains over baseline MQA. In some implementations, Mobile MQA can offer improved performance over multiheaded self attention (MHSA) and over unmodified MQA. Table 8 records the results of a test in which additional attention blocks are added to the last stage of a base model of MNv4-Conv-L. Attention blocks are inserted interleaving with UIB blocks, as shown in Figures 4 and 5. Models are trained with 400 epochs with distillation. Table 8. An ablation study of attention blocks.

[0206] Figure 7 is a block diagram illustrating an example system for generating a plurality of training datasets that can be dynamically mixed for training a student model based on outputs of a teacher model 700. A first dataset 702-1 can include a first plurality of training examples 704-1. The training examples 704-1 can be processed using a first set of augmentation operations 706-1 to generate a first augmented dataset 708-1 containing a first plurality of augmented training examples 710-1. These augmented training examples 710-1can be input to teacher model 700 to generate training labels 712-1 that correspond to augmented training examples 710-1.

[0207] To generate a diverse corpus of training examples, this procedure can be repeated with different configurations (e.g., different base dataset, different augmentations, etc.). For instance, the procedure can be repeated N times to generate N datasets. For example, an N-th dataset 702-N can include an N-th plurality of training examples 704-N. Training examples 704-N can be processed using an N-th set of augmentation operations 706-N to generate an N-th augmented dataset 708-N containing an N-th plurality of augmented training examples 710-N. These augmented training examples 710-N can be input to teacher model 700 to generate training labels 712-N that correspond to augmented training examples 710-N. Base dataset 702-N can be the same as or different from base dataset 702-1.

[0208] In an example, N is equal to three, such that three sets of augmented datasets are generated. In an example, two of the three base datasets are the same. In an example described in the attached Appendix, one base dataset is ImageNet-1K, and one base dataset is JFT-300M.

[0209] In an example, augmentation operations can include random cropping, where portions of the image are randomly selected and cropped to focus on different regions of the image; rotation, where the image is rotated by a random angle within a specified range; flipping, where the image is flipped horizontally or vertically to simulate different orientations; scaling, where the image is resized to different scales to simulate varying distances from the object of interest; translation, where the image is shifted horizontally or vertically to simulate off-center positioning; color jittering, which alters the brightness, contrast, saturation, and hue of the image to simulate different lighting conditions; noise injection, where random pixel noise is added to the image to simulate sensor noise or environmental interference; perspective transformation, which warps the image to simulate different camera angles and viewpoints; and elastic deformation, where the image is locally stretched or squeezed to simulate non-rigid transformations. These operations can be applied in isolation or in combination and can be parameterized to control the intensity and frequency of the augmentations.

[0210] As an example, geometric transformations can alter the spatial arrangement of pixels in an image, effectively simulating changes in the camera perspective, object orientation, or scene layout. Common geometric transformations include rotation, where the image is rotated by a random angle within a specified range; translation, where the image is shifted horizontally and / or vertically by a random number of pixels; scaling, where the imageis resized to a different scale factor to simulate zooming in or out; and flipping, where the image is mirrored along the vertical or horizontal axis. More complex geometric transformations might involve shearing, which skews the image along one axis, or elastic deformations that apply a localized warping effect to simulate non-rigid transformations.

[0211] As an example, photometric transformations can modify the visual appearance of the image in terms of color, brightness, contrast, and texture. Examples include brightness adjustment, where the intensity of all pixels is increased or decreased to simulate variations in lighting conditions; contrast adjustment, which alters the difference between the darkest and lightest parts of the image; saturation adjustment, which modifies the intensity of colors; and hue adjustment, which shifts the colors along the color spectrum. Noise injection, such as adding Gaussian or salt-and-pepper noise, can simulate sensor noise or graininess in images. Color jittering randomly changes the brightness, contrast, saturation, and hue of the image in a single augmentation step.

[0212] As an example, random erasing or occlusion can introduce information loss by masking out random regions of the image. This can simulate occlusions that might occur in real-world scenarios, such as objects being partially blocked by other objects. The size, shape, and location of the occluded region can be chosen randomly, and the region can be filled with a constant value (such as black or gray), random noise, or pixel values sampled from other parts of the image. Cutout augmentation is a specific type of random erasing where a square region can be removed from the image.

[0213] Perspective warping can change the apparent perspective of the scene, which can simulate the effect of taking a photo from a different angle or height. Barrel and pincushion distortions can mimic the optical distortions introduced by camera lenses, where the image appears to bulge outwards or pinch inwards, respectively.

[0214] Synthetic combinations can involve blending multiple images to create composite images. This can be done, for example, through alpha blending, where two images are combined with varying transparency levels, or through mixup, where two images and their corresponding labels are combined linearly.

[0215] In some applications, domain-specific augmentations can be used to reflect particular aspects of downstream tasks. For instance, in medical imaging, augmentations can include simulating different imaging modalities or artifacts specific to medical equipment. In satellite imagery, augmentations could involve changes in the angle of sunlight, cloud cover, or seasonal variations in vegetation.

[0216] Different combinations of augmentation operations can be used in any of augmentation operation(s) 706-1, ... , 706-N, etc.

[0217] Table 9 reports a set of example datasets. D1 applies Inception Crop and RandAugument, with 2 layers and magnitude 9, on 500 replicas of ImageNet-1k. D2 applies Inception Crop and extreme Mixup on a dataset which is 1000 replicas of ImageNet-1K. D3 applies InceptionCrop and RandAugment to a resampled JFT-300M dataset. The resampled dataset was resampled to create a new class-balanced dataset with 130K images per class, resulting in 130M images.

[0218] EfficientNet-B0, trained on ImageNet-1K, was used to select relevant images with a threshold of 0.3. For frequently-appearing classes, the top 130K samples were selected. For rarely-appearing classes, images were randomly selected images to replicate to upsample. The dataset was replicated 10X for distillation. Table 9. Distillation datasets created with an EfficientNet L2 teacher. Dataset Source Augmentations D1500 × ImageNet 1k Inception Crop + RandAug 12m9 D21000 × ImageNet 1k Inception Crop + Extreme Mixup D310 × JFT subset Inception Crop + RandAug 12m5

[0219] Training batches can be dynamically mixed. Figure 8 is a block diagram of an example system for dynamically mixing training inputs for an input batch for training a student model. A dynamic batch generator 800 can receive training examples from augmented datasets 708-1, ... , 708-N to generate training inputs 306. Dynamic batch generator 800 can include examples from different datasets in different proportions.

[0220] Batch mixing can occur at inference time. For instance, batch generator 800 can generate a batch on demand according to a configured proportion or ratio.

[0221] To perform distillation for some example experiments, images were first augmented and then compressed using JPEG encoding to create the offline distillation dataset. The images were then decoded, and inference was run using the teacher.

[0222] The teacher’s predicted probabilities (e.g., the softmax output) of the 1k classes were recorded and were to be used later as soft labels to train the student. EfficientNet L2 was used as the teacher model. This teacher model had 480 M parameters and 290 G MACs. It achieved 87.5% top-1 accuracy on ImageNet 1k. During student training, all data augmentation except left-right flip was disabled. Dropout in the final fully connected layer was also applied as the only regularization method. Cross-entropy between the teacher’s softlabels and the student’s predicted class probability vectors was used as the training loss. The student was trained with the AdamW optimizer. A cosine learning rate schedule with warm- up was used. Table 10. Distillation results using MNv4-Conv-L as student. Top 1 Acc (Val / Train) Mixing ratio 400 epochs 2000 epochs D1- 83.8 / 86.6 - D2- 84.1 / 85.6 - D3- 81.8 / 84.1 - D1+ D21:1 84.4 / 85.0 - D2+ D31:1 84.7 / 82.7 - D1+ D2+ D31:1:2 84.9 / 82.6 85.9 / 85.5

[0223] As shown in Table 10, the distillation result obtained with D1 is inferior to D2 - 83.8% vs 84.1%. However, when D1 and D2 are dynamically mixed during student training, performance improves to 84.4% (+0.3%). Training with dynamic mixing multiple datasets performed better than the best individual dataset training. For instance, the high training accuracy of the D1-only student suggests that this set produces inferior performance because the augmented images are too easy for the student. However, since D1 contains a very different set of augmented images, combining it with images from D2 may help to enlarge the volume of the overall augmented image space, which may lead to better results than either dataset alone.

[0224] In one implementation of a full distillation recipe (represented in the last row of Table 3), datasets D1, D2, and D3 were dynamically mixed in training batches, with a ratio of 1:1:2. With this distillation recipe, 84.9% top-1 accuracy on ImageNet-1k was achieved by MNv4-Conv-L, with 400 epochs of training. With 2000 training epochs, the top-1 accuracy is increased by another 1% to 85.9%. This student model has only 31M parameters and 6G MACs, making it 15x smaller than the teacher in terms of parameter size, and 48x smaller in terms of MACs. Despite this, a top-1 accuracy that is only 1.6% lower than the teacher is achieved, demonstrating the effectiveness of the distillation method. Besides the significant boost in model accuracy, other benefits of using distillation training are observed. For example, distillation training was found to be 2.5x faster than training on ImageNet-1k. With a batch size of 16k and 400 training epochs, distillation training on MNv4-Conv-Large took only 2.3 hours on 128 TPU v5.

[0225] Further tests were performed to evaluate the improvements of the distillation techniques described herein on existing architectures. Table 11. Retraining results–Various MobileNet versions with the top-1 accuracy originally reported in the literature, top-1 accuracy results from the improved training recipes disclosed herein, and top-1 accuracy results from training with more resources than just ImageNet-1k e.g. distillation and / or pre-training on JFT. Model Prior SOTAWith Present TrainingWith JFT Recipeadded V1-1.0x 70.60% 74.00% - V2-1.0x 72.00% 73.40% - V3-1.0x 75.20% 75.50% - V4-Conv-S - 73.80% 75.50% V4-Conv-M - 79.90% 82.70% V4-Hybrid- M- 80.70% 83.70%V4-Conv-L - 82.90% 85.90% V4-Hybrid-L - 83.40% 87.00% Table 12. Top-1 Accuracy Comparison Across Training Approaches: This table contrasts baseline ImageNet-1k training, SOTA distillation, and our distillation. SOTA PresentImageNet-1konlyDistillation* Distillation V4-Conv-S 73.8 - 75.5 V4-Conv-M 79.9 81.5 82.7 V4-Hybrid-M 80.7 82.7 83.7 V4-Conv-L 82.9 84.4 85.9 V4-Hybrid-L 83.4 85.7 86.6 * Patient Teacher approach from Lucas Beyer, Xiaohua Zhai, Amélie Royer, Larisa Markeeva, Rohan Anil, and Alexander Kolesnikov. Knowledge distillation: A good teacher is patient and consistent. In Proceedings of the IEEE / CVF conference on computer vision and pattern recognition, pages 10925–10934, 2022.Example Results

[0226] An example implementation of searched models were evaluated on the ImageNet-1K dataset. Top-1 accuracy is reported on the validation split of the dataset. "Top- 1 Open" is also reported, which refers to ImageNet-1K results without restrictions on the training data. This includes models that were trained via distillation and on additional data beyond the ImageNet-1K training set, which was deduplicated against the ImageNet validation split. In the ImageNet-only setting, the models were trained for 500 epochs. In the open setting, the models were trained for 2K epochs with the distillation dataset. Additionally, the largest model, MNv4-Hybrid-L, was pretrained with JFT-300M before distillation. For a fair comparison, MobileNets V1, V2, and V3 were re-trained with the same ImageNet-only training recipe. These results are provided in Table 13.

[0227] Expanded results are provided in Table 14. For these results, the latency analysis includes a representative selection of mobile hardware, including ARM Cortex CPUs (Pixel 6, Samsung S23), Qualcomm Hexagon DSP (Pixel 4), ARM Mali GPU (Pixel 7), Qualcomm Snapdragon (S23 GPU), Apple Neural Engine, and Google EdgeTPU.

[0228] The effectiveness of MNv4 backbones for object detection tasks was evaluated on the COCO 17 dataset. M-sized MNv4 backbones were compared against SOTA efficient backbones with a similar number of Madds. For each backbone, an object detector using the RetinaNet framework was built. A 256-d FPN decoder was attached to the P3 - P7 endpoints, along with a 256-d prediction head, which included 4 convolutional layers. Depth-separable convolutions were employed, as usual for mobile detectors, to reduce the computational complexity of both the FPN decoder and the box prediction head. All models were trained on the COCO 17 training set for 600 epochs. All images were resized to 384px and augmented using random horizontal flip, random scale, as well as Randaug. Shear and Rotate augmentations from Randaug were excluded, as those deformations were found to deteriorate the detection AP of small objects. Training was performed using a 2048 batch size, Adam, and an L2 weight decay of 0.00003. A cosine learning rate schedule with 24 epochs warm-up was used, and the learning rate for each model was separately tuned. For all baselines, the filter multiplier was set so that MACs were roughly comparable. Following classification experiments, MobileNet V4 backbones were trained using a stochastic drop rate of 0.2. All MobileNet baselines were trained using the official Tensorflow Model Garden implementation. EfficientFormer was re-implemented in Tensorflow.

[0229] A consistent benchmarking strategy was used across various mobile platforms, with an exception for the Apple Neural Engine. To enhance efficiency, models wereconverted to TensorFlow Lite format and quantized to INT8 for mobile CPUs, Hexagon, and EdgeTPUs, while FP16 was used for mobile GPUs. Each model was run approximately 1000 times, and the mean latency of all runs is captured. That process was repeated 5 times for each model and the median of those means is reported.

[0230] The CPU affinity was set to the fastest core and used the XNNPACK backend for CPU evaluations.

[0231] For benchmarks on the Apple Neural Engine (conducted on an iPhone 13 with iOS 16.6.1, CoreMLTools 7.1, and Xcode 15.0.1 for profiling), PyTorch models were converted to CoreML’s MLProgram format in Float16 precision, with float16 MultiArray inputs to minimize input copying.

[0232] Experimental results were reported in Table 15. Parameters, MACs, and benchmarks were computed using the entire detector at the 384px input resolution.Table 13. ImageNet-1K classification results. Median latency is reported. †=GeLu, otherwise ReLU is used. * indicates additional data was used. ** indicates additional pretraining with JFT. Latency (ms) Top-1Top-1Params MACs Pixel 6 Pixel 8 Open(M) (G) GPU EdgeTPU MNv3-0.5xL 69.2 - 2.7 0.1 2.4 0.8 MobileOne-S0 71.4 72.5 2.1 0.3 4.2 0.7 MNv4-Conv-S 73.8 75.5* 3.8 0.2 2.4 0.7 MobileNet-v1 74 - 4.2 0.6 6.1 0.9 FastViT-T8† 75.6 76.7 3.6 0.7 49.3 1.3 MobileOne-S4 79.4 81.4 14.8 1.5 26.7 1.7 FastViT-S12† 79.8 80.9 8.8 1.8 86.5 1.8 E-FormerV2-S2† - 82 12.6 1.3 57.2 - MNv4-Conv-M 79.9 82.7* 9.2 1 11.4 1.1 FastViT-SA12† 80.6 81 10.9 1.9 86.5 2 MNv4-Hybrid-M 80.7 83.7* 10.5 1.2 14.3 1.5 E-FormerV2-L† - 83.5 26.1 2.6 117.7 - FastViT-SA24† 82.6 83.4 20.6 3.8 171.6 3.2 MNv4-Conv-L 82.9 85.9* 31 5.9 59.9 2.4 ConvNext-S† 83.1 84.6* 50 8.7 314.9 3.7 NextViT-B† 83.2 85.1 44.8 8.3 - - MNv4-Hybrid-L† 83.486.6* / 35.9 7.2 87.6 3.8FastViT-SA36† 83.6 84.2 30.4 5.6 241.6 4.3Table 14. Expanded Classification results on ImageNet-1K. Median latency is reported. A hyphen indicates that a model was not benchmarked due to missing corresponding model files for a platform. Failed indicates that the model is not supported by the platform. Latency (ms) Para. MACPixelPixel iPhone Pixel Pixel Samsung 6 8 13 4 7 S23 ModelTop-1(M) (G) CPUGPU CPU GPUMobileNet-V2- 0.5x66 2.0 0.1 2.4 0.7 0.5 2.9 8.3 1.8 1.9MobileNet- V3L-0.5x69.2 2.7 0.1 2.4 0.8 0.45 3.5 9.9 2.0 2.1MobileOne-S0 71.4 2.1 0.3 4.2 0.7 0.5 2.9 10.7 3.3 1.7 MobileNet-V2 73.4 3.5 0.3 5.0 0.7 0.7 3.9 13.6 4.1 2.5 MNv4-Conv-S 73.8 3.8 0.2 2.4 0.7 0.6 2.4 8.4 1.8 2.0 MobileNet-V1 74.0 4.2 0.6 6.1 0.8 0.7 3.2 13.0 4.6 2.1 FastViT-T8† 75.6 3.6 0.7 49.3 1.3 0.7 Failed 40.7 43.6 24.7 MobileNet-V2- 1.5x76.8 6.8 0.7 9.3 0.9 1.0 5.6 16.4 7.3 3.3MultiHardware- MAX-1.5x77.9 8.9 0.8 9.8 1.0 - 5.7 23.2 - 4.1MultiHardware- AVG-1.5x78.2 10.0 1.0 12.0 1.1 - 6.1 20.3 - 4.5MobileNet-V2- 2.0x78.4 11.2 1.1 13.9 1.1 1.5 6.9 19.1 10.6 4.2MobileOne-S4 79.4 14.8 1.5 26.7 1.7 1.5 9.0 28.6 19.4 5.9 FastViT-S12† 79.8 8.8 1.8 83.0 1.8 1.6 Failed 75.0 69.2 47.0 MIT- EfficientViT- 79.4 9.1 0.5 - - 2.4 - - 18.1 5.0 B1-r224 MNv4-Conv-M 79.9 9.2 1.0 11.4 1.1 1.1 7.3 18.1 8.6 4.1 FastViT-SA12 80.6 10.9 1.9 86.5 2.0 1.6 Failed 79.6 69.5 52.180.7 10.5 1.2 14.3 1.5 - Failed 17.9 10.8 5.9FastViT-SA24 82.6 20.6 3.8 171.6 3.2 2.4 Failed 131.9 136.3 107.5 MIT- EfficientViT- 82.7 24.0 2.1 - - 5.4 - - 64.9 9.5 B2-r256 MNv4-Conv-L 82.9 31 5.9 59.9 2.4 3.0 20.8 37.6 43.0 13.2 ConvNext-S 83.1 50 8.7 314.9 3.7 - Failed 45.2 243.9 18.5 NextViT-B 83.2 44.8 8.3 - - - - - - -83.4 35.9 7.2 87.6 3.8 - Failed 61.3 61.8 18.1MIT- EfficientViT- 83.5 49.0 4.0 - - 12.2 - - 125.9 18.4 B3-r224 FastViT-SA36 83.6 30.4 5.6 241.6 4.3 - Failed 186.5 206.3 138.1Table 15. Object detection results on the COCO-17 Val. set. The width-multiplier is reported next to the MobileNet backbones that were scaled-up. COCO MACs Params Pixel 6 CPU Backbone Val AP (G) (M) latency (ms) EfficientFormer L1 29.5 6.54 12.77 84.3 MobileNet v1 @ 1.5 31.0 6.68 9.05 66.4 MNv4-Conv-M 32.6 5.06 9.79 51.3 MobileNet Multi-AVG @ 1.5 32.7 5.42 9.51 58.1 MobileNet v2 @ 2.0 32.9 5.81 10.15 66.4 MobileNet v3 Large@ 2.0 33.2 4.99 17.92 59.9 MNv4-Hybrid-M 34.0 5.62 11.15 60.5

[0233] Figure 9 charts further example results achieved by example implementations of the present disclosure. An example implementation of the present disclosure is labeled “MobileNetV4” in the chart. Example Methods

[0234] Figure 10 depicts a flowchart of a method 1000 for using a twice-mixed inverted bottleneck block to efficiently process spatial features in input data array. One or more portion(s) of example method 1000 can be implemented by a computing system that includes one or more computing devices such as, for example, computing systems described with reference to the other figures. Each respective portion of example method 1000 can be performed by any (or any combination) of one or more computing devices. Moreover, one or more portion(s) of example method 1000 can be implemented on the hardware components of the device(s) described herein, for example, to train one or more systems or models. Figure 10 depicts elements performed in a particular order for purposes of illustration and discussion. Those of ordinary skill in the art, using the disclosures provided herein, will understand that the elements of any of the methods discussed herein can be adapted, rearranged, expanded, omitted, combined, or modified in various ways without deviating from the scope of the present disclosure. Figure 10 is described with reference to elements / terms described with respect to other systems and figures for exemplary illustrated purposes and is not meant to be limiting. One or more portions of example method 1000 can be performed additionally, or alternatively, by other systems.

[0235] At 1002, example method 1000 can include receiving, at an input to the twice- mixed inverted bottleneck block, an input state of a neural network having a first channel depth (e.g., input state 102).

[0236] At 1004, example method 1000 can include convolving, using a first depthwise convolutional layer of the twice-mixed inverted bottleneck block, the input state to obtain a spatially-mixed input state (e.g., spatially-mixed input state 104).

[0237] At 1006, example method 1000 can include expanding, using an up-projection pointwise convolutional layer of the twice-mixed inverted bottleneck block, the spatially- mixed input state to generate an expanded state having a second channel depth greater than the first channel depth (e.g., expanded state 106).

[0238] At 1008, example method 1000 can include convolving, using a second depthwise convolutional layer of the twice-mixed inverted bottleneck block, the expanded state to obtain a spatially-mixed expanded state (e.g., spatially-mixed expanded state 108).

[0239] At 1010, example method 1000 can include projecting, using a down- projection pointwise convolutional layer of the twice-mixed inverted bottleneck block, the spatially-mixed expanded state to obtain a projected state having a third channel depth less than the second channel depth (e.g., state 110).

[0240] At 1012, example method 1000 can include using the projected state to perform a prediction task with the neural network. For example, in an implementation, the input state represents image data, and the prediction task comprises an image processing task. Example prediction tasks can include any one or more of the tasks or outputs described herein with respect to machine-learned model 1 or 4.

[0241] In some implementations of example method 1000, the neural network comprises a plurality of instances of the twice-mixed inverted bottleneck block in contiguous sequential order.

[0242] In some implementations of example method 1000, the first depthwise convolutional layer comprises a first kernel size. In some implementations of example method 1000, the second depthwise convolutional layer comprises a second kernel size, wherein the first kernel size is different than the second kernel size.

[0243] In some implementations of example method 1000, the first depthwise convolutional layer comprises a first kernel size. In some implementations of example method 1000, the second depthwise convolutional layer comprises a second kernel size, wherein the first kernel size is equal to the second kernel size.

[0244] In some implementations of example method 1000, the neural network can include an inverted bottleneck block with a different configuration (e.g., a singly-mixed block). In some implementations of example method 1000, processing an output of the twice-mixed inverted residual block can be performed using a singly-mixed inverted residual bottleneck block that performs spatial mixing in the expanded state.

[0245] In some implementations of example method 1000, the neural network processes an output of the twice-mixed inverted residual block using a singly-mixed inverted residual bottleneck block. In some implementations, example method 1000 includes receiving, at an input to the singly-mixed inverted bottleneck block, a downstream input state having a fourth channel depth. In some implementations, example method 1000 includes expanding, using an up-projection pointwise convolutional layer of the singly-mixed inverted bottleneck block, the downstream input state to generate a downstream expanded state having a fifth channel depth greater than the fourth channel depth. In some implementations, example method 1000 includes convolving, using a first and only depthwise convolutional layer of the singly-mixed inverted bottleneck block, the downstream expanded state to obtain a spatially-mixed downstream expanded state. In some implementations, example method 1000 includes projecting, using a down-projection pointwise convolutional layer of the singly-mixed inverted bottleneck block, the spatially-mixed downstream expanded state to obtain a downstream projected state having a sixth channel depth less than the fifth channel depth.

[0246] In some implementations of example method 1000, the neural network processes an output of the twice-mixed inverted residual block using a singly-mixed inverted residual bottleneck block. In some implementations, example method 1000 includes receiving, at an input to the singly-mixed inverted bottleneck block, a downstream input state having a fourth channel depth. In some implementations, example method 1000 includes convolving, using a first and only depthwise convolutional layer of the singly-mixed inverted bottleneck block, the downstream input state to obtain a spatially-mixed downstream input state. In some implementations, example method 1000 includes expanding, using an up- projection pointwise convolutional layer of the singly-mixed inverted bottleneck block, the spatially-mixed downstream input state to generate a downstream expanded state having a fifth channel depth greater than the fourth channel depth. In some implementations, example method 1000 includes projecting, using a down-projection pointwise convolutional layer of the singly-mixed inverted bottleneck block, the downstream expanded state to obtain a downstream projected state having a sixth channel depth less than the fifth channel depth.

[0247] In some implementations of example method 1000, the neural network comprises a transformer (e.g., a transformer encoder, a transformer block), and wherein thetwice-mixed inverted bottleneck block processes a downstream output of an attention block. For example, Figures 4 and 5 illustrate an example of such a configuration.

[0248] In some implementations of example method 1000, the neural network comprises a transformer encoder, and wherein the twice-mixed inverted bottleneck block processes a downstream output of a multi-query attention block. For example, Figure 5 illustrates an example of such a configuration.

[0249] In some implementations of example method 1000, the neural network comprises a transformer encoder, and wherein the twice-mixed inverted bottleneck block processes a downstream output of a mobile multi-query attention block. For example, Figure 6 illustrates an example of such a configuration. In some implementations of example method 1000, the mobile multi-query attention block implements processing techniques further described with respect to example methods 1200 and 1300.

[0250] In some implementations of example method 1000, the neural network comprises at least one attention block that implements multi-query attention over inputs to generate an array of attention values which are used to generate an output. For instance, an output of an attention block can correspond to an updated version of an input. For example, one or more positions in an input tensor X can be updated based on a combination of values corresponding to the positions in M in respective amounts corresponding to the weight of attention applied to those positions.

[0251] In some implementations of example method 1000, the neural network processes an output of the twice-mixed inverted residual block using at least one attention block that implements multi-query attention over inputs to generate an array of attention values which are used to generate an output. For instance, an output of an attention block can correspond to an updated version of an input. For example, one or more positions in an input tensor X can be updated based on a combination of values corresponding to the positions in M in respective amounts corresponding to the weight of attention applied to those positions.

[0252] In some implementations of example method 1000, at least one attention block projects the inputs to a projection dimension of at least 48. In some implementations of example method 1000, at least one attention block projects the inputs to a projection dimension of at least 64.

[0253] In some implementations of example method 1000, at least one dimension of a state is decomposed into two smaller dimensions for input to a convolution operator.

[0254] In some implementations of example method 1000, the neural network implementing the block is a search outcome of example method 1100.

[0255] Figure 11 depicts a flowchart of a method 1100 for performing efficient neural network architecture search. One or more portion(s) of example method 1100 can be combined with or independent from any one or more of the example methods described above. One or more portion(s) of example method 1100 can be implemented by a computing system that includes one or more computing devices such as, for example, computing systems described with reference to the other figures. Each respective portion of example method 1100 can be performed by any (or any combination) of one or more computing devices. Moreover, one or more portion(s) of example method 1100 can be implemented on the hardware components of the device(s) described herein, for example, to train one or more systems or models. Figure 11 depicts elements performed in a particular order for purposes of illustration and discussion. Those of ordinary skill in the art, using the disclosures provided herein, will understand that the elements of any of the methods discussed herein can be adapted, rearranged, expanded, omitted, combined, or modified in various ways without deviating from the scope of the present disclosure. Figure 11 is described with reference to elements / terms described with respect to other systems and figures for exemplary illustrated purposes and is not meant to be limiting. One or more portions of example method 1100 can be performed additionally, or alternatively, by other systems.

[0256] At 1102, example method 1100 can include initializing a search space comprising a plurality of configurations of a universal inverted bottleneck block. For example, search system 300 can initialize the search space. The search space can be defined as described herein with respect to Figures 2A and 2B.

[0257] For example, the universal bottleneck block can include a first placeholder for a depthwise convolutional layer that is discretely activated during search using a first parameter (e.g., parameter 202). The universal bottleneck block can include an up-projection pointwise convolutional layer that follows the first placeholder and expands a channel depth of an internal state of the network (e.g., state 104) to an expanded state (e.g., state 106). The universal bottleneck block can include a second placeholder for a depthwise convolutional layer that is discretely activated during search using a second parameter (e.g., parameter 204), the second placeholder following the expanding pointwise convolutional layer. The universal bottleneck block can include a down-projection pointwise convolutional layer that follows the second placeholder and decreases a channel depth of an internal state of the network from the expanded state to a projected state (e.g., state 110). In an example, the placeholder can perform a unit transform or passthrough if not activated by the first or second parameter.

[0258] At 1104, example method 1100 can include constructing a plurality of candidate models, wherein constructing the plurality of candidate models comprises, for each candidate model of the plurality of candidate models: selecting, for each respective network block of a plurality of network blocks of a neural network architecture, a respective value of the first parameter (e.g., parameter 202) and a respective value of the second parameter (e.g., parameter 204) to instantiate a respective inverted bottleneck block at the respective network block. For example, a search system 300 can use a search policy to select (e.g., based on a sampled value based on the policy) values for selectable parameters 202, 204, 206, etc. to configure blocks of candidate architectures using the universal inverted bottleneck block framework.

[0259] At 1106, example method 1100 can include training the plurality of candidate models over a training dataset. The training dataset can include labeled training examples. The training dataset can include a distillation dataset as described herein. The training dataset can be composed as described in further detail with respect to example method 1400.

[0260] At 1108, example method 1100 can include evaluating performance of the plurality of candidate models. For example, search system 300 can compute a loss or evaluate other evaluation metrics to characterize a performance of a candidate.

[0261] At 1110, example method 1100 can include outputting a selected candidate model based on the evaluating.

[0262] In some implementations of example method 1100, the neural network architecture is configured for image processing tasks, such as image classification or image segmentation, multimodal tasks, such as image and audio processing tasks or image and text processing tasks.

[0263] In some implementations of example method 1100, the evaluating comprises determining a prediction quality of the plurality of candidate models on prediction tasks.

[0264] In some implementations of example method 1100, the evaluating comprises determining a first latency of execution of the plurality of candidate models on a first hardware configuration. In some implementations of example method 1100, the evaluating comprises determining a second latency of execution of the plurality of candidate models on a second hardware configuration.

[0265] In some implementations of example method 1100, the latency search can be expanded to an arbitrary number of latencies. The latency cost value can be generated taking the geometric mean of the number of latencies. Other metrics can be used as a proxy forlatency, such as counting a total number of FLOPs. These proxies can be used alone or in combination with estimated or actual measures of latency.

[0266] In some implementations of example method 1100, the search policy can be learned using a reward value computed using a prediction quality term. The prediction quality term can be reduced using a first latency divergence term indicating a measure of divergence of the first latency from a first latency target (e.g., a magnitude of a difference between unity and a ratio of a test time and a target time). The prediction quality term can be reduced using a second latency divergence term indicating a measure of divergence of the second latency from a second latency target.

[0267] In some implementations, example method 1100 comprises sampling, probabilistically, the respective value of the first parameter and the respective value of the second parameter using a policy learned based on prior evaluation of the performance of the plurality of candidate models.

[0268] In some implementations of example method 1100, the policy is learned using a reward value computed using a prediction quality term that is reduced by a latency divergence term, the latency divergence term indicating a measure of divergence of the first latency or the second latency from a latency target.

[0269] In some implementations of example method 1100, the policy is learned using a reward value computed using a prediction quality term that is reduced by a first latency divergence term indicating a measure of divergence of the first latency from a first latency target. In some implementations of example method 1100, the policy is learned using a reward value computed using a prediction quality term that is reduced by a second latency divergence term indicating a measure of divergence of the second latency from a second latency target.

[0270] In some implementations of example method 1100, determining the first latency comprises estimating execution time on a first processor, and determining the second latency comprises estimating execution time on a second, different processor.

[0271] In some implementations of example method 1100, different latency targets for different hardware can result in a search outcome that provides architectures performant over a broad range of hardware.

[0272] In some implementations of example method 1100, the first processor is a processor having a thermal design power (TDP) of less than 5 watts, and the second processor is a different processor having a TDP of less than 5 watts.

[0273] In some implementations of example method 1100, the first processor is a processor having a thermal design power (TDP) of less than 5 watts, and the second processor is a different processor having a TDP of greater than 5 watts.

[0274] In some implementations of example method 1100, the first processor is a processor having a thermal design power (TDP) of less than 5 watts, and the second processor is a different processor having a TDP of greater than 5 watts and less than 50 watts.

[0275] In some implementations of example method 1100, the first processor is a processor having a thermal design power (TDP) of greater than 5 watts and less than 10 watts, and the second processor is a different processor having a TDP of greater than 5 watts and less than 10 watts.

[0276] In some implementations of example method 1100, the first processor is a processor having a thermal design power (TDP) of less than 10 watts, and the second processor is a different processor having a TDP of greater than 10 watts.

[0277] In some implementations of example method 1100, the first processor is a processor having a thermal design power (TDP) of less than 10 watts, and the second processor is a different processor having a TDP of greater than 10 watts and less than 50 watts.

[0278] In some implementations of example method 1100, the first processor is a processor having a thermal design power (TDP) of less than 50 watts, and the second processor is a different processor having a TDP of less than 50 watts.

[0279] In some implementations of example method 1100, the first processor is a processor having a thermal design power (TDP) of greater than 50 watts, and the second processor is a different processor having a TDP of greater than 50 watts.

[0280] In some implementations of example method 1100, the first processor or the second processor is an accelerator device specifically adapted for parallelized matrix multiplication and designed to work cooperatively with a central processing unit (CPU) of a computing system.

[0281] In some implementations of example method 1100, the latency target is less than 20 ms, such as less than 15 ms, such as less than 10 ms, such as less than 5 ms, such as less than 3 ms, such as less than 2 ms, such as less than 1 ms.

[0282] In some implementations of example method 1100, for instances of a given network block across the plurality of candidate models: a first shared set of parameters are used for the corresponding up-projection pointwise convolutional layers; and a second sharedset of parameters are used for the corresponding down-projection pointwise convolutional layers.

[0283] In some implementations of example method 1100, for instances of all network blocks across the plurality of candidate models: a first shared set of parameters are used for the corresponding up-projection pointwise convolutional layers; and a second shared set of parameters are used for the corresponding down-projection pointwise convolutional layers.

[0284] In some implementations of example method 1100, the plurality of candidate models share greater than 95% of learnable network parameters.

[0285] In some implementations of example method 1100, the search space comprises an attention block that implements multi-query attention over inputs to generate an array of attention values which are used to generate an output. For instance, an output of an attention block can correspond to an updated version of an input. For example, one or more positions in an input tensor X can be updated based on a combination of values corresponding to the positions in M in respective amounts corresponding to the weight of attention applied to those positions.

[0286] Figure 12 depicts a flowchart of a method 1200 for efficient computation of multi-query attention in a machine-learned sequence processing model. One or more portion(s) of example method 1200 can be combined with or independent from any one or more of the example methods described above. One or more portion(s) of example method 1200 can be implemented by a computing system that includes one or more computing devices such as, for example, computing systems described with reference to the other figures. Each respective portion of example method 1200 can be performed by any (or any combination) of one or more computing devices. Moreover, one or more portion(s) of example method 1200 can be implemented on the hardware components of the device(s) described herein, for example, to train one or more systems or models. Figure 12 depicts elements performed in a particular order for purposes of illustration and discussion. Those of ordinary skill in the art, using the disclosures provided herein, will understand that the elements of any of the methods discussed herein can be adapted, rearranged, expanded, omitted, combined, or modified in various ways without deviating from the scope of the present disclosure. Figure 12 is described with reference to elements / terms described with respect to other systems and figures for exemplary illustrated purposes and is not meant to be limiting. One or more portions of example method 1200 can be performed additionally, or alternatively, by other systems.

[0287] At 1202, example method 1200 can include obtaining a current input tensor representing an input sequence of one or more elements. For example, an input tensor X can represent N elements, each element corresponding to a d-dimensional embedded representation.

[0288] At 1204, example method 1200 can include computing attention values over the current input tensor using one or more tensor contractions. As an example, attention can be performed between X and a tensor M. M can represent M elements, each element corresponding to a d-dimensional embedded representation. In self-attention, attention can be performed between X and itself.

[0289] In some implementations of example method 1200, at least a portion of respective tensor contractions of the one or more tensor contractions comprise an output tensor having non-contracting indices inherited from a first tensor and non-contracting indices inherited from a second tensor. The non-contracting indices inherited from the first tensor can be contiguous and ordered as in the first tensor, and the non-contracting indices inherited from the second tensor can be contiguous and ordered as in the second tensor.

[0290] For example, in the tensor contraction Q = tf.einsum("bnd,dhk->bnhk", X, P_q) the input tensor X can be constructed such that the non-contracting indices b and n are contiguous and ordered as in the desired output Q, which has dimensions bnhk. Similarly, a projection tensor P_q can be constructed such that the non-contracting indices h and k are contiguous and ordered as in the desired output Q.

[0291] At 1206, example method 1200 can include generating an output of the machine-learned sequence processing model based on the attention values. Example outputs are described herein with respect to machine-learned models 1 and 4.

[0292] In some implementations of example method 1200, the portion of respective tensor contractions of the plurality of tensor contractions comprises all the plurality of tensor contractions. For example, all the tensor contractions can be optimized according to example method 1200 to avoid unnecessary transpositions which can incur unnecessary memory read costs.

[0293] In some implementations of example method 1200, non-contracting indices for the first tensor are adjacent indices of the first tensor. In some implementations of example method 1200, non-contracting indices for the second tensor are adjacent indices of the second tensor.

[0294] In some implementations, example method 1200 includes reshaping, without transposition, the first tensor into a first matrix by transforming the adjacent indices of the first tensor into a single index. In some implementations, example method 1200 includes reshaping, without transposition, the second tensor into a second matrix by transforming the adjacent indices of the second tensor into a single index.

[0295] In some implementations, example method 1200 includes reshaping at least one dimension of a tensor into at least two smaller dimensions for performing a convolution operation.

[0296] Figure 13 depicts a flowchart of a method 1300 for efficient computation of multi-query attention in a machine-learned sequence processing model. One or more portion(s) of example method 1300 can be combined with or independent from any one or more of the example methods described above. One or more portion(s) of example method 1300 can be implemented by a computing system that includes one or more computing devices such as, for example, computing systems described with reference to the other figures. Each respective portion of example method 1300 can be performed by any (or any combination) of one or more computing devices. Moreover, one or more portion(s) of example method 1300 can be implemented on the hardware components of the device(s) described herein, for example, to train one or more systems or models. Figure 13 depicts elements performed in a particular order for purposes of illustration and discussion. Those of ordinary skill in the art, using the disclosures provided herein, will understand that the elements of any of the methods discussed herein can be adapted, rearranged, expanded, omitted, combined, or modified in various ways without deviating from the scope of the present disclosure. Figure 13 is described with reference to elements / terms described with respect to other systems and figures for exemplary illustrated purposes and is not meant to be limiting. One or more portions of example method 1300 can be performed additionally, or alternatively, by other systems.

[0297] At 1302, example method 1300 includes obtaining an input tensor representing at least one input sequence of one or more elements, the input tensor having a sequence length dimension and an embedding dimension. For example, an input tensor X can represent N elements, each element corresponding to a d-dimensional embedded representation. In einsum notation, for instance, X can correspond to indices bnd, corresponding to batch, sequence length, and embedding dimension, respectively. In some cases, the batch dimension can be unitary, although multiple inputs can be batched.

[0298] At 1304, example method 1300 includes obtaining a query projection tensor configured to generate a projection of the input tensor using a plurality of learned weights, the query projection tensor having a query head dimension, an embedding dimension, and a projection dimension, wherein adjacent indices of the query projection tensor respectively correspond to the query head dimension and the projection dimension. For example, a query projection tensor P_q can have indices dhk corresponding to an embedding dimension, a head dimension, and a projection dimension.

[0299] At 1306, example method 1300 includes contracting the input tensor and the query projection tensor in the embedding dimension.

[0300] Contracting the input tensor and the query projection tensor in the embedding dimension can be performed by, at 1306-1, reshaping the query projection tensor into a query matrix by combining the query head dimension and the projection dimension into a single dimension. For example, a query matrix can be characterized by the non-contracting indices hk of the query projection tensor.

[0301] Contracting the input tensor and the query projection tensor in the embedding dimension can be performed by, at 1306-2, obtaining an input matrix comprising at least a portion of the input tensor.

[0302] Contracting the input tensor and the query projection tensor in the embedding dimension can be performed by, at 1306-3, multiplying the input matrix and the query matrix to obtain an output.

[0303] Contracting the input tensor and the query projection tensor in the embedding dimension can be performed by, at 1306-4, reshaping the output into a query tensor having a query head dimension, a sequence length dimension, and a projection dimension, wherein adjacent indices of the query tensor respectively correspond to the query head dimension and the projection dimension. For instance, a query tensor Q can be generated that is characterized by indices bnhk, with h and k being adjacent.

[0304] At 1308, example method 1300 includes using the query tensor to compute attention values over the at least one input sequence. For example, attention values can include self-attention values that map interrelatedness of elements in a sequence represented by X with the other elements in X. Attention values can be cross-attention values that map interrelatedness of elements in X with elements in another sequence, such as M.

[0305] In some implementations of example method 1300, the attention values can be used to perform a prediction task. Example prediction tasks are described herein with respect to machine-learned model 1, 4.

[0306] For instance, an output of an attention block can correspond to an updated version of an input. For example, one or more positions in an input tensor X can be updated based on a combination of values corresponding to the positions in M in respective amounts corresponding to the weight of attention applied to those positions. This updated version of X can be provided to downstream network components, such as classification heads, language decoders, etc. In some implementations, a classification token position in X is updated over one or more passes through an attention block, and the classification token value is used as input to a downstream network component.

[0307] In some implementations of example method 1300, the contracting of the input tensor and the query projection tensor is performed without transposing any portion of the query projection tensor or any portion of the query tensor.

[0308] In some implementations of example method 1300, the contracting of the input tensor and the query projection tensor is performed without generating an intermediate array in memory. For instance, data copying can be avoided, with the reshaping only using changed mapping to memory addresses.

[0309] In some implementations of example method 1300, using the query tensor to attend over the at least one input sequence comprises obtaining an input tensor representing an input sequence of one or more elements, the input tensor having a sequence length dimension and an embedding dimension. For example, an input tensor X can represent N elements, each element corresponding to a d-dimensional embedded representation. In einsum notation, for instance, X can correspond to indices bnd, corresponding to batch, sequence length, and embedding dimension, respectively. In some cases, the batch dimension can be unitary, although multiple inputs can be batched.

[0310] In some implementations of example method 1300, using the query tensor to attend over the at least one input sequence comprises obtaining a key projection tensor configured to generate a projection of the input tensor using a plurality of learned weights, the key projection tensor having an embedding dimension and a projection dimension. For example, a key projection tensor can be expressed as P_k with einsum indices dk.

[0311] In some implementations of example method 1300, using the query tensor to attend over the at least one input sequence comprises contracting the input tensor and the key projection tensor in the embedding dimension to obtain a key tensor. For example, a key tensor can be expressed as K with einsum indices bmk.

[0312] In some implementations of example method 1300, using the query tensor to attend over the at least one input sequence comprises obtaining a value projection tensorconfigured to generate a projection of the input tensor using a plurality of learned weights, the value projection tensor having an embedding dimension and a projection dimension. For example, a value projection tensor can be expressed as P_v with einsum indices dv.

[0313] In some implementations of example method 1300, using the query tensor to attend over the at least one input sequence comprises contracting the input tensor and the value projection tensor in the embedding dimension to obtain a value tensor. For example, a value tensor can be expressed as V with einsum indices bmv.

[0314] In some implementations of example method 1300, using the query tensor to attend over the at least one input sequence comprises contracting the query tensor and the key tensor along the projection dimension to obtain a logit tensor. For example, a logit tensor can be expressed as L with einsum indices bnhm.

[0315] In some implementations of example method 1300, using the query tensor to attend over the at least one input sequence comprises applying a softmax operation over the logit tensor to obtain a weight tensor. For example, a weight tensor can be expressed as W with einsum indices bnhm. The weight tensor can represent the attention weight on a particular position in M from the perspective of a particular position in X.

[0316] In some implementations of example method 1300, using the query tensor to attend over the at least one input sequence comprises obtaining an output tensor by contracting the weight tensor and the value tensor over the sequence length dimension. For example, an output tensor can be expressed as O with einsum indices bnhv.

[0317] In some implementations of example method 1300, using the query tensor to attend over the at least one input sequence comprises obtaining an updated tensor for the input sequence re-projecting the output tensor from the projection dimension to the embedding dimension using an output projection tensor. For example, each position in X updated based on values corresponding to the positions in M in an amount corresponding to the weight of attention applied to that position. The updated tensor can be represented as Y with indices bnd.

[0318] In some implementations, example method 1300 includes contracting the output tensor and the output projection tensor along the projection dimension. In some implementations of example method 1300, non-contracting indices of the output tensor are adjacent, and non-contracting indices of the output projection tensor are adjacent.

[0319] In some implementations, example method 1300 includes reshaping at least one dimension of a tensor into at least two smaller dimensions for performing a convolution operation.

[0320] In some implementations, example method 1300 includes reshaping at least one dimension of the output tensor into at least two smaller dimensions, wherein the output projection tensor is defined with at least two corresponding decomposed dimensions.

[0321] In some implementations of example method 1300, the output tensor is reshaped into a four-dimensional tensor, and wherein the output projection tensor is defined as a four-dimensional tensor.

[0322] In some implementations, example method 1300 includes computing the query tensor using a downsampled input state.

[0323] In some implementations, example method 1300 includes computing the query tensor using a downsampled input state obtained using a depthwise convolution having a first stride parameter.

[0324] In some implementations, example method 1300 includes computing the key tensor, the value tensor, or both using a downsampled input state.

[0325] In some implementations, example method 1300 includes computing the key tensor, the value tensor, or both using a downsampled input state obtained using a depthwise convolution having a second stride parameter.

[0326] In some implementations, example method 1300 includes computing the query tensor using a first downsampled input state obtained using a depthwise convolution having a first stride parameter; and computing the key tensor, the value tensor, or both using a second downsampled input state obtained using a depthwise convolution having a second stride parameter. In some implementations of example method 1300, the first stride parameter is different from the second stride parameter, or the first stride parameter is equal to the second stride parameter.

[0327] Figure 14 depicts a flowchart of a method 1400 for improved training efficiency of machine-learned image processing models. One or more portion(s) of example method 1400 can be combined with or independent from any one or more of the example methods described above. One or more portion(s) of example method 1400 can be implemented by a computing system that includes one or more computing devices such as, for example, computing systems described with reference to the other figures. Each respective portion of example method 1400 can be performed by any (or any combination) of one or more computing devices. Moreover, one or more portion(s) of example method 1400 can be implemented on the hardware components of the device(s) described herein, for example, to train one or more systems or models. Figure 14 depicts elements performed in a particular order for purposes of illustration and discussion. Those of ordinary skill in the art,using the disclosures provided herein, will understand that the elements of any of the methods discussed herein can be adapted, rearranged, expanded, omitted, combined, or modified in various ways without deviating from the scope of the present disclosure. Figure 14 is described with reference to elements / terms described with respect to other systems and figures for exemplary illustrated purposes and is not meant to be limiting. One or more portions of example method 1400 can be performed additionally, or alternatively, by other systems.

[0328] At 1402, example method 1400 can include obtaining a distillation dataset that comprises a plurality of training examples labeled by a teacher model, wherein the plurality of training examples comprises: a first plurality of training examples generated using a first set of augmentation routines; and a second plurality of training examples generated using a second set of augmentation routines different from the first set of augmentation routines.

[0329] At 1404, example method 1400 can include generating a training batch for training a student model by combining one or more of the first plurality of training examples and one or more of the second plurality of training examples.

[0330] At 1406, example method 1400 can include processing batchwise the training batch using the student model to obtain a batch of outputs.

[0331] At 1408, example method 1400 can include evaluating the batch of outputs using corresponding labels for the training batch generated by the teacher model.

[0332] At 1410, example method 1400 can include updating one or more learnable parameters of the student model based on the evaluating.

[0333] In some implementations of example method 1400, the student model is a supernet implemented by a neural network architecture search, wherein the supernet comprises multiple different computation pathways corresponding to subnets that represent different candidate neural network architectures.

[0334] In some implementations of example method 1400, the first set of augmentation routines and the second set of augmentation routines share a common augmentation routine.

[0335] In some implementations of example method 1400, the common augmentation routine implements random cropping.

[0336] In some implementations of example method 1400, the first set of augmentation routines induces less distortion than the second set of augmentation routines.

[0337] In some implementations of example method 1400, the plurality of training examples comprises a third plurality of training examples generated using a third set ofaugmentation routines applied to a base dataset different from a base dataset used for the first plurality of training examples or the second plurality of training examples.

[0338] In some implementations of example method 1400, the third plurality of training examples is class-balanced using a labeling network to classify individual training examples.

[0339] In some implementations of example method 1400, the third plurality of training examples is class-balanced using a labeling network to classify individual training examples, wherein the third plurality of training examples comprises an equal class population for each class. In some implementations of example method 1400, for a first class, the equal class population is less than a total number of training examples in the first class in the base dataset. In some implementations of example method 1400, for a second class, the equal class population is greater than a total number of training examples in the second class in the base dataset, the total number of training examples in the second class being supplemented by replicating random training examples in the second class to reach the equal class population.

[0340] In some implementations of example method 1400, generating the training batch comprises combining one or more of the first plurality of training examples and one or more of the second plurality of training examples and one or more of the third plurality of training examples.

[0341] In some implementations of example method 1400, generating the training batch comprises combining, in equal proportion, one or more of the first plurality of training examples and one or more of the second plurality of training examples and one or more of the third plurality of training examples.

[0342] In some implementations of example method 1400, generating the training batch comprises combining one or more of the first plurality of training examples, one or more of the second plurality of training examples, and one or more of the third plurality of training examples. In some implementations of example method 1400, the batch comprises more of the third plurality of training examples than either of the first plurality of training examples or the second plurality of training examples.

[0343] In some implementations of example method 1400, generating the training batch comprises combining one or more of the first plurality of training examples and one or more of the second plurality of training examples and one or more of the third plurality of training examples. In some implementations of example method 1400, the batch comprises anumber of the third plurality of training examples at least equal to a number of the first plurality of training examples and the second plurality of training examples combined.

[0344] In some implementations of example method 1400, generating the training batch comprises combining, in a ratio of 1:1:2, one or more of the first plurality of training examples, one or more of the second plurality of training examples, and one or more of the third plurality of training examples.

[0345] Figure 15 depicts a flowchart of a method 1500 for training one or more machine-learned models according to aspects of the present disclosure. For instance, an example machine-learned model can include a machine-learned model 100, teacher model 700, etc.

[0346] One or more portion(s) of example method 1500 can be implemented by a computing system that includes one or more computing devices such as, for example, computing systems described with reference to the other figures. Each respective portion of example method 1500 can be performed by any (or any combination) of one or more computing devices. Moreover, one or more portion(s) of example method 1500 can be implemented on the hardware components of the device(s) described herein, for example, to train one or more systems or models. Figure 15 depicts elements performed in a particular order for purposes of illustration and discussion. Those of ordinary skill in the art, using the disclosures provided herein, will understand that the elements of any of the methods discussed herein can be adapted, rearranged, expanded, omitted, combined, or modified in various ways without deviating from the scope of the present disclosure. Figure 15 is described with reference to elements / terms described with respect to other systems and figures for exemplary illustrated purposes and is not meant to be limiting. One or more portions of example method 1500 can be performed additionally, or alternatively, by other systems.

[0347] At 1502, example method 1500 can include obtaining a training instance. A set of training data can include a plurality of training instances divided between multiple datasets (e.g., a training dataset, a validation dataset, or testing dataset). A training instance can be labeled or unlabeled. Although referred to in example method 1500 as a “training” instance, it is to be understood that runtime inferences can form training instances when a model is trained using an evaluation of the model’s performance on that runtime instance (e.g., online training / learning). Example data types for the training instance and various tasks associated therewith are described throughout the present disclosure.

[0348] At 1504, example method 1500 can include processing, using one or more machine-learned models, the training instance to generate an output. The output can be directly obtained from the one or more machine-learned models or can be a downstream result of a chain of processing operations that includes an output of the one or more machine- learned models.

[0349] At 1506, example method 1500 can include receiving an evaluation signal associated with the output. The evaluation signal can be obtained using a loss function. Various determinations of loss can be used, such as mean squared error, likelihood loss, cross entropy loss, hinge loss, contrastive loss, or various other loss functions. The evaluation signal can be computed using known ground-truth labels (e.g., supervised learning), predicted or estimated labels (e.g., semi- or self-supervised learning), or without labels (e.g., unsupervised learning). The evaluation signal can be a reward (e.g., for reinforcement learning). The reward can be computed using a machine-learned reward model configured to generate rewards based on output(s) received. The reward can be computed using feedback data describing human feedback on the output(s).

[0350] At 1508, example method 1500 can include updating the machine-learned model using the evaluation signal. For example, values for parameters of the machine-learned model(s) can be learned, in some embodiments, using various training or learning techniques, such as, for example, backwards propagation. For example, the evaluation signal can be backpropagated from the output (or another source of the evaluation signal) through the machine-learned model(s) to update one or more parameters of the model(s) (e.g., based on a gradient of the evaluation signal with respect to the parameter value(s)). For example, system(s) containing one or more machine-learned models can be trained in an end-to-end manner. Gradient descent techniques can be used to iteratively update the parameters over a number of training iterations. In some implementations, performing backwards propagation of errors can include performing truncated backpropagation through time. Example method 1500 can include implementing a number of generalization techniques (e.g., weight decays, dropouts, etc.) to improve the generalization capability of the models being trained.

[0351] In some implementations, example method 1500 can be implemented for training a machine-learned model from an initialized state to a fully trained state (e.g., when the model exhibits a desired performance profile, such as based on accuracy, precision, recall, etc.).

[0352] In some implementations, example method 1500 can be implemented for particular stages of a training procedure. For instance, in some implementations, examplemethod 1500 can be implemented for pre-training a machine-learned model. Pre-training can include, for instance, large-scale training over potentially noisy data to achieve a broad base of performance levels across a variety of tasks / data types. In some implementations, example method 1500 can be implemented for fine-tuning a machine-learned model. Fine-tuning can include, for instance, smaller-scale training on higher-quality (e.g., labeled, curated, etc.) data. Fine-tuning can affect all or a portion of the parameters of a machine-learned model. For example, various portions of the machine-learned model can be “frozen” for certain training stages. For example, parameters associated with an embedding space can be “frozen” during fine-tuning (e.g., to retain information learned from a broader domain(s) than present in the fine-tuning dataset(s)). An example fine-tuning approach includes reinforcement learning. Reinforcement learning can be based on user feedback on model performance during use. Example Machine-Learned Models

[0353] Figure 16 is a block diagram of an example processing flow for using machine-learned model(s) 1 to process input(s) 2 to generate output(s) 3.

[0354] Machine-learned model(s) 1 can be or include one or multiple machine- learned models or model components. Example machine-learned models can include neural networks (e.g., deep neural networks). Example machine-learned models can include non- linear models or linear models. Example machine-learned models can use other architectures in lieu of or in addition to neural networks. Example machine-learned models can include decision tree based models, support vector machines, hidden Markov models, Bayesian networks, linear regression models, k-means clustering models, etc.

[0355] Machine-learned model(s) 1 can be or include any one of or any part of machine-learned models referenced with respect to the preceding figures (e.g., models 100, 700, etc.). For example, any one or multiple of machine-learned models 100, 700, etc. can be a machine-learned model 1. Features and variations described herein with respect to machine- learned model 1 are to be understood as describing features and variations of any of the machine-learned models described herein. Where this description references machine-learned model 1 it is to be understood that implementations of each of the other models described herein are implicitly referenced and represented thereby.

[0356] Example neural networks can include feed-forward neural networks, recurrent neural networks (RNNs), including long short-term memory (LSTM) based recurrent neural networks, convolutional neural networks (CNNs), diffusion models, generative-adversarialnetworks, or other forms of neural networks. Example neural networks can be deep neural networks. Some example machine-learned models can leverage an attention mechanism such as self-attention. For example, some example machine-learned models can include multi- headed self-attention components. For example, some example machine-learned models can include multi-query self-attention components, such as an M-MQA block as described herein.

[0357] Machine-learned model(s) 1 can include a single or multiple instances of the same model configured to operate on data from input(s) 2. Machine-learned model(s) 1 can include an ensemble of different models that can cooperatively interact to process data from input(s) 2. For example, machine-learned model(s) 1 can employ a mixture-of-experts structure. See, e.g., Zhou et al., Mixture-of-Experts with Expert Choice Routing, ARXIV:2202.09368v2 (Oct.14, 2022).

[0358] Input(s) 2 can generally include or otherwise represent various types of data. Input(s) 2 can include one type or many different types of data. Output(s) 3 can be data of the same type(s) or of different types of data as compared to input(s) 2. Output(s) 3 can include one type or many different types of data.

[0359] Example data types for input(s) 2 or output(s) 3 include natural language text data, software code data (e.g., source code, object code, machine code, or any other form of computer-readable instructions or programming languages), machine code data (e.g., binary code, assembly code, or other forms of machine-readable instructions that can be executed directly by a computer's central processing unit), assembly code data (e.g., low-level programming languages that use symbolic representations of machine code instructions to program a processing unit), genetic data or other chemical or biochemical data, image data, audio data, audiovisual data, haptic data, biometric data, medical data, financial data, statistical data, geographical data, astronomical data, historical data, sensor data generally (e.g., digital or analog values, such as voltage or other absolute or relative level measurement values from a real or artificial input, such as from an audio sensor, light sensor, displacement sensor, etc.), and the like. Data can be raw or processed and can be in any format or schema.

[0360] In multimodal inputs 2 or outputs 3, example combinations of data types include image data and audio data, image data and natural language data, natural language data and software code data, image data and biometric data, sensor data and medical data, etc. It is to be understood that any combination of data types in an input 2 or an output 3 can be present.

[0361] An example input 2 can include one or multiple data types, such as the example data types noted above. An example output 3 can include one or multiple data types,such as the example data types noted above. The data type(s) of input 2 can be the same as or different from the data type(s) of output 3. It is to be understood that the example data types noted above are provided for illustrative purposes only. Data types contemplated within the scope of the present disclosure are not limited to those examples noted above. Example Machine-Learned Sequence Processing Models

[0362] Figure 17 is a block diagram of an example implementation of an example machine-learned model configured to process sequences of information. For instance, an example implementation of machine-learned model(s) 1 can include machine-learned sequence processing model(s) 4. An example system can pass input(s) 2 to sequence processing model(s) 4. Sequence processing model(s) 4 can include one or more machine- learned components. Sequence processing model(s) 4 can process the data from input(s) 2 to obtain an input sequence 5. Input sequence 5 can include one or more input elements 5-1, 5- 2, ... , 5-M, etc. obtained from input(s) 2. Sequence processing model 4 can process input sequence 5 using prediction layer(s) 6 to generate an output sequence 7. Output sequence 7 can include one or more output elements 7-1, 7-2, ... , 7-N, etc. generated based on input sequence 5. The system can generate output(s) 3 based on output sequence 7.

[0363] Sequence processing model(s) 4 can include one or multiple machine-learned model components configured to ingest, generate, or otherwise reason over sequences of information. For example, some example sequence processing models in the text domain are referred to as “Large Language Models,” or LLMs. See, e.g., PaLM 2 Technical Report, GOOGLE, https: / / ai.google / static / documents / palm2techreport.pdf (n.d.). Other example sequence processing models can operate in other domains, such as image domains, see, e.g., Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, ARXIV:2010.11929v2 (Jun.3, 2021), audio domains, see, e.g., Agostinelli et al., MusicLM: Generating Music From Text, ARXIV:2301.11325v1 (Jan.26, 2023), biochemical domains, see, e.g., Jumper et al., Highly accurate protein structure prediction with AlphaFold, 596 Nature 583 (Aug.26, 2021), by way of example. Sequence processing model(s) 4 can process one or multiple types of data simultaneously. Sequence processing model(s) 4 can include relatively large models (e.g., more parameters, computationally expensive, etc.), relatively small models (e.g., fewer parameters, computationally lightweight, etc.), or both.

[0364] In general, sequence processing model(s) 4 can obtain input sequence 5 using data from input(s) 2. For instance, input sequence 5 can include a representation of data from input(s) 2 in a format understood by sequence processing model(s) 4. One or more machine-learned components of sequence processing model(s) 4 can ingest the data from input(s) 2, parse the data into pieces compatible with the processing architectures of sequence processing model(s) 4 (e.g., via “tokenization”), and project the pieces into an input space associated with prediction layer(s) 6 (e.g., via “embedding”).

[0365] Sequence processing model(s) 4 can ingest the data from input(s) 2 and parse the data into a sequence of elements to obtain input sequence 5. For example, a portion of input data from input(s) 2 can be broken down into pieces that collectively represent the content of the portion of the input data. The pieces can provide the elements of the sequence.

[0366] Elements 5-1, 5-2, ... , 5-M can represent, in some cases, building blocks for capturing or expressing meaningful information in a particular data domain. For instance, the elements can describe “atomic units” across one or more domains. For example, for textual input source(s), the elements can correspond to groups of one or more words or sub-word components, such as sets of one or more characters.

[0367] For example, elements 5-1, 5-2, ... , 5-M can represent tokens obtained using a tokenizer. For instance, a tokenizer can process a given portion of an input source and output a series of tokens (e.g., corresponding to input elements 5-1, 5-2, ... , 5-M) that represent the portion of the input source. Various approaches to tokenization can be used. For instance, textual input source(s) can be tokenized using a byte-pair encoding (BPE) technique. See, e.g., Kudo et al., SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing, PROCEEDINGS OF THE 2018 CONFERENCE ON EMPIRICAL METHODS IN NATURAL LANGUAGE PROCESSING (System Demonstrations), pages 66–71 (October 31–November 4, 2018), https: / / aclanthology.org / D18-2012.pdf. Image-based input source(s) can be tokenized by extracting and serializing patches from an image.

[0368] In general, arbitrary data types can be serialized and processed into input sequence 5. It is to be understood that element(s) 5-1, 5-2, ... , 5-M depicted in Figure 17 can be the tokens or can be the embedded representations thereof.

[0369] Prediction layer(s) 6 can predict one or more output elements 7-1, 7-2, ... , 7- N based on the input elements. Prediction layer(s) 6 can include one or more machine-learned model architectures, such as one or more layers of learned parameters that manipulate and transform the input(s) to extract higher-order meaning from, and relationships between, input element(s) 5-1, 5-2, ... , 5-M. In this manner, for instance, example prediction layer(s) 6 can predict new output element(s) in view of the context provided by input sequence 5.

[0370] Prediction layer(s) 6 can evaluate associations between portions of input sequence 5 and a particular output element. These associations can inform a prediction of the likelihood that a particular output follows the input context. For example, consider the textual snippet, “The carpenter’s toolbox was small and heavy. It was full of ___.” Example prediction layer(s) 6 can identify that “It” refers back to “toolbox” by determining a relationship between the respective embeddings. Example prediction layer(s) 6 can also link “It” to the attributes of the toolbox, such as “small” and “heavy.” Based on these associations, prediction layer(s) 6 can, for instance, assign a higher probability to the word “nails” than to the word “sawdust.”

[0371] A transformer is an example architecture that can be used in prediction layer(s) 4. See, e.g., Vaswani et al., Attention Is All You Need, ARXIV:1706.03762v7 (Aug.2, 2023). A transformer is an example of a machine-learned model architecture that uses an attention mechanism to compute associations between items within a context window. The context window can include a sequence that contains input sequence 5 and potentially one or more output element(s) 7-1, 7-2, ... , 7-N. A transformer block can include one or more attention layer(s) and one or more post-attention layer(s) (e.g., feedforward layer(s), such as a multi-layer perceptron).

[0372] An example attention mechanism can include an M-MQA block as described herein.

[0373] Prediction layer(s) 6 can include other machine-learned model architectures in addition to or in lieu of transformer-based architectures. For example, recurrent neural networks (RNNs) and long short-term memory (LSTM) models can also be used, as well as convolutional neural networks (CNNs). In general, prediction layer(s) 6 can leverage various kinds of artificial neural networks that can understand or generate sequences of information.

[0374] Output sequence 7 can include or otherwise represent the same or different data types as input sequence 5. For instance, input sequence 5 can represent textual data, and output sequence 7 can represent textual data. Input sequence 5 can represent image, audio, or audiovisual data, and output sequence 7 can represent textual data (e.g., describing the image, audio, or audiovisual data). It is to be understood that prediction layer(s) 6, and any other interstitial model components of sequence processing model(s) 4, can be configured to receive a variety of data types in input sequence(s) 5 and output a variety of data types in output sequence(s) 7.

[0375] Output sequence 7 can have various relationships to input sequence 5. Output sequence 7 can be a continuation of input sequence 5. Output sequence 7 can becomplementary to input sequence 5. Output sequence 7 can translate, transform, augment, or otherwise modify input sequence 5. Output sequence 7 can answer, evaluate, confirm, or otherwise respond to input sequence 5. Output sequence 7 can implement (or describe instructions for implementing) an instruction provided via input sequence 5.

[0376] Output sequence 7 can be generated autoregressively. For instance, for some applications, an output of one or more prediction layer(s) 6 can be passed through one or more output layers (e.g., softmax layer) to obtain a probability distribution over an output vocabulary (e.g., a textual or symbolic vocabulary) conditioned on a set of input elements in a context window. In this manner, for instance, output sequence 7 can be autoregressively generated by sampling a likely next output element, adding that element to the context window, and re-generating the probability distribution based on the updated context window, and sampling a likely next output element, and so forth.

[0377] Output sequence 7 can also be generated non-autoregressively. For instance, multiple output elements of output sequence 7 can be predicted together without explicit sequential conditioning on each other. See, e.g., Saharia et al., Non-Autoregressive Machine Translation with Latent Alignments, ARXIV:2004.07437v3 (Nov.16, 2020).

[0378] Output sequence 7 can include one or multiple portions or elements. In an example content generation configuration, output sequence 7 can include multiple elements corresponding to multiple portions of a generated output sequence (e.g., a textual sentence, values of a discretized waveform, computer code, etc.). In an example classification configuration, output sequence 7 can include a single element associated with a classification output. For instance, an output “vocabulary” can include a set of classes into which an input sequence is to be classified. For instance, a vision transformer block can pass latent state information to a multilayer perceptron that outputs a likely class value associated with an input image.

[0379] Figure 18 is a block diagram of an example technique for populating an example input sequence 8. Input sequence 8 can include various functional elements that form part of the model infrastructure, such as an element 8-0 obtained from a task indicator 9 that signals to any model(s) that process input sequence 8 that a particular task is being performed (e.g., to help adapt a performance of the model(s) to that particular task). Input sequence 8 can include various data elements from different data modalities. For instance, an input modality 10-1 can include one modality of data. A data-to-sequence model 11-1 can process data from input modality 10-1 to project the data into a format compatible with input sequence 8 (e.g., one or more vectors dimensioned according to the dimensions of inputsequence 8) to obtain elements 8-1, 8-2, 8-3. Another input modality 10-2 can include a different modality of data. A data-to-sequence model 11-2 can project data from input modality 10-2 into a format compatible with input sequence 8 to obtain elements 8-4, 8-5, 8- 6. Another input modality 10-3 can include yet another different modality of data. A data-to- sequence model 11-3 can project data from input modality 10-3 into a format compatible with input sequence 8 to obtain elements 8-7, 8-8, 8-9.

[0380] Input sequence 8 can be the same as or different from input sequence 5. Input sequence 8 can be a multimodal input sequence that contains elements that represent data from different modalities using a common dimensional representation. For instance, an embedding space can have P dimensions. Input sequence 8 can be configured to contain a plurality of elements that have P dimensions. In this manner, for instance, example implementations can facilitate information extraction and reasoning across diverse data modalities by projecting data into elements in the same embedding space for comparison, combination, or other computations therebetween.

[0381] For example, elements 8-0, ... , 8-9 can indicate particular locations within a multidimensional embedding space. Some elements can map to a set of discrete locations in the embedding space. For instance, elements that correspond to discrete members of a predetermined vocabulary of tokens can map to discrete locations in the embedding space that are associated with those tokens. Other elements can be continuously distributed across the embedding space. For instance, some data types can be broken down into continuously defined portions (e.g., image patches) that can be described using continuously distributed locations within the embedding space.

[0382] In some implementations, the expressive power of the embedding space may not be limited to meanings associated with any particular set of tokens or other building blocks. For example, a continuous embedding space can encode a spectrum of high-order information. An individual piece of information (e.g., a token) can map to a particular point in that space: for instance, a token for the word “dog” can be projected to an embedded value that points to a particular location in the embedding space associated with canine-related information. Similarly, an image patch of an image of a dog on grass can also be projected into the embedding space. In some implementations, the projection of the image of the dog can be similar to the projection of the word “dog” while also having similarity to a projection of the word “grass,” while potentially being different from both. In some implementations, the projection of the image patch may not exactly align with any single projection of a single word. In some implementations, the projection of the image patch can align with acombination of the projections of the words “dog” and “grass.” In this manner, for instance, a high-order embedding space can encode information that can be independent of data modalities in which the information is expressed.

[0383] Task indicator 9 can include a model or model component configured to identify a task being performed and inject, into input sequence 8, an input value represented by element 8-0 that signals which task is being performed. For instance, the input value can be provided as a data type associated with an input modality and projected along with that input modality (e.g., the input value can be a textual task label that is embedded along with other textual data in the input; the input value can be a pixel-based representation of a task that is embedded along with other image data in the input; etc.). The input value can be provided as a data type that differs from or is at least independent from other input(s). For instance, the input value represented by element 8-0 can be a learned within a continuous embedding space.

[0384] Input modalities 10-1, 10-2, and 10-3 can be associated with various different data types (e.g., as described above with respect to input(s) 2 and output(s) 3).

[0385] Data-to-sequence models 11-1, 11-2, and 11-3 can be the same or different from each other. Data-to-sequence models 11-1, 11-2, and 11-3 can be adapted to each respective input modality 10-1, 10-2, and 10-3. For example, a textual data-to-sequence model can subdivide a portion of input text and project the subdivisions into element(s) in input sequence 8 (e.g., elements 8-1, 8-2, 8-3, etc.). An image data-to-sequence model can subdivide an input image and project the subdivisions into element(s) in input sequence 8 (e.g., elements 8-4, 8-5, 8-6, etc.). An arbitrary datatype data-to-sequence model can subdivide an input of that arbitrary datatype and project the subdivisions into element(s) in input sequence 8 (e.g., elements 8-7, 8-8, 8-9, etc.).

[0386] Data-to-sequence models 11-1, 11-2, and 11-3 can form part of machine- learned sequence processing model(s) 4. Data-to-sequence models 11-1, 11-2, and 11-3 can be jointly trained with or trained independently from machine-learned sequence processing model(s) 4. Data-to-sequence models 11-1, 11-2, and 11-3 can be trained end-to-end with machine-learned sequence processing model(s) 4.

[0387] Example data-to-sequence models can include example configurations of a universal inverted bottleneck block, such as a twice-mixed inverted bottleneck block.Example Machine-Learned Model Development Platform

[0388] Figure 19 is a block diagram of an example model development platform 12 that can facilitate creation, adaptation, and refinement of example machine-learned models (e.g., machine-learned model(s) 1, sequence processing model(s) 4, etc.). Model development platform 12 can provide a number of different toolkits that developer systems can employ in the development of new or adapted machine-learned models.

[0389] Model development platform 12 can provide one or more model libraries 13 containing building blocks for new models. Model libraries 13 can include one or more pre- trained foundational models 13-1, which can provide a backbone of processing power across various tasks. Model libraries 13 can include one or more pre-trained expert models 13-2, which can be focused on performance in particular domains of expertise. Model libraries 13 can include various model primitives 13-3, which can provide low-level architectures or components (optionally pre-trained), which can be assembled in various arrangements as desired.

[0390] Model development platform 12 can receive selections of various model components 14. Model development platform 12 can pass selected model components 14 to a workbench 15 that combines selected model components 14 into a development model 16.

[0391] Workbench 15 can facilitate further refinement and adaptation of development model 16 by leveraging a number of different toolkits integrated with model development platform 12. For example, workbench 15 can facilitate alignment of the development model 16 with a desired performance profile on various tasks using a model alignment toolkit 17.

[0392] Model alignment toolkit 17 can provide a number of tools for causing development model 16 to generate outputs aligned with desired behavioral characteristics. Alignment can include increasing an accuracy, precision, recall, etc. of model outputs. Alignment can include enforcing output styles, schema, or other preferential characteristics of model outputs. Alignment can be general or domain-specific. For instance, a pre-trained foundational model 13-1 can begin with an initial level of performance across multiple domains. Alignment of the pre-trained foundational model 13-1 can include improving a performance in a particular domain of information or tasks (e.g., even at the expense of performance in another domain of information or tasks).

[0393] Model alignment toolkit 17 can integrate one or more dataset(s) 17-1 for aligning development model 16. Curated dataset(s) 17-1 can include labeled or unlabeled training data. Dataset(s) 17-1 can be obtained from public domain datasets. Dataset(s) 17-1can be obtained from private datasets associated with one or more developer system(s) for the alignment of bespoke machine-learned model(s) customized for private use-cases.

[0394] Pre-training pipelines 17-2 can include a machine-learned model training workflow configured to update development model 16 over large-scale, potentially noisy datasets. For example, pre-training can leverage unsupervised learning techniques (e.g., de- noising, etc.) to process large numbers of training instances to update model parameters from an initialized state and achieve a desired baseline performance. Pre-training pipelines 17-2 can leverage unlabeled datasets in dataset(s) 17-1 to perform pre-training. Workbench 15 can implement a pre-training pipeline 17-2 to pre-train development model 16.

[0395] Fine-tuning pipelines 17-3 can include a machine-learned model training workflow configured to refine the model parameters of development model 16 with higher- quality data. Fine-tuning pipelines 17-3 can update development model 16 by conducting supervised training with labeled dataset(s) in dataset(s) 17-1. Fine-tuning pipelines 17-3 can update development model 16 by conducting reinforcement learning using reward signals from user feedback signals. Workbench 15 can implement a fine-tuning pipeline 17-3 to fine- tune development model 16.

[0396] Prompt libraries 17-4 can include sets of inputs configured to induce behavior aligned with desired performance criteria. Prompt libraries 17-4 can include few-shot prompts (e.g., inputs providing examples of desired model outputs for prepending to a desired runtime query), chain-of-thought prompts (e.g., inputs providing step-by-step reasoning within the exemplars to facilitate thorough reasoning by the model), and the like.

[0397] Example prompts can be retrieved from an available repository of prompt libraries 17-4. Example prompts can be contributed by one or more developer systems using workbench 15.

[0398] In some implementations, pre-trained or fine-tuned models can achieve satisfactory performance without exemplars in the inputs. For instance, zero-shot prompts can include inputs that lack exemplars. Zero-shot prompts can be within a domain within a training dataset or outside of the training domain(s).

[0399] Prompt libraries 17-4 can include one or more prompt engineering tools. Prompt engineering tools can provide workflows for retrieving or learning optimized prompt values. Prompt engineering tools can facilitate directly learning prompt values (e.g., input element values) based one or more training iterations. Workbench 15 can implement prompt engineering tools in development model 16.

[0400] Prompt libraries 17-4 can include pipelines for prompt generation. For example, inputs can be generated using development model 16 itself or other machine- learned models. In this manner, for instance, a first model can process information about a task and output a input for a second model to process in order to perform a step of the task. The second model can be the same as or different from the first model. Workbench 15 can implement prompt generation pipelines in development model 16.

[0401] Prompt libraries 17-4 can include pipelines for context injection. For instance, a performance of development model 16 on a particular task can improve if provided with additional context for performing the task. Prompt libraries 17-4 can include software components configured to identify desired context, retrieve the context from an external source (e.g., a database, a sensor, etc.), and add the context to the input prompt. Workbench 15 can implement context injection pipelines in development model 16.

[0402] Although various training examples described herein with respect to model development platform 12 refer to “pre-training” and “fine-tuning,” it is to be understood that model alignment toolkit 17 can generally support a wide variety of training techniques adapted for training a wide variety of machine-learned models. Example training techniques can correspond to the example training method 1500 described above.

[0403] Model development platform 12 can include a model plugin toolkit 18. Model plugin toolkit 18 can include a variety of tools configured for augmenting the functionality of a machine-learned model by integrating the machine-learned model with other systems, devices, and software components. For instance, a machine-learned model can use tools to increase performance quality where appropriate. For instance, deterministic tasks can be offloaded to dedicated tools in lieu of probabilistically performing the task with an increased risk of error. For instance, instead of autoregressively predicting the solution to a system of equations, a machine-learned model can recognize a tool to call for obtaining the solution and pass the system of equations to the appropriate tool. The tool can be a traditional system of equations solver that can operate deterministically to resolve the system of equations. The output of the tool can be returned in response to the original query. In this manner, tool use can allow some example models to focus on the strengths of machine-learned models—e.g., understanding an intent in an unstructured request for a task—while augmenting the performance of the model by offloading certain tasks to a more focused tool for rote application of deterministic algorithms to a well-defined problem.

[0404] Model plugin toolkit 18 can include validation tools 18-1. Validation tools 18- 1 can include tools that can parse and confirm output(s) of a machine-learned model.Validation tools 18-1 can include engineered heuristics that establish certain thresholds applied to model outputs. For example, validation tools 18-1 can ground the outputs of machine-learned models to structured data sources (e.g., to mitigate “hallucinations”).

[0405] Model plugin toolkit 18 can include tooling packages 18-2 for implementing one or more tools that can include scripts or other executable code that can be executed alongside development model 16. Tooling packages 18-2 can include one or more inputs configured to cause machine-learned model(s) to implement the tools (e.g., few-shot prompts that induce a model to output tool calls in the proper syntax, etc.). Tooling packages 18-2 can include, for instance, fine-tuning training data for training a model to use a tool.

[0406] Model plugin toolkit 18 can include interfaces for calling external application programming interfaces (APIs) 18-3. For instance, in addition to or in lieu of implementing tool calls or tool code directly with development model 16, development model 16 can be aligned to output instruction that initiate API calls to send or obtain data via external systems.

[0407] Model plugin toolkit 18 can integrate with prompt libraries 17-4 to build a catalog of available tools for use with development model 16. For instance, a model can receive, in an input, a catalog of available tools, and the model can generate an output that selects a tool from the available tools and initiates a tool call for using the tool.

[0408] Model development platform 12 can include a computational optimization toolkit 19 for optimizing a computational performance of development model 16. For instance, tools for model compression 19-1 can allow development model 16 to be reduced in size while maintaining a desired level of performance. For instance, model compression 19-1 can include quantization workflows, weight pruning and sparsification techniques, etc. Tools for hardware acceleration 19-2 can facilitate the configuration of the model storage and execution formats to operate optimally on different hardware resources. For instance, hardware acceleration 19-2 can include tools for optimally sharding models for distributed processing over multiple processing units for increased bandwidth, lower unified memory requirements, etc. Tools for distillation 19-3 can provide for the training of lighter-weight models based on the knowledge encoded in development model 16. For instance, development model 16 can be a highly performant, large machine-learned model optimized using model development platform 12. To obtain a lightweight model for running in resource-constrained environments, a smaller model can be a “student model” that learns to imitate development model 16 as a “teacher model.” In this manner, for instance, the investment in learning the parameters and configurations of development model 16 can be efficiently transferred to a smaller model for more efficient inference.

[0409] Workbench 15 can implement one, multiple, or none of the toolkits implemented in model development platform 12. Workbench 15 can output an output model 20 based on development model 16. Output model 20 can be a deployment version of development model 16. Output model 20 can be a development or training checkpoint of development model 16. Output model 20 can be a distilled, compressed, or otherwise optimized version of development model 16.

[0410] Figure 20 is a block diagram of an example training flow for training a machine-learned development model 16. One or more portion(s) of the example training flow can be implemented by a computing system that includes one or more computing devices such as, for example, computing systems described with reference to the other figures. Each respective portion of the example training flow can be performed by any (or any combination) of one or more computing devices. Moreover, one or more portion(s) of the example training flow can be implemented on the hardware components of the device(s) described herein, for example, to train one or more systems or models. Figure 20 depicts elements performed in a particular order for purposes of illustration and discussion. Those of ordinary skill in the art, using the disclosures provided herein, will understand that the elements of any of the methods discussed herein can be adapted, rearranged, expanded, omitted, combined, or modified in various ways without deviating from the scope of the present disclosure. Figure 20 is described with reference to elements / terms described with respect to other systems and figures for exemplary illustrated purposes and is not meant to be limiting. One or more portions of the example training flow can be performed additionally, or alternatively, by other systems.

[0411] Initially, development model 16 can persist in an initial state as an initialized model 21. Development model 16 can be initialized with weight values. Initial weight values can be random or based on an initialization schema. Initial weight values can be based on prior pre-training for the same or for a different model.

[0412] Initialized model 21 can undergo pre-training in a pre-training stage 22. Pre- training stage 22 can be implemented using one or more pre-training pipelines 17-2 over data from dataset(s) 17-1. Pre-training can be omitted, for example, if initialized model 21 is already pre-trained (e.g., development model 16 contains, is, or is based on a pre-trained foundational model or an expert model).

[0413] Pre-trained model 23 can then be a new version of development model 16, which can persist as development model 16 or as a new development model. Pre-trained model 23 can be the initial state if development model 16 was already pre-trained. Pre-trainedmodel 23 can undergo fine-tuning in a fine-tuning stage 24. Fine-tuning stage 24 can be implemented using one or more fine-tuning pipelines 17-3 over data from dataset(s) 17-1. Fine-tuning can be omitted, for example, if a pre-trained model as satisfactory performance, if the model was already fine-tuned, or if other tuning approaches are preferred.

[0414] Fine-tuned model 29 can then be a new version of development model 16, which can persist as development model 16 or as a new development model. Fine-tuned model 29 can be the initial state if development model 16 was already fine-tuned. Fine-tuned model 29 can undergo refinement with user feedback 26. For instance, refinement with user feedback 26 can include reinforcement learning, optionally based on human feedback from human users of fine-tuned model 25. As reinforcement learning can be a form of fine-tuning, it is to be understood that fine-tuning stage 24 can subsume the stage for refining with user feedback 26. Refinement with user feedback 26 can produce a refined model 27. Refined model 27 can be output to downstream system(s) 28 for deployment or further development.

[0415] In some implementations, computational optimization operations can be applied before, during, or after each stage. For instance, initialized model 21 can undergo computational optimization 29-1 (e.g., using computational optimization toolkit 19) before pre-training stage 22. Pre-trained model 23 can undergo computational optimization 29-2 (e.g., using computational optimization toolkit 19) before fine-tuning stage 24. Fine-tuned model 25 can undergo computational optimization 29-3 (e.g., using computational optimization toolkit 19) before refinement with user feedback 26. Refined model 27 can undergo computational optimization 29-4 (e.g., using computational optimization toolkit 19) before output to downstream system(s) 28. Computational optimization(s) 29-1, ... , 29-4 can all be the same, all be different, or include at least some different optimization techniques. Example Machine-Learned Model Inference System

[0416] Figure 21 is a block diagram of an inference system for operating one or more machine-learned model(s) 1 to perform inference (e.g., for training, for deployment, etc.). A model host 31 can receive machine-learned model(s) 1. Model host 31 can host one or more model instance(s) 31-1, which can be one or multiple instances of one or multiple models. Model host 31 can host model instance(s) 31-1 using available compute resources 31-2 associated with model host 31.

[0417] Model host 31 can perform inference on behalf of one or more client(s) 32. Client(s) 32 can transmit an input request 33 to model host 31. Using input request 33, modelhost 31 can obtain input(s) 2 for input to machine-learned model(s) 1. Machine-learned model(s) 1 can process input(s) 2 to generate output(s) 3. Using output(s) 3, model host 31 can return an output payload 34 for responding to input request 33 from client(s) 32. Output payload 34 can include or be based on output(s) 3.

[0418] Model host 31 can leverage various other resources and tools to augment the inference task. For instance, model host 31 can communicate with tool interfaces 35 to facilitate tool use by model instance(s) 31-1. Tool interfaces 35 can include local or remote APIs. Tool interfaces 35 can include integrated scripts or other software functionality. Model host 31 can engage online learning interface(s) 36 to facilitate ongoing improvements to machine-learned model(s) 1. For instance, online learning interface(s) 36 can be used within reinforcement learning loops to retrieve user feedback on inferences served by model host 31. Model host 31 can access runtime data source(s) 37 for augmenting input(s) 2 with additional contextual information. For instance, runtime data source(s) 37 can include a knowledge graph 37-1 that facilitates structured information retrieval for information associated with input request(s) 33 (e.g., a search engine service). Runtime data source(s) 37 can include public or private, external or local database(s) 37-2 that can store information associated with input request(s) 33 for augmenting input(s) 2. Runtime data source(s) 37 can include account data 37-3 which can be retrieved in association with a user account corresponding to a client 32 for customizing the behavior of model host 31 accordingly.

[0419] Model host 31 can be implemented by one or multiple computing devices or systems. Client(s) 2 can be implemented by one or multiple computing devices or systems, which can include computing devices or systems shared with model host 31.

[0420] For example, model host 31 can operate on a server system that provides a machine-learning service to client device(s) that operate client(s) 32 (e.g., over a local or wide-area network). Client device(s) can be end-user devices used by individuals. Client device(s) can be server systems that operate client(s) 32 to provide various functionality as a service to downstream end-user devices.

[0421] In some implementations, model host 31 can operate on a same device or system as client(s) 32. Model host 31 can be a machine-learning service that runs on-device to provide machine-learning functionality to one or multiple applications operating on a client device, which can include an application implementing client(s) 32. Model host 31 can be a part of a same application as client(s) 32. For instance, model host 31 can be a subroutine or method implemented by one part of an application, and client(s) 32 can be another subroutine or method that engages model host 31 to perform inference functions within the application.It is to be understood that model host 31 and client(s) 32 can have various different configurations.

[0422] Model instance(s) 31-1 can include one or more machine-learned models that are available for performing inference. Model instance(s) 31-1 can include weights or other model components that are stored on in persistent storage, temporarily cached, or loaded into high-speed memory. Model instance(s) 31-1 can include multiple instance(s) of the same model (e.g., for parallel execution of more requests on the same model). Model instance(s) 31-1 can include instance(s) of different model(s). Model instance(s) 31-1 can include cached intermediate states of active or inactive model(s) used to accelerate inference of those models. For instance, an inference session with a particular model may generate significant amounts of computational results that can be re-used for future inference runs (e.g., using a KV cache for transformer-based models). These computational results can be saved in association with that inference session so that session can be executed more efficiently when resumed.

[0423] Compute resource(s) 31-2 can include one or more processors (central processing units, graphical processing units, tensor processing units, machine-learning accelerators, etc.) connected to one or more memory devices. Compute resource(s) 31-2 can include a dynamic pool of available resources shared with other processes. Compute resource(s) 31-2 can include memory devices large enough to fit an entire model instance in a single memory instance. Compute resource(s) 31-2 can also shard model instance(s) across multiple memory devices (e.g., using data parallelization or tensor parallelization, etc.). This can be done to increase parallelization or to execute a large model using multiple memory devices which individually might not be able to fit the entire model into memory.

[0424] Input request 33 can include data for input(s) 2. Model host 31 can process input request 33 to obtain input(s) 2. Input(s) 2 can be obtained directly from input request 33 or can be retrieved using input request 33. Input request 33 can be submitted to model host 31 via an API.

[0425] Model host 31 can perform inference over batches of input requests 33 in parallel. For instance, a model instance 31-1 can be configured with an input structure that has a batch dimension. Separate input(s) 2 can be distributed across the batch dimension (e.g., rows of an array). The separate input(s) 2 can include completely different contexts. The separate input(s) 2 can be multiple inference steps of the same task. The separate input(s) 2 can be staggered in an input structure, such that any given inference cycle can be operating on different portions of the respective input(s) 2. In this manner, for instance, model host 31can perform inference on the batch in parallel, such that output(s) 3 can also contain the batch dimension and return the inference results for the batched input(s) 2 in parallel. In this manner, for instance, batches of input request(s) 33 can be processed in parallel for higher throughput of output payload(s) 34.

[0426] Output payload 34 can include or be based on output(s) 3 from machine- learned model(s) 1. Model host 31 can process output(s) 3 to obtain output payload 34. This can include chaining multiple rounds of inference (e.g., iteratively, recursively, across the same model(s) or different model(s)) to arrive at a final output for a task to be returned in output payload 34. Output payload 34 can be transmitted to client(s) 32 via an API.

[0427] Online learning interface(s) 36 can facilitate reinforcement learning of machine-learned model(s) 1. Online learning interface(s) 36 can facilitate reinforcement learning with human feedback (RLHF). Online learning interface(s) 36 can facilitate federated learning of machine-learned model(s) 1.

[0428] Model host 31 can execute machine-learned model(s) 1 to perform inference for various tasks using various types of data. For example, various different input(s) 2 and output(s) 3 can be used for various different tasks. In some implementations, input(s) 2 can be or otherwise represent image data. Machine-learned model(s) 1 can process the image data to generate an output. As an example, machine-learned model(s) 1 can process the image data to generate an image recognition output (e.g., a recognition of the image data, a latent embedding of the image data, an encoded representation of the image data, a hash of the image data, etc.). As another example, machine-learned model(s) 1 can process the image data to generate an image segmentation output. As another example, machine-learned model(s) 1 can process the image data to generate an image classification output. As another example, machine-learned model(s) 1 can process the image data to generate an image data modification output (e.g., an alteration of the image data, etc.). As another example, machine- learned model(s) 1 can process the image data to generate an encoded image data output (e.g., an encoded and / or compressed representation of the image data, etc.). As another example, machine-learned model(s) 1 can process the image data to generate an upscaled image data output. As another example, machine-learned model(s) 1 can process the image data to generate a prediction output.

[0429] In some implementations, the task is a computer vision task. In some cases, input(s) 2 includes pixel data for one or more images and the task is an image processing task. For example, the image processing task can be image classification, where the output is a set of scores, each score corresponding to a different object class and representing thelikelihood that the one or more images depict an object belonging to the object class. The image processing task may be object detection, where the image processing output identifies one or more regions in the one or more images and, for each region, a likelihood that region depicts an object of interest. As another example, the image processing task can be image segmentation, where the image processing output defines, for each pixel in the one or more images, a respective likelihood for each category in a predetermined set of categories. For example, the set of categories can be foreground and background. As another example, the set of categories can be object classes. As another example, the image processing task can be depth estimation, where the image processing output defines, for each pixel in the one or more images, a respective depth value. As another example, the image processing task can be motion estimation, where the network input includes multiple images, and the image processing output defines, for each pixel of one of the input images, a motion of the scene depicted at the pixel between the images in the network input.

[0430] In some implementations, input(s) 2 can be or otherwise represent natural language data. Machine-learned model(s) 1 can process the natural language data to generate an output. As an example, machine-learned model(s) 1 can process the natural language data to generate a language encoding output. As another example, machine-learned model(s) 1 can process the natural language data to generate a latent text embedding output. As another example, machine-learned model(s) 1 can process the natural language data to generate a translation output. As another example, machine-learned model(s) 1 can process the natural language data to generate a classification output. As another example, machine-learned model(s) 1 can process the natural language data to generate a textual segmentation output. As another example, machine-learned model(s) 1 can process the natural language data to generate a semantic intent output. As another example, machine-learned model(s) 1 can process the natural language data to generate an upscaled text or natural language output (e.g., text or natural language data that is higher quality than the input text or natural language, etc.). As another example, machine-learned model(s) 1 can process the natural language data to generate a prediction output (e.g., one or more predicted next portions of natural language content).

[0431] In some implementations, input(s) 2 can be or otherwise represent speech data (e.g., data describing spoken natural language, such as audio data, textual data, etc.). Machine-learned model(s) 1 can process the speech data to generate an output. As an example, machine-learned model(s) 1 can process the speech data to generate a speech recognition output. As another example, machine-learned model(s) 1 can process the speechdata to generate a speech translation output. As another example, machine-learned model(s) 1 can process the speech data to generate a latent embedding output. As another example, machine-learned model(s) 1 can process the speech data to generate an encoded speech output (e.g., an encoded and / or compressed representation of the speech data, etc.). As another example, machine-learned model(s) 1 can process the speech data to generate an upscaled speech output (e.g., speech data that is higher quality than the input speech data, etc.). As another example, machine-learned model(s) 1 can process the speech data to generate a textual representation output (e.g., a textual representation of the input speech data, etc.). As another example, machine-learned model(s) 1 can process the speech data to generate a prediction output.

[0432] In some implementations, input(s) 2 can be or otherwise represent latent encoding data (e.g., a latent space representation of an input, etc.). Machine-learned model(s) 1 can process the latent encoding data to generate an output. As an example, machine- learned model(s) 1 can process the latent encoding data to generate a recognition output. As another example, machine-learned model(s) 1 can process the latent encoding data to generate a reconstruction output. As another example, machine-learned model(s) 1 can process the latent encoding data to generate a search output. As another example, machine- learned model(s) 1 can process the latent encoding data to generate a reclustering output. As another example, machine-learned model(s) 1 can process the latent encoding data to generate a prediction output.

[0433] In some implementations, input(s) 2 can be or otherwise represent statistical data. Statistical data can be, represent, or otherwise include data computed and / or calculated from some other data source. Machine-learned model(s) 1 can process the statistical data to generate an output. As an example, machine-learned model(s) 1 can process the statistical data to generate a recognition output. As another example, machine-learned model(s) 1 can process the statistical data to generate a prediction output. As another example, machine- learned model(s) 1 can process the statistical data to generate a classification output. As another example, machine-learned model(s) 1 can process the statistical data to generate a segmentation output. As another example, machine-learned model(s) 1 can process the statistical data to generate a visualization output. As another example, machine-learned model(s) 1 can process the statistical data to generate a diagnostic output.

[0434] In some implementations, input(s) 2 can be or otherwise represent sensor data. Machine-learned model(s) 1 can process the sensor data to generate an output. As an example, machine-learned model(s) 1 can process the sensor data to generate a recognitionoutput. As another example, machine-learned model(s) 1 can process the sensor data to generate a prediction output. As another example, machine-learned model(s) 1 can process the sensor data to generate a classification output. As another example, machine-learned model(s) 1 can process the sensor data to generate a segmentation output. As another example, machine-learned model(s) 1 can process the sensor data to generate a visualization output. As another example, machine-learned model(s) 1 can process the sensor data to generate a diagnostic output. As another example, machine-learned model(s) 1 can process the sensor data to generate a detection output.

[0435] In some implementations, machine-learned model(s) 1 can be configured to perform a task that includes encoding input data for reliable and / or efficient transmission or storage (and / or corresponding decoding). For example, the task may be an audio compression task. The input may include audio data and the output may comprise compressed audio data. In another example, the input includes visual data (e.g. one or more images or videos), the output comprises compressed visual data, and the task is a visual data compression task. In another example, the task may comprise generating an embedding for input data (e.g. input audio or visual data). In some cases, the input includes audio data representing a spoken utterance and the task is a speech recognition task. The output may comprise a text output which is mapped to the spoken utterance. In some cases, the task comprises encrypting or decrypting input data. In some cases, the task comprises a microprocessor performance task, such as branch prediction or memory address translation.

[0436] In some implementations, the task is a generative task, and machine-learned model(s) 1 can be configured to output content generated in view of input(s) 2. For instance, input(s) 2 can be or otherwise represent data of one or more modalities that encodes context for generating additional content.

[0437] In some implementations, the task can be a text completion task. Machine- learned model(s) 1 can be configured to process input(s) 2 that represent textual data and to generate output(s) 3 that represent additional textual data that completes a textual sequence that includes input(s) 2. For instance, machine-learned model(s) 1 can be configured to generate output(s) 3 to complete a sentence, paragraph, or portion of text that follows from a portion of text represented by input(s) 2.

[0438] In some implementations, the task can be an instruction following task. Machine-learned model(s) 1 can be configured to process input(s) 2 that represent instructions to perform a function and to generate output(s) 3 that advance a goal of satisfying the instruction function (e.g., at least a step of a multi-step procedure to perform thefunction). Output(s) 3 can represent data of the same or of a different modality as input(s) 2. For instance, input(s) 2 can represent textual data (e.g., natural language instructions for a task to be performed) and machine-learned model(s) 1 can process input(s) 2 to generate output(s) 3 that represent textual data responsive to the instructions (e.g., natural language responses, programming language responses, machine language responses, etc.). Input(s) 2 can represent image data (e.g., image-based instructions for a task to be performed, optionally accompanied by textual instructions) and machine-learned model(s) 1 can process input(s) 2 to generate output(s) 3 that represent textual data responsive to the instructions (e.g., natural language responses, programming language responses, machine language responses, etc.). One or more output(s) 3 can be iteratively or recursively generated to sequentially process and accomplish steps toward accomplishing the requested functionality. For instance, an initial output can be executed by an external system or be processed by machine-learned model(s) 1 to complete an initial step of performing a function. Multiple steps can be performed, with a final output being obtained that is responsive to the initial instructions.

[0439] In some implementations, the task can be a question answering task. Machine- learned model(s) 1 can be configured to process input(s) 2 that represent a question to answer and to generate output(s) 3 that advance a goal of returning an answer to the question (e.g., at least a step of a multi-step procedure to perform the function). Output(s) 3 can represent data of the same or of a different modality as input(s) 2. For instance, input(s) 2 can represent textual data (e.g., natural language instructions for a task to be performed) and machine- learned model(s) 1 can process input(s) 2 to generate output(s) 3 that represent textual data responsive to the question (e.g., natural language responses, programming language responses, machine language responses, etc.). Input(s) 2 can represent image data (e.g., image-based instructions for a task to be performed, optionally accompanied by textual instructions) and machine-learned model(s) 1 can process input(s) 2 to generate output(s) 3 that represent textual data responsive to the question (e.g., natural language responses, programming language responses, machine language responses, etc.). One or more output(s) 3 can be iteratively or recursively generated to sequentially process and accomplish steps toward answering the question. For instance, an initial output can be executed by an external system or be processed by machine-learned model(s) 1 to complete an initial step of obtaining an answer to the question (e.g., querying a database, performing a computation, executing a script, etc.). Multiple steps can be performed, with a final output being obtained that is responsive to the question.

[0440] In some implementations, the task can be an image generation task. Machine- learned model(s) 1 can be configured to process input(s) 2 that represent context regarding a desired portion of image content. The context can include text data, image data, audio data, etc. Machine-learned model(s) 1 can be configured to generate output(s) 3 that represent image data that depicts imagery related to the context. For instance, machine-learned model(s) 1 can be configured to generate pixel data of an image. Values for channel(s) associated with the pixels in the pixel data can be selected based on the context (e.g., based on a probability determined based on the context).

[0441] In some implementations, the task can be an audio generation task. Machine- learned model(s) 1 can be configured to process input(s) 2 that represent context regarding a desired portion of audio content. The context can include text data, image data, audio data, etc. Machine-learned model(s) 1 can be configured to generate output(s) 3 that represent audio data related to the context. For instance, machine-learned model(s) 1 can be configured to generate waveform data in the form of an image (e.g., a spectrogram). Values for channel(s) associated with pixels of the image can be selected based on the context. Machine- learned model(s) 1 can be configured to generate waveform data in the form of a sequence of discrete samples of a continuous waveform. Values of the sequence can be selected based on the context (e.g., based on a probability determined based on the context).

[0442] In some implementations, the task can be a data generation task. Machine- learned model(s) 1 can be configured to process input(s) 2 that represent context regarding a desired portion of data (e.g., data from various data domains, such as sensor data, image data, multimodal data, statistical data, etc.). The desired data can be, for instance, synthetic data for training other machine-learned models. The context can include arbitrary data type(s). Machine-learned model(s) 1 can be configured to generate output(s) 3 that represent data that aligns with the desired data. For instance, machine-learned model(s) 1 can be configured to generate data values for populating a dataset. Values for the data object(s) can be selected based on the context (e.g., based on a probability determined based on the context). Example Computing Systems and Devices

[0443] Figure 22 is a block diagram of an example networked computing system that can perform aspects of example implementations of the present disclosure. The system can include a number of computing devices and systems that are communicatively coupled over a network 49. An example computing device 50 is described to provide an example of a computing device that can perform any aspect of the present disclosure (e.g., implementingmodel host 31, client(s) 32, or both). An example server computing system 60 is described as an example of a server computing system that can perform any aspect of the present disclosure (e.g., implementing model host 31, client(s) 32, or both). Computing device 50 and server computing system(s) 60 can cooperatively interact (e.g., over network 49) to perform any aspect of the present disclosure (e.g., implementing model host 31, client(s) 32, or both). Model development platform system 70 is an example system that can host or serve model development platform(s) 12 for development of machine-learned models. Third-party system(s) 80 are example system(s) with which any of computing device 50, server computing system(s) 60, or model development platform system(s) 70 can interact in the performance of various aspects of the present disclosure (e.g., engaging third-party tools, accessing third-party databases or other resources, etc.).

[0444] Network 49 can be any type of communications network, such as a local area network (e.g., intranet), wide area network (e.g., Internet), or some combination thereof and can include any number of wired or wireless links. In general, communication over network 49 can be carried via any type of wired or wireless connection, using a wide variety of communication protocols (e.g., TCP / IP, HTTP, SMTP, FTP), encodings or formats (e.g., HTML, XML), or protection schemes (e.g., VPN, secure HTTP, SSL). Network 49 can also be implemented via a system bus. For instance, one or more devices or systems of Figure 22 can be co-located with, contained by, or otherwise integrated into one or more other devices or systems.

[0445] Computing device 50 can be any type of computing device, such as, for example, a personal computing device (e.g., laptop or desktop), a mobile computing device (e.g., smartphone or tablet), a gaming console or controller, a wearable computing device, an embedded computing device, a server computing device, a virtual machine operating on a host device, or any other type of computing device. Computing device 50 can be a client computing device. Computing device 50 can be an end-user computing device. Computing device 50 can be a computing device of a service provided that provides a service to an end user (who may use another computing device to interact with computing device 50).

[0446] Computing device 50 can include one or more processors 51 and a memory 52. Processor(s) 51 can be any suitable processing device (e.g., a processor core, a microprocessor, an ASIC, an FPGA, a controller, a microcontroller, etc.) and can be one processor or a plurality of processors that are operatively connected. Memory 52 can include one or more non-transitory computer-readable storage media, such as HBM, RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, etc., and combinations thereof.Memory 52 can store data 53 and instructions 54 which can be executed by processor(s) 51 to cause computing device 50 to perform operations. The operations can implement any one or multiple features described herein. The operations can implement example methods and techniques described herein.

[0447] Computing device 50 can also include one or more input components that receive user input. For example, a user input component can be a touch-sensitive component (e.g., a touch-sensitive display screen or a touch pad) that is sensitive to the touch of a user input object (e.g., a finger or a stylus). The touch-sensitive component can serve to implement a virtual keyboard. Other example user input components include a microphone, camera, LIDAR, a physical keyboard or other buttons, or other means by which a user can provide user input.

[0448] Computing device 50 can store or include one or more machine-learned models 55. Machine-learned models 55 can include one or more machine-learned model(s) 1, such as a sequence processing model 4. Machine-learned models 55 can include one or multiple model instance(s) 31-1. Machine-learned model(s) 55 can be received from server computing system(s) 60, model development platform system 70, third party system(s) 80 (e.g., an application distribution platform), or developed locally on computing device 50. Machine-learned model(s) 55 can be loaded into memory 52 and used or otherwise implemented by processor(s) 51. Computing device 50 can implement multiple parallel instances of machine-learned model(s) 55.

[0449] Server computing system(s) 60 can include one or more processors 61 and a memory 62. Processor(s) 61 can be any suitable processing device (e.g., a processor core, a microprocessor, an ASIC, an FPGA, a controller, a microcontroller, etc.) and can be one processor or a plurality of processors that are operatively connected. Memory 62 can include one or more non-transitory computer-readable storage media, such as HBM, RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, etc., and combinations thereof. Memory 62 can store data 63 and instructions 64 which can be executed by processor(s) 61 to cause server computing system(s) 60 to perform operations. The operations can implement any one or multiple features described herein. The operations can implement example methods and techniques described herein.

[0450] In some implementations, server computing system 60 includes or is otherwise implemented by one or multiple server computing devices. In instances in which server computing system 60 includes multiple server computing devices, such server computingdevices can operate according to sequential computing architectures, parallel computing architectures, or some combination thereof.

[0451] Server computing system 60 can store or otherwise include one or more machine-learned models 65. Machine-learned model(s) 65 can be the same as or different from machine-learned model(s) 55. Machine-learned models 65 can include one or more machine-learned model(s) 1, such as a sequence processing model 4. Machine-learned models 65 can include one or multiple model instance(s) 31-1. Machine-learned model(s) 65 can be received from computing device 50, model development platform system 70, third party system(s) 80, or developed locally on server computing system(s) 60. Machine-learned model(s) 65 can be loaded into memory 62 and used or otherwise implemented by processor(s) 61. Server computing system(s) 60 can implement multiple parallel instances of machine-learned model(s) 65.

[0452] In an example configuration, machine-learned models 65 can be included in or otherwise stored and implemented by server computing system 60 to establish a client-server relationship with computing device 50 for serving model inferences. For instance, server computing system(s) 60 can implement model host 31 on behalf of client(s) 32 on computing device 50. For instance, machine-learned models 65 can be implemented by server computing system 60 as a portion of a web service (e.g., remote machine-learned model hosting service, such as an online interface for performing machine-learned model operations over a network on server computing system(s) 60). For instance, server computing system(s) 60 can communicate with computing device 50 over a local intranet or internet connection. For instance, computing device 50 can be a workstation or endpoint in communication with server computing system(s) 60, with implementation of machine-learned models 65 being managed by server computing system(s) 60 to remotely perform inference (e.g., for runtime or training operations), with output(s) returned (e.g., cast, streamed, etc.) to computing device 50. Machine-learned models 65 can work cooperatively or interoperatively with machine- learned models 55 on computing device 50 to perform various tasks.

[0453] Model development platform system(s) 70 can include one or more processors 71 and a memory 72. Processor(s) 71 can be any suitable processing device (e.g., a processor core, a microprocessor, an ASIC, an FPGA, a controller, a microcontroller, etc.) and can be one processor or a plurality of processors that are operatively connected. Memory 72 can include one or more non-transitory computer-readable storage media, such as HBM, RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, etc., and combinations thereof. Memory 72 can store data 73 and instructions 74 which can be executed byprocessor(s) 71 to cause model development platform system(s) 70 to perform operations. The operations can implement any one or multiple features described herein. The operations can implement example methods and techniques described herein. Example operations include the functionality described herein with respect to model development platform 12. This and other functionality can be implemented by developer tool(s) 75.

[0454] Third-party system(s) 80 can include one or more processors 81 and a memory 82. Processor(s) 81 can be any suitable processing device (e.g., a processor core, a microprocessor, an ASIC, an FPGA, a controller, a microcontroller, etc.) and can be one processor or a plurality of processors that are operatively connected. Memory 82 can include one or more non-transitory computer-readable storage media, such as HBM, RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, etc., and combinations thereof. Memory 82 can store data 83 and instructions 84 which can be executed by processor(s) 81 to cause third-party system(s) 80 to perform operations. The operations can implement any one or multiple features described herein. The operations can implement example methods and techniques described herein. Example operations include the functionality described herein with respect to tools and other external resources called when training or performing inference with machine-learned model(s) 1, 4, 16, 20, 55, 65, etc. (e.g., third-party resource(s) 85).

[0455] Figure 22 illustrates one example arrangement of computing systems that can be used to implement the present disclosure. Other computing system configurations can be used as well. For example, in some implementations, one or both of computing system 50 or server computing system(s) 60 can implement all or a portion of the operations of model development platform system 70. For example, computing system 50 or server computing system(s) 60 can implement developer tool(s) 75 (or extensions thereof) to develop, update / train, or refine machine-learned models 1, 4, 16, 20, 55, 65, etc. using one or more techniques described herein with respect to model alignment toolkit 17. In this manner, for instance, computing system 50 or server computing system(s) 60 can develop, update / train, or refine machine-learned models based on local datasets (e.g., for model personalization / customization, as permitted by user data preference selections).

[0456] Figure 23 is a block diagram of an example computing device 98 that performs according to example embodiments of the present disclosure. Computing device 98 can be a user computing device or a server computing device (e.g., computing device 50, server computing system(s) 60, etc.). Computing device 98 can implement model host 31. For instance, computing device 98 can include a number of applications (e.g., applications 1through N). Each application can contain its own machine learning library and machine- learned model(s). For example, each application can include a machine-learned model. Example applications include a text messaging application, an email application, a dictation application, a virtual keyboard application, a browser application, etc. As illustrated in Figure 23, each application can communicate with a number of other components of the computing device, such as, for example, one or more sensors, a context manager, a device state component, or additional components. In some implementations, each application can communicate with each device component using an API (e.g., a public API). In some implementations, the API used by each application is specific to that application.

[0457] Figure 24 is a block diagram of an example computing device 99 that performs according to example embodiments of the present disclosure. Computing device 99 can be the same as or different from computing device 98. Computing device 99 can be a user computing device or a server computing device (e.g., computing device 50, server computing system(s) 60, etc.). Computing device 98 can implement model host 31. For instance, computing device 99 can include a number of applications (e.g., applications 1 through N). Each application can be in communication with a central intelligence layer. Example applications include a text messaging application, an email application, a dictation application, a virtual keyboard application, a browser application, etc. In some implementations, each application can communicate with the central intelligence layer (and model(s) stored therein) using an API (e.g., a common API across all applications).

[0458] The central intelligence layer can include a number of machine-learned models. For example, as illustrated in Figure 24, a respective machine-learned model can be provided for each application and managed by the central intelligence layer. In other implementations, two or more applications can share a single machine-learned model. For example, in some implementations, the central intelligence layer can provide a single model for all of the applications. In some implementations, the central intelligence layer is included within or otherwise implemented by an operating system of computing device 99.

[0459] The central intelligence layer can communicate with a central device data layer. The central device data layer can be a centralized repository of data for computing device 99. As illustrated in Figure 24, the central device data layer can communicate with a number of other components of the computing device, such as, for example, one or more sensors, a context manager, a device state component, or additional components. In some implementations, the central device data layer can communicate with each device component using an API (e.g., a private API).Additional Disclosure

[0460] The technology discussed herein makes reference to servers, databases, software applications, and other computer-based systems, as well as actions taken and information sent to and from such systems. The inherent flexibility of computer-based systems allows for a great variety of possible configurations, combinations, and divisions of tasks and functionality between and among components. For instance, processes discussed herein can be implemented using a single device or component or multiple devices or components working in combination. Databases and applications can be implemented on a single system or distributed across multiple systems. Distributed components can operate sequentially or in parallel.

[0461] While the present subject matter has been described in detail with respect to various specific example embodiments thereof, each example is provided by way of explanation, not limitation of the disclosure. Those skilled in the art, upon attaining an understanding of the foregoing, can readily produce alterations to, variations of, and equivalents to such embodiments. Accordingly, the subject disclosure does not preclude inclusion of such modifications, variations or additions to the present subject matter as would be readily apparent to one of ordinary skill in the art. For instance, features illustrated or described as part of one embodiment can be used with another embodiment to yield a still further embodiment. Thus, it is intended that the present disclosure cover such alterations, variations, and equivalents.

[0462] Aspects of the disclosure have been described in terms of illustrative embodiments thereof. Any and all features in the following claims can be combined or rearranged in any way possible, including combinations of claims not explicitly enumerated in combination together, as the example claim dependencies listed herein should not be read as limiting the scope of possible combinations of features disclosed herein. Accordingly, the scope of the present disclosure is by way of example rather than by way of limitation, and the subject disclosure does not preclude inclusion of such modifications, variations or additions to the present subject matter as would be readily apparent to one of ordinary skill in the art. Moreover, terms are described herein using lists of example elements joined by conjunctions such as “and,” “or,” “but,” etc. It should be understood that such conjunctions are provided for explanatory purposes only. Clauses and other sequences of items joined by a particular conjunction such as “or,” for example, can refer to “and / or,” “at least one of”, “anycombination of” example elements listed therein, etc. Terms such as “based on” should be understood as “based at least in part on.”

[0463] The term “can” should be understood as referring to a possibility of a feature in various implementations and not as prescribing an ability that is necessarily present in every implementation. For example, the phrase “X can perform Y” should be understood as indicating that, in various implementations, X has the potential to be configured to perform Y, and not as indicating that in every instance X must always be able to perform Y. It should be understood that, in various implementations, X might be unable to perform Y and remain within the scope of the present disclosure.

[0464] The term “may” should be understood as referring to a possibility of a feature in various implementations and not as prescribing an ability that is necessarily present in every implementation. For example, the phrase “X may perform Y” should be understood as indicating that, in various implementations, X has the potential to be configured to perform Y, and not as indicating that in every instance X must always be able to perform Y. It should be understood that, in various implementations, X might be unable to perform Y and remain within the scope of the present disclosure.

Claims

WHAT IS CLAIMED IS:

1. A computer-implemented method for using a twice-mixed inverted bottleneck block to efficiently process spatial features in input data arrays, the method comprising: receiving, at an input to the twice-mixed inverted bottleneck block, an input state of a neural network having a first channel depth; convolving, using a first depthwise convolutional layer of the twice-mixed inverted bottleneck block, the input state to obtain a spatially-mixed input state; expanding, using an up-projection pointwise convolutional layer of the twice-mixed inverted bottleneck block, the spatially-mixed input state to generate an expanded state having a second channel depth greater than the first channel depth; convolving, using a second depthwise convolutional layer of the twice-mixed inverted bottleneck block, the expanded state to obtain a spatially-mixed expanded state; projecting, using a down-projection pointwise convolutional layer of the twice-mixed inverted bottleneck block, the spatially-mixed expanded state to obtain a projected state having a third channel depth less than the second channel depth; and using the projected state to perform a prediction task with the neural network.

2. The computer-implemented method of claim 1, wherein the input state represents image data, and wherein the prediction task comprises an image processing task.

3. The computer-implemented method of any of the preceding claims, wherein the neural network comprises a plurality of instances of the twice-mixed inverted bottleneck block in contiguous sequential order.

4. The computer-implemented method of any of the preceding claims, wherein the first depthwise convolutional layer comprises a first kernel size, and wherein the second depthwise convolutional layer comprises a second kernel size, wherein the first kernel size is different than the second kernel size.

5. The computer-implemented method of any of the preceding claims, wherein the first depthwise convolutional layer comprises a first kernel size, and wherein the second depthwise convolutional layer comprises a second kernel size, wherein the first kernel size is equal to the second kernel size.

6. The computer-implemented method of any of the preceding claims, wherein the neural network processes an output of the twice-mixed inverted residual block using a singly-mixed inverted residual bottleneck block, wherein the method comprises: receiving, at an input to the singly-mixed inverted bottleneck block, a downstream input state having a fourth channel depth; expanding, using an up-projection pointwise convolutional layer of the singly-mixed inverted bottleneck block, the downstream input state to generate a downstream expanded state having a fifth channel depth greater than the fourth channel depth; convolving, using a first and only depthwise convolutional layer of the singly-mixed inverted bottleneck block, the downstream expanded state to obtain a spatially-mixed downstream expanded state; and projecting, using a down-projection pointwise convolutional layer of the singly-mixed inverted bottleneck block, the spatially-mixed downstream expanded state to obtain a downstream projected state having a sixth channel depth less than the fifth channel depth.

7. The computer-implemented method of any of the preceding claims, wherein the neural network processes an output of the twice-mixed inverted residual block using a singly-mixed inverted residual bottleneck block, wherein the method comprises: receiving, at an input to the singly-mixed inverted bottleneck block, a downstream input state having a fourth channel depth; convolving, using a first and only depthwise convolutional layer of the singly-mixed inverted bottleneck block, the downstream input state to obtain a spatially-mixed downstream input state; and expanding, using an up-projection pointwise convolutional layer of the singly-mixed inverted bottleneck block, the spatially-mixed downstream input state to generate a downstream expanded state having a fifth channel depth greater than the fourth channel depth; and projecting, using a down-projection pointwise convolutional layer of the singly-mixed inverted bottleneck block, the downstream expanded state to obtain a downstream projected state having a sixth channel depth less than the fifth channel depth.

8. A computer-implemented method of any of the preceding claims, wherein the neural network comprises a transformer encoder, and wherein the twice-mixed inverted bottleneck block processes a downstream output of an attention block.

9. A computer-implemented method of any of the preceding claims, wherein the neural network comprises a transformer encoder, and wherein the twice-mixed inverted bottleneck block processes a downstream output of a multi-query attention block.

10. The computer-implemented method of any of the preceding claims, wherein the neural network comprises at least one attention block that implements multi-query attention over inputs to generate an output.

11. The computer-implemented method of any of the preceding claims, wherein the neural network processes an output of the twice-mixed inverted residual block using at least one attention block that implements multi-query attention over inputs to generate an output.

12. The computer-implemented method of any of the preceding claims, wherein at least one attention block projects the inputs to a projection dimension of at least 48.

13. The computer-implemented method of any of the preceding claims, wherein at least one attention block projects the inputs to a projection dimension of at least 64.

14. A computer-implemented method for performing efficient neural network architecture search, comprising: initializing a search space comprising a plurality of configurations of a universal inverted bottleneck block, wherein the universal bottleneck block comprises: a first placeholder for a depthwise convolutional layer that is discretely activated during search using a first parameter; an up-projection pointwise convolutional layer that follows the first placeholder and expands a channel depth of an internal state of the network to an expanded state; a second placeholder for a depthwise convolutional layer that is discretely activated during search using a second parameter, the second placeholder following the expanding pointwise convolutional layer; and a down-projection pointwise convolutional layer that follows the second placeholder and decreases a channel depth of an internal state of the network from the expanded state to a projected state;constructing a plurality of candidate models, wherein constructing the plurality of candidate models comprises, for each candidate model of the plurality of candidate models: selecting, for each respective network block of a plurality of network blocks of a neural network architecture, a respective value of the first parameter and a respective value of the second parameter to instantiate a respective inverted bottleneck block at the respective network block; training the plurality of candidate models over a training dataset; evaluating performance of the plurality of candidate models; and outputting a selected candidate model based on the evaluating.

15. The computer-implemented method of claim 14, wherein the neural network architecture is configured for image processing tasks, such as image classification or image segmentation, multimodal tasks, such as image and audio processing tasks or image and text processing tasks.

16. The computer-implemented method of any of the preceding claims, wherein the evaluating comprises: determining a prediction quality of the plurality of candidate models on prediction tasks; determining a first latency of execution of the plurality of candidate models on a first hardware configuration; and determining a second latency of execution of the plurality of candidate models on a second hardware configuration.

17. The computer-implemented method of any of the preceding claims, comprising: sampling, probabilistically, the respective value of the first parameter and the respective value of the second parameter using a policy learned based on prior evaluation of the performance of the plurality of candidate models.

18. The computer-implemented method of any of the preceding claims, wherein the policy is learned using a reward value computed using a prediction quality term that is reduced by a latency divergence term, the latency divergence term indicating a measure of divergence of the first latency or the second latency from a latency target.

19. The computer-implemented method of any of the preceding claims, wherein the policy is learned using a reward value computed using a prediction quality term that is reduced by: a first latency divergence term indicating a measure of divergence of the first latency from a first latency target; and a second latency divergence term indicating a measure of divergence of the second latency from a second latency target.

20. The computer-implemented method of any of the preceding claims, wherein determining the first latency comprises estimating execution time on a first processor, and wherein determining the second latency comprises estimating execution time on a second, different processor.

21. The computer-implemented method of any of the preceding claims, wherein the first processor is a processor having a thermal design power (TDP) of less than 5 watts, and wherein the second processor is a different processor having a TDP of less than 5 watts.

22. The computer-implemented method of any of the preceding claims, wherein the first processor is a processor having a thermal design power (TDP) of less than 5 watts, and wherein the second processor is a different processor having a TDP of greater than 5 watts.

23. The computer-implemented method of any of the preceding claims, wherein the first processor is a processor having a thermal design power (TDP) of less than 5 watts, and wherein the second processor is a different processor having a TDP of greater than 5 watts and less than 50 watts.

24. The computer-implemented method of any of the preceding claims, wherein the first processor is a processor having a thermal design power (TDP) of greater than 5 watts and less than 10 watts, and wherein the second processor is a different processor having a TDP of greater than 5 watts and less than 10 watts.

25. The computer-implemented method of any of the preceding claims, wherein the first processor is a processor having a thermal design power (TDP) of less than 10 watts,and wherein the second processor is a different processor having a TDP of greater than 10 watts.

26. The computer-implemented method of any of the preceding claims, wherein the first processor is a processor having a thermal design power (TDP) of less than 10 watts, and wherein the second processor is a different processor having a TDP of greater than 10 watts and less than 50 watts.

27. The computer-implemented method of any of the preceding claims, wherein the first processor is a processor having a thermal design power (TDP) of less than 50 watts, and wherein the second processor is a different processor having a TDP of less than 50 watts.

28. The computer-implemented method of any of the preceding claims, wherein the first processor is a processor having a thermal design power (TDP) of greater than 50 watts, and wherein the second processor is a different processor having a TDP of greater than 50 watts.

29. The computer-implemented method of any of the preceding claims, wherein the first processor or the second processor is an accelerator device specifically adapted for parallelized matrix multiplication and designed to work cooperatively with a central processing unit (CPU) of a computing system.

30. The computer-implemented method of any of the preceding claims, wherein the latency target is less than 20 ms, such as less than 15 ms, such as less than 10 ms, such as less than 5 ms, such as less than 3 ms, such as less than 2 ms, such as less than 1 ms.

31. The computer-implemented method of any of the preceding claims, wherein, for instances of a given network block across the plurality of candidate models: a first shared set of parameters are used for the corresponding up-projection pointwise convolutional layers; and a second shared set of parameters are used for the corresponding down-projection pointwise convolutional layers.

32. The computer-implemented method of any of the preceding claims, wherein, for instances of all network blocks across the plurality of candidate models: a first shared set of parameters are used for the corresponding up-projection pointwise convolutional layers; and a second shared set of parameters are used for the corresponding down-projection pointwise convolutional layers.

33. The computer-implemented method of any of the preceding claims, wherein the plurality of candidate models share greater than 95% of learnable network parameters.

34. The computer-implemented method of any of the preceding claims, wherein the search space comprises an attention block that implements multi-query attention over inputs to generate an output.

35. A computer-implemented method for efficient computation of multi-query attention in a machine-learned sequence processing model, comprising: obtaining a current input tensor representing an input sequence of one or more elements; computing attention values over the current input tensor using one or more tensor contractions, wherein at least a portion of respective tensor contractions of the one or more tensor contractions comprise an output tensor having non-contracting indices inherited from a first tensor and non-contracting indices inherited from a second tensor, wherein: the non-contracting indices inherited from the first tensor are contiguous and ordered as in the first tensor; and the non-contracting indices inherited from the second tensor are contiguous and ordered as in the second tensor; and generating an output of the machine-learned sequence processing model based on the attention values.

36. The computer-implemented method of any of the preceding claims, wherein the portion of respective tensor contractions of the plurality of tensor contractions comprises all the plurality of tensor contractions.

37. The computer-implemented method of any of the preceding claims, wherein: non-contracting indices for the first tensor are adjacent indices of the first tensor; and non-contracting indices for the second tensor are adjacent indices of the second tensor.

38. The computer-implemented method of any of the preceding claims, comprising: reshaping, without transposition, the first tensor into a first matrix by transforming the adjacent indices of the first tensor into a single index; and reshaping, without transposition, the second tensor into a second matrix by transforming the adjacent indices of the second tensor into a single index.

39. A computer-implemented method for efficient computation of multi-query attention in a machine-learned sequence processing model, comprising: obtaining an input tensor representing an input sequence of one or more elements, the input tensor having a sequence length dimension and an embedding dimension; obtaining a query projection tensor configured to generate a projection of the input tensor using a plurality of learned weights, the query projection tensor having a query head dimension, an embedding dimension, and a projection dimension, wherein adjacent indices of the query projection tensor respectively correspond to the query head dimension and the projection dimension; contracting the input tensor and the query projection tensor in the embedding dimension by: reshaping the query projection tensor into a query matrix by combining the query head dimension and the projection dimension into a single dimension; obtaining an input matrix comprising at least a portion of the input tensor; multiplying the input matrix and the query matrix to obtain an output; and reshaping the output into a query tensor having a query head dimension, a sequence length dimension, and a projection dimension, wherein adjacent indices of the query tensor respectively correspond to the query head dimension and the projection dimension; and using the query tensor to attend over the at least one input sequence.

40. The computer implemented method of any of the preceding claims, wherein the contracting of the input tensor and the query projection tensor is performed without transposing any portion of the query projection tensor or any portion of the query tensor.

41. The computer implemented method of any of the preceding claims, wherein the contracting of the input tensor and the query projection tensor is performed without generating an intermediate array in memory.

42. The computer implemented method of any of the preceding claims, comprising: obtaining a second input tensor having a sequence length dimension and an embedding dimension; obtaining a key projection tensor configured to generate a projection of the second input tensor using a plurality of learned weights, the key projection tensor having an embedding dimension and a projection dimension; contracting the second input tensor and the key projection tensor over the embedding dimension to obtain a key tensor; obtaining a value projection tensor configured to generate a projection of the second input tensor using a plurality of learned weights, the value projection tensor having an embedding dimension and a projection dimension; contracting the second input tensor and the value projection tensor over the embedding dimension to obtain a value tensor; contracting the query tensor and the key tensor over the projection dimension to obtain a logit tensor; applying a softmax operation over the logit tensor to obtain a weight tensor; obtaining an output tensor by contracting the weight tensor and the value tensor over the sequence length dimension; and re-projecting the output tensor from the projection dimension to the embedding dimension using an output projection tensor.

43. The computer-implemented method of claim 42, wherein the second input tensor is the first input tensor.

44. The computer-implemented method of any of the preceding claims, comprising: contracting the output tensor and the output projection tensor along the projection dimension; wherein: non-contracting indices of the output tensor are adjacent, and non-contracting indices of the output projection tensor are adjacent.

45. The computer-implemented method of any of the preceding claims, comprising: reshaping at least one dimension of the output tensor into at least two smaller dimensions, wherein the output projection tensor is defined with at least two corresponding decomposed dimensions.

46. The computer-implemented method of any of the preceding claims, wherein the output tensor is reshaped into a four-dimensional tensor, and wherein the output projection tensor is defined as a four-dimensional tensor.

47. The computer-implemented method of any of the preceding claims, comprising: computing the query tensor using a downsampled input state.

48. The computer-implemented method of any of the preceding claims, comprising: computing the query tensor using a downsampled input state obtained using a depthwise convolution having a first stride parameter.

49. The computer-implemented method of any of the preceding claims, comprising: computing the key tensor, the value tensor, or both using a downsampled input state.

50. The computer-implemented method of any of the preceding claims, comprising:computing the key tensor, the value tensor, or both using a downsampled input state obtained using a depthwise convolution having a second stride parameter.

51. The computer-implemented method of any of the preceding claims, comprising: computing the query tensor using a first downsampled input state obtained using a depthwise convolution having a first stride parameter; and computing the key tensor, the value tensor, or both using a second downsampled input state obtained using a depthwise convolution having a second stride parameter; wherein: the first stride parameter is different from the second stride parameter, or the first stride parameter is equal to the second stride parameter.

52. A computer-implemented method for improved training efficiency of machine-learned image processing models, the method comprising: obtaining a distillation dataset that comprises a plurality of training examples labeled by a teacher model, wherein the plurality of training examples comprises: a first plurality of training examples generated using a first set of augmentation routines; and a second plurality of training examples generated using a second set of augmentation routines different from the first set of augmentation routines; generating a training batch for training a student model by combining one or more of the first plurality of training examples and one or more of the second plurality of training examples; processing batchwise the training batch using the student model to obtain a batch of outputs; evaluating the batch of outputs using corresponding labels for the training batch generated by the teacher model; and updating one or more learnable parameters of the student model based on the evaluating.

53. The computer-implemented method of any of the preceding claims, wherein the student model is a supernet implemented by a neural network architecture search, whereinthe supernet comprises multiple different computation pathways corresponding to subnets that represent different candidate neural network architectures.

54. The computer-implemented method of any of the preceding claims, wherein the first set of augmentation routines and the second set of augmentation routines share a common augmentation routine.

55. The computer-implemented method of any of the preceding claims, wherein the common augmentation routine implements random cropping.

56. The computer-implemented method of any of the preceding claims, wherein the first set of augmentation routines induces less distortion than the second set of augmentation routines.

57. The computer-implemented method of any of the preceding claims, wherein the plurality of training examples comprises: a third plurality of training examples generated using a third set of augmentation routines applied to a base dataset different from a base dataset used for the first plurality of training examples or the second plurality of training examples.

58. The computer-implemented method of any of the preceding claims, wherein the third plurality of training examples is class-balanced using a labeling network to classify individual training examples.

59. The computer-implemented method of any of the preceding claims, wherein the third plurality of training examples is class-balanced using a labeling network to classify individual training examples, wherein the third plurality of training examples comprises an equal class population for each class; wherein for a first class, the equal class population is less than a total number of training examples in the first class in the base dataset; and wherein for a second class, the equal class population is greater than a total number of training examples in the second class in the base dataset, the total number of training examples in the second class being supplemented by replicating random training examples in the second class to reach the equal class population.

60. The computer-implemented method of any of the preceding claims, wherein generating the training batch comprises: combining one or more of the first plurality of training examples and one or more of the second plurality of training examples and one or more of the third plurality of training examples.

61. The computer-implemented method of any of the preceding claims, wherein generating the training batch comprises: combining, in equal proportion, one or more of the first plurality of training examples and one or more of the second plurality of training examples and one or more of the third plurality of training examples.

62. The computer-implemented method of any of the preceding claims, wherein generating the training batch comprises: combining one or more of the first plurality of training examples, one or more of the second plurality of training examples, and one or more of the third plurality of training examples; wherein the batch comprises more of the third plurality of training examples than either of the first plurality of training examples or the second plurality of training examples.

63. The computer-implemented method of any of the preceding claims, wherein generating the training batch comprises: combining one or more of the first plurality of training examples and one or more of the second plurality of training examples and one or more of the third plurality of training examples; wherein the batch comprises a number of the third plurality of training examples at least equal to a number of the first plurality of training examples and the second plurality of training examples combined.

64. The computer-implemented method of any of the preceding claims, wherein generating the training batch comprises: combining, in a ratio of 1:1:2, one or more of the first plurality of training examples, one or more of the second plurality of training examples, and one or more of the third plurality of training examples.

65. One or more non-transitory, computer-readable media storing instructions that are executable by one or more processors to cause a computing system to perform operations, the operations comprising: a computer-implemented method of any of the preceding claims.

66. A computing system, comprising: one or more processors; and one or more non-transitory, computer-readable media storing instructions that are executable by the one or more processors to cause the computing system to perform operations, the operations comprising: a computer-implemented method of any of the preceding claims.

67. A neural network configured to implement a computer-implemented method of any of the preceding claims.

Images