Neural networks with nested mixture-of-experts layers

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for processing inputs using a neural network. In particular, the neural network includes one or more nested mixture of experts (MoE) layers that each include a routing layer and a respective set of nested expert layer blocks.

Description

[0001]Neural Networks with Nested Mixture-of-Experts Layers CROSS REFERENCE TO RELATED APPLICATIONS This application claims priority to Indian Patent Application No.202411042404, filed on May 31, 2024, the entire contents of which are hereby incorporated by reference. BACKGROUND This specification relates processing data using machine learning models. As one example, neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to another layer in the network, e.g., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current values of a respective set of weights. SUMMARY This specification describes a system implemented as computer programs on one or more computers in one or more locations that processes a visual input, e.g., an image or a video, using a computer vision neural network to generate a network output for the visual input. More specifically, the computer vision neural network includes one or more mixture of experts (MoE) layers that each include (i) a routing layer and (ii) a plurality of nested expert layer blocks each having respective parameters, e.g., that are a respective subset of a shared set of parameters that is shared between the nested expert layer blocks. The expert layer blocks are referred to as “nested” because, for each nested expert layer other than the largest nested expert layer (i.e., the nested expert layer having the most parameters), the parameters of the nested expert layer are a proper subset of, i.e., less than all of, the parameters of one or more other nested expert layer blocks in the plurality of nested experts. While the above description describes that the neural network is a computer vision neural network and that the sequence of tokens represents a visual input, more generally, the described MoE layers can be inserted within any appropriate neural network that processes a sequence of tokens representing any appropriate type of input. The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following advantages. Visual modalities, e.g., images and videos, contain significant amounts of information redundancy, which can be leveraged to efficiently process inputs of the modality for preforming a variety of computer vision tasks. While Vision Transformer based models and other models that process sequences of tokens extracted from visual inputs scale effectively to large data regimes due to their ability to capture global information, they fail to utilize this inherent redundancy, leading to higher computational costs. This specification describes techniques that address these issues by dynamically choosing experts at the token level when processing a given visual input. Advantageously, the expert for each token is selected from a mixture of nested models falling on an increasing compute-accuracy curve, resulting in superior performance at any given provided compute budget. That is, for each token in a given input to a nested mixture of experts layer, the layer routes the token to one of multiple different nested experts with different compute requirements. In particular, by dynamically prioritizing tokens at inference time, the described techniques can route redundant tokens through cheaper models while routing more important or unique tokens through more expensive models, thus reducing inference-time compute requirements while maintaining or increasing accuracy. For example, the described techniques can reduce compute requirements by up to three-fold relative to a baseline architecture while maintaining or even increasing accuracy. Moreover, the nested experts share parameters, so that the total number of parameters of the computer vision neural network remains comparable to a baseline computer vision neural network that does not include any mixture of experts layers. Thus, the computer vision neural network can be deployed with approximately the same memory requirements as the baseline computer vision neural network. Moreover, in some cases, the capacity bounds for nested expert layer blocks within a given MoE layer can be computed to optimize for the inference compute requirements, e.g., in terms of FLOPs, for a given instance of the computer vision neural network when deployed on particular hardware. Thus, the capacity bounds can be adjusted to optimize the neural network for deployment on different hardware devices with different compute requirements and budgets. The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims. BRIEF DESCRIPTION OF DRAWINGS FIG.1 shows an example computer vision system. FIG.2 is a flow diagram of an example process for processing a visual input using a computer vision neural network FIG.3 is a flow diagram of an example process for processing an input sequence using a nested MoE layer. FIG.4 shows an example of the operation of a nested MoE layer when the nested expert layer blocks include respective instances of self-attention layer blocks. FIG.5 shows another example of the operation of the nested MoE layer when the nested expert layer blocks include respective instances of self-attention layer blocks. FIG.6 shows an example of the performance of the described techniques on an image classification task. FIG.7 shows an example of the performance of the described techniques on a video classification task. Like reference numbers and designations in the various drawings indicate like elements. DETAILED DESCRIPTION FIG.1 shows an example computer vision system 100. The computer vision system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented. The system 100 processes a visual input 102, e.g., an image or a video, using a computer vision neural network 110 to generate a network output 112 for the visual input 102. As used in this specification, processing a visual input 102 generally refers to processing the intensity values of the pixels of the visual input 102, i.e., of the pixels of the image or the pixels of the video frames of the video. The computer vision neural network 110 is a neural network that receives an input sequence 120 of tokens representing a visual input 102 and processes the input sequence tokens is provided as input to the neural network 110, the sequence of tokens can also be referred to as a “network input” to the neural network 110. For example, the output 112 for the visual input 102 can be an embedding of the visual input 102 in an embedding space. An “embedding” as used in this specification is a vector of numeric values, e.g., floating point values or other values, having a pre-determined dimensionality. The space of possible vectors having the pre-determined dimensionality is referred to as the “embedding space.” In other words, the output can be an encoding of the visual input 102. As another example, the output 112 for the visual input 102 can be a classification output for a classification task for the visual input 102. For example, the output can include a respective score for each of a set of object categories, with the score for an object category representing the likelihood that the image or video depicts an object belonging to the object category. As another example, the output can be a text sequence that represents a label for an object category to which an object depicted in the image or video belongs. As another example, the task can be object detection, e.g., open vocabulary object detection, where the output 112 for a given visual input 102 identifies locations of one or more bounding boxes in the visual input 102 and, optionally, for each bounding box, an object category to which an object depicted in each bounding box belongs. As another example, the task can be image segmentation. That is, the neural network can be configured to generate an element-level classification output 112 (e.g., a pixel-level classification output) that includes, for each element in the visual input 102, a respective score corresponding to each of multiple categories. For a given element (e.g., for a given pixel), the score for a category indicates a likelihood that element belongs to the category. In some cases, the categories may be classes of objects, and an element may belong to a category if it is part on an object included in the object class corresponding to the category. As another example, the task can be image depth prediction. In a depth prediction task, the output 112 generated by the neural network identifies, for each pixel in the visual input 102, a predicted depth of the scene at the pixel. As another example, when the visual input 102 is a video, the task can be a video understanding task. That is, the neural network can be configured to process a sequence of video frames to generate an output 112 that characterizes the video frames, e.g., by characterizing whether the video frames depict a person or other agent performing a particular action, by generating a caption that describes the content of the video, by classifying one or more objects in the video and so on. As another example, the task can be a multi-modal text and vision task, e.g., image captioning, where the input is an image and the output 112 is a text caption describing the input image, or visual question answering, where the input is an image or a video and a question about the image or video and the output is an answer to the question. The computer vision neural network 110 can have any appropriate architecture that processes a sequence of tokens representing a visual input 102 to generate the output for the visual input. For example, the neural network 110 can have an architecture that includes multiple self-attention network blocks that perform self-attention to update the tokens in the sequence. Examples of such architectures include Vision Transformers (ViTs), Video Vision Transformers (ViViTs), and other ViT variants. As another example, the neural network 110 can have an architecture that includes multiple network blocks that perform a different type of operation to update the tokens in the sequence. An example of such an architecture is an MLP-mixer architecture in which each block includes a token mixing multi-layer perceptron (MLP) that mixes information across tokens and a channel mixing MLP that mixes information across the channels within each token. More specifically, the system obtains a visual input 102, e.g., an image or a video. The system 100 generates, from the visual input 102, a sequence 120 of tokens. A “token” as used in this specification is vector of numerical values having a fixed dimensionality. For example, the system 100 can generate the sequence 120 of tokens by partitioning the visual input 102 into patches and then processing each patch using an embedding subnetwork, e.g., a linear layer or a multi-layer perceptron, to generate a token representing the patch. The system 100 then processes the sequence 120 of tokens using the computer vision neural network 110 to generate a network output 112 for the visual input 102. Generally, the computer vision neural network 110 includes a sequence of layers 130 that each update the sequence 120 of tokens. The sequence of layers 130 includes one or more mixture of experts (MoE) layers 140 that each include (i) a routing layer 150 and (ii) a plurality of nested expert layer blocks 160 each having respective parameters. The expert layer blocks 160 are referred to as “nested” because, for each nested expert layer block other than the largest nested expert layer block (i.e., the nested expert layer having the most parameters), the parameters of the nested expert layer block are a proper subset of, i.e., less than all of, the parameters of one or more other nested expert layer blocks in the plurality of nested expert layer blocks 160. That is, each of the nested expert layer blocks has a different subset of a shared set of parameters, with the largest nested expert layer blocks having all of the shared set of parameters, and each other nested expert layer block having a different proper subset of the shared set. In some implementations, each nested expert layer block includes all of the parameters of all smaller nested expert layer blocks. During the processing, each MoE layer 140 receives the sequence 120 of tokens and updates the sequence 120 of tokens, by, for each token, processing the token using the routing layer 150 to assign the token to one of the nested expert layer blocks 160 and then processing the token using the assigned nested expert layer 160 to generate an initial updated token. Thus, the MoE layer 140 dynamically assigns each token to a respective expert for processing, with different tokens potentially being assigned to different experts. Because the different experts are nested and therefore have different computational costs, the MoE layer 140 can assign more computational resources for processing “important” tokens to experts with a higher computational cost (as a result of having more parameters) while assigning less important or redundant tokens to experts with a lower computational cost (as a result of having fewer parameters). Because the different experts are nested and therefore share parameters, i.e., each expert’s parameters are a different subset of a shared set of parameters, the computer vision neural network 110 has approximately the same number of parameters as a neural network that includes only the largest expert block for each MoE layer, with the only additional parameters coming from the parameters of the routing layers 150 of the MoE layers. Updating tokens using a “nested” MoE layer is 140 described below with reference to FIGS.2-5. The system 100 can process the initial updated token through any subsequent layers of the computer vision neural network 110 to generate the network output 112. In some implementations, each layer in the sequence of layers 130 is a MoE layer 140. In some other implementations, only a subset of the sequence of layers 130 are MoE layers 140 and the remaining layer(s) are other types of layers, e.g., attention layers with no mixture of experts, dense feed-forward layers, and so on. As a particular example, in a ViVit architecture or other architecture with factorized attention layers, each spatial attention layer can be a MoE layer 140 while each temporal attention layer can be a conventional temporal attention layer. While the above description describes that the neural network is a computer vision neural network and that the sequence of tokens represents a visual input, more generally, the described MoE layers can be inserted within any appropriate neural network that processes a sequence of tokens representing any appropriate type of input. For example, the described MoE layer can be inserted in place of an attention layer block in any appropriate Transformer neural network, e.g., an encoder-only Transformer neural network, a decoder-only Transformer neural network, or an encoder-decoder Transformer neural network. In this example, the input represented by the sequence of tokens can be any appropriate type of input, e.g., a text sequence, an audio signal, a visual input, a tabular or other structured input or a multi-modal sequence that includes multiple modalities of data, e.g., two or more of text, audio, visual, or structured inputs. Moreover, when the MoE layer is included within a decoder of a Transformer neural network or another type of neural network that auto-regressively generates the output tokens in an output sequence one output after the other, the input at any given processing time step can include the most recently generated token at the preceding processing time step and the output can the token that follows the most recently generated token in the output sequence. In some of these cases, the input includes the most recently generated token and the preceding tokens in the output sequence (and, in some cases, a conditioning input) while in others of these cases, the input includes only the most recently generated token and the system caches the outputs from processing previous tokens, e.g., keys and values for attention heads or outputs of layer blocks, necessary to perform the processing of the most recently generated tokens. FIG.2 is a flow diagram of an example process 200 for processing a visual input using a computer vision neural network. For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a computer vision system, e.g., the computer vision system 100 of FIG.1, appropriately programmed in accordance with this specification, can perform the process 200. The system obtains a visual input (step 202). As described above, the visual input can be an image or a video. The system generates, from the visual input, a sequence of tokens (step 204). The system processes the sequence of tokens using a computer vision neural network to generate a network output for the visual input (step 206). As described above, the computer vision neural network includes a sequence of layers that each update the sequence of tokens. That is, each layer receives the sequence of tokens, e.g., after being updated by the preceding layer or, for the first layer in the sequence, the input sequence of tokens, and processes the sequence of tokens to update the sequence of tokens before providing the sequence to the next layer. Generally, the sequence of layers includes one or more nested MoE layers. Thus, processing the sequence of tokens using a computer vision neural network to generate a network output for the visual input generally includes, for each of the nested MoE layers, processing the sequence of tokens, e.g., after having been updated by the preceding layer in the sequence, using the nested MoE layer, to update the sequence of tokens (step 208). As described above, each nested MoE layer generally includes (i) a routing layer and (ii) a plurality of nested expert layer blocks each having respective parameters. The expert layer blocks are “nested” because, for each nested expert layer other than the largest nested expert layer of the plurality of nested expert layer blocks, the parameters of the nested expert layer are a proper subset of the parameters of one or more other nested expert layer blocks in the plurality of nested experts. In other words, for each nested expert layer other than the smallest nested expert layer of the plurality of nested expert layer blocks, the parameters of the nested expert layer include the parameters of one or more smaller nested expert layer blocks in the plurality of nested experts and additional parameters that are not in the smaller nested expert layer blocks. More specifically, each nested expert layer block has a different subset of a shared set of parameters, with the largest nested expert layer having all of the shared set of parameters and each other nested expert layer block having a different proper subset of the shared set of parameters. In some implementations, each nested expert layer block includes all of the parameters of all smaller nested expert layer blocks. Processing an input sequence using a nested MoE layer is described in more detail below with reference to FIGS.3-5. Generally, the computer vision neural network processes the sequence of tokens after being updated by the last layer in the sequence to generate the network output. For example, when the network output includes an embedding of the visual input in an embedding space, the system can process the sequence of tokens after being updated by the last layer in the sequence to generate the embedding. As a particular example of this, the system can apply global average pooling to the sequence of tokens after being updated by the last layer in the sequence. As another example, when the network output includes a classification output for a classification task, the system can process the sequence of tokens after being updated by the last layer in the sequence to generate an embedding of the visual input, e.g., as described above, and then process the embedding to generate the classification output. As one example of generating a classification output, the system can process the embedding using one or more output layers to generate the classification output, e.g., as a probability distribution over a set of possible object classes. As another example of generating a classification output, the system can process the embedding and respective embeddings for each of a plurality of object classes to generate the classification output. For example, the system can compute respective similarity measures, e.g., cosine similarity, dot product similarity, or Euclidean distance, between the embedding of the visual input and each of the respective embeddings for the object classes, and then select the object class having the most similar embedding according to the similarity measures. As another example, the network output can also include, instead of or in addition to the classification output, an object detection output that is generated by processing each token after being updated using the last layer in the sequence using a set of one or more output layers. As another example, the network output can be a segmentation output that is generated by processing each token after being updated using the last layer in the sequence using a set of one or more output layers. FIG.3 is a flow diagram of an example process 300 for processing an input sequence using a nested MoE layer. For convenience, the process 300 will be described as being performed by a system of one or more computers located in one or more locations. For example, a computer vision system, e.g., the computer vision system 100 of FIG.1, appropriately programmed in accordance with this specification, can perform the process 300. The system can perform a respective instance of the process 300 for each nested MoE layer in the computer vision neural network as part processing the input sequence of tokens to generate the network output for the visual input. The system receives a sequence of tokens (step 302). For example, the sequence of tokens can have been updated by the preceding layer in the sequence of layers. As another example, when the nested MoE layer is the first layer in the sequence, the sequence of tokens can be the input sequence of tokens as generated from the visual input. The system then updates the sequence of tokens by performing steps 304, 306, and optionally 308 for each token in the sequence. The system processes the token using the routing layer of the nested MoE layer to assign the token to one of the nested expert layer blocks of the nested MoE layer (step 304). In general, the system processes the token using the routing layer to generate a respective score for each of the nested expert layer blocks and then assigns the token to one of the nested expert layer blocks based on the respective scores. The routing layer can be any appropriate learned function that maps a token to a score distribution, e.g., a probability distribution or a set of logits, over the expert layer blocks. For example, the routing layer can be a single feed-forward neural network layer or can be a multi-layer perceptron (MLP). The function is referred to as “learned” because the parameters of the function are learned during the training of the computer vision neural network. As a result, the routing function learns to generate routing scores that result in more important or informative tokens being more likely to be routed to larger nested experts. In some implementations, to assign the token to one of the nested expert layer blocks based on the respective scores, the system assigns the token to the expert with the highest score or samples an expert in accordance with the scores in the distribution. In some other implementations, the system assigns the token based on the respective scores and on respective capacity bounds for each of the nested expert layer blocks. That is, each expert block can be associated with a respective capacity bound that specifies the maximum number of tokens that can be routed to the expert block. For example, the respective capacity bounds can be values between zero and one and the respective capacity bounds for the expert blocks can sum to one, so that the respective capacity bound for each expert represents the maximum fraction of the tokens in the sequence that can be routed to the expert block. In some cases, the respective capacity bounds are received as input by the system. In some other cases, the system determines the respective capacity bounds for each of the nested expert layer blocks based on a respective model dimension for each of the nested expert layer blocks and an inference compute requirement for the computer vision neural network. As will be described in more detail below, the model dimension for a nested expert defines the size of the subset of the parameters of the nested expert. For example, each model dimension can be associated with a respective number of FLOPs required to process an input of the corresponding model dimension, and the inference compute requirement can be a requirement on the maximum number of FLOPs that is required to generate a network output for a visual input using the computer vision neural network. In this example, the system can set the capacity bounds so that, if each expert in the computer vision neural network processes a number of tokens equal to the capacity bound for the expert, the total number of FLOPs performed to generate an output does not exceed the maximum number while incentivizing the use of larger experts to maintain output quality. For example, the system can set the capacity bounds by performing an optimization that incentivizes the usage of larger experts while staying within the inference FLOPs requirement and also rewards uniformity of capacity across experts. As a particular example, this can be expressed using the following optimization problem: In determined through a hyperparameter search, e.g., (8,4), (12, 6), (10, 2), or (14,8) and use a Sequential Least Squares Programming (SLSQP) algorithm or other optimization algorithm to solve for the capacities ciof the experts given an effective capacity ecfor the experts that is defined based on the FLOP requirement. That is, the effective capacity ec for the experts is defined by the FLOP requirement for the neural network and can be received as input by the system. The system can generally assign the token using the respective scores generated by the router and the respective capacity bounds in any appropriate manner. As a particular example, the system can assign the token based on the respective scores and on respective capacity bounds for each of the nested expert layer blocks through expert preferred routing. Expert preferred routing (EPR) refers to a greedy assignment approach that gives higher preference to larger nested experts, aiming to identify the most important tokens first. To perform EPR, the system orders the experts by size and proceeds through the order starting from the largest expert. For a given expert j, the system assigns to the expert^^^ ൌ ⌊^^^^^⌋ of the remaining tokens that have not yet been assigned to earlier experts in theorder, where ^^^is the capacity bound for the j-th nested expert and N is the total number of tokens. Any remaining tokens, arising from integer packing constraints, are assigned to the smallest expert. A specific example of EPR is described below in Table 1: Table 1 The system processes the token using the assigned nested expert layer block to generate an initial updated token (step 306). As described above, because different experts have different numbers of parameters, the amount of compute, e.g., the number of FLOPs, required to process the token using the assigned nested expert layer block, will differ based on which nested expert layer block the token has been assigned to. As a particular example, each nested expert layer block can include a respective instance of each of one or more linear neural network layers that each have a corresponding weight matrix. In this example, for each nested expert layer block and for each of the one or more linear neural network layers, the respective instance of the linear neural network layer for the expert has a different slice of the corresponding weight matrix for the linear neural network layer. That is, because nested expert layer blocks all share the same weight matrix (but use different slices of the weight matrix), the parameters are shared across all of the nested expert layer blocks and each nested expert layer block has a respective subset of the shared parameters. In this specification, a “slice” of size m of a token refers to a portion of the token that includes only the first D / m values of the token, where m is greater than or equal to one and less than or equal to the total number of values of the token D. In this specification, a “slice” of size m of a D1 x D2 matrix can refer to (i) a sub- matrix that includes only the first D1 / m rows of the matrix, (ii) a sub-matrix that includes only the first D2 / m columns of the matrix, or (iii) a sub-matrix that includes only the first D1 / m rows of the matrix and the first D2 / m columns of the matrix. Thus, because different experts have different “slices” of the same weight matrix, the experts share parameters and no additional parameters are added to the computer vision neural network relative to maintaining only the largest expert that has the entire weight matrix. To perform the processing of the linear neural network layers in this example, for each nested expert layer block and for each of the one or more linear neural network layers, the respective instance of the linear neural network layer is configured to apply the slice of the corresponding weight matrix for the linear neural network layer to an input slice derived from a given token to generate the output of the respective instance of the linear neural network layer for the given token. Thus, different instances of the layer operate on different slices of a corresponding input token and use different slices of the weight matrix of the layer. Thus, for each linear neural network layer that has operations that can be represented by Wx, where W is a weight matrix and x is an input token for the linear neural network layer, and for a given expert that has a slice of the matrix of size m, the output of the respective instance of the layer for the expert can be a partial projection of the output of the layer: ^^^;^^^^^^. ^^^ ^; ^^^ if there are four nested expert layer blocks, one that corresponds to the has m = 1, one that has m = 2, one that has m = 4, and one that has m = 8, for a given weight matrix that is D1 x D, the expert with m = 1 will have the full D1 x D matrix, the expert with m = 2 will have a D1 x D / 2 slice of the full matrix, the expert with m = 4 will have a D1 x D / 4 slice of the full matrix, and the expert with m = 8 will have a D1 x D / 8 slice of the full matrix. Thus, each expert has a respective “nested” slice of the same D1 x D matrix. As a particular example of this, each nested expert layer block can include a respective instance of a self-attention layer block. This example is described in more detail below with reference to FIGS.4 and 5. In some cases, the system uses the initial updated token as the output for the nested MoE layer. In some other cases, the system then performs one or more additional operations on the initial updated token to generate a final updated token for the nested MoE layer (step 308). For example, the system can multiply the initial updated token by a value that is defined by the routing score r for the assigned nested MoE layer. Performing this multiplication can ensure that gradients propagate to the router weights during training. In some cases, the value is equal to the routing score. In some other cases, the value is also based on a learned parameter α that is initialized to zero and that ensures proper gradient flow during early training stages, e.g., when fine-tuning the neural network from a pre-trained neural network. For example, the value can be equal to rα + 1. As another example, after performing the multiplication, the system can apply a residual connection, e.g., with the input to the block or an intermediate output to the block. As another example, the system can apply a normalization operation, e.g., a LayerNorm (LN) operation. FIG.4 shows an example 400 of the operation of a nested MoE layer when the nested expert layer blocks include respective instances of self-attention layer blocks. Generally, in the example 400, each nested expert block is a respective instance of a self-attention layer block that includes (i) a self-attention mechanism that has a set of one or more attention heads and (ii) an MLP layer. In the example 400, each nested block also includes one or more additional types of operations, e.g., LayerNorm (LN), and residual connections (denoted with the ^ symbol). More specifically, in the example 400 each block in the neural network operates on D dimensional tokens, where D is a fixed value that is also referred to as the model dimension. Each of the nested expert blocks has a respective value of m that defines the sizes of the weight matrix slices and the token slices that are extracted from the D dimensional tokens by the “Extract” operation. Thus, the value of m controls the parameter count and the FLOPs of a given block, with smaller values of m having more parameters and more FLOPs because they operate on larger tokens (and therefore use larger slices of the shared weight matrices of the layer block) than larger values of m because they operate on smaller tokens. The “pad” operation shown in the example 400 is used when the D / m tokens for the block need to be padded to have D dimension before certain operations (e.g., residual connections or the application of QKV attention by an attention head). FIG.5 shows an example 500 of the operation of the nested MoE layer when the nested expert layer blocks include respective instances of self-attention layer blocks. In the example 500, the layer with index l of the neural network has four nested expert layer blocks, one that corresponds to the “full” model and has m = 1, one that has m = 2, one that has m = 4, and one that has m = 8. As shown in the example 500, the router can route different tokens to different nested expert layer blocks. In the example 500, one token, i.e., token xi, is routed to a nested expert layer block (referred to as the “nested model”) with model dimension D / 4, whereas the next token, i.e., token xi+1, gets routed to the full expert layer block (referred to as the “full model”) that has model dimension D. Shading in a given weight matrix in FIG.5 indicates a sliced matrix to construct the “nestedness” of different expert blocks. That is, shaded regions of a weight matrix belong to a given slice while un-shaded regions of the weight matrix do not. In particular, in the example 500, the self-attention block includes, for each attention head in a set of one or more attention heads, a set of input linear neural network layers that process an input token to generate a query, key, and value from the input token and a self- attention mechanism that applies self-attention across queries, keys, and values for each of the tokens in the sequence to generate an initial attended token for each of the tokens in the sequence. Thus, each nested expert layer block includes a respective instance of the set of input linear neural network layers and the self-attention mechanism is shared across all of the nested expert layer blocks. In other words, information exchange between different tokens of different dimension happens in the self-attention mechanism, where they are always projected to the same dimension D. To allow the self-attention mechanism to be performed at the dimension D, each nested expert layer block is configured to, prior to providing a query, key, and value for a given token as input to the shared self-attention mechanism, pad the query, key, and value to each have a fixed dimensionality, i.e., D, that is shared across the nested expert layer blocks. “Padding” a token can refer to extending the dimensionality of the token by appending zeros or “noise” values sampled from a specified distribution, e.g., a Gaussian distribution, to the token. In the example 500, the attention mechanism is a “multi-head” attention mechanism and the set of attention heads has a plurality of attention heads. As a result, the self- attention layer block also includes an output linear neural network layer that, for a given token, processes a combined token generated from the initial attended tokens for the given token to generate an attended token. Thus, each nested expert layer block includes a respective instance of the output linear neural network layer. Further, as described above, the self-attention layer block further includes a multi- layer perceptron (MLP) layer. The MLP layer has one or more linear neural network layers and each nested expert layer block includes a respective instance of the MLP layer. Thus, in the example 500, the parameters of each nested expert layer block include (i) a respective slice of a query weight matrix ^^^^, (ii) a respective slice of a key weight matrix ^^^^, (iii) a respective slice of a value weight matrix ^^^^, (iv) a respective slice of a self-attention output weight matrix ^^^^^^^^for combining the outputs of the attention heads, (v) a respective slice of a first weight matrix ^^^^of a first layer of the MLP, and (vi) a respective slice of a second weight matrix ^^^^of a second layer of the MLP. Prior to using the computer vision neural network to generate outputs for new visual inputs, the system 100 or another training system trains the computer vision neural network to determine trained values of the parameters of the neural network, including the parameters of the routing functions and the shared set of parameters of the nested experts of the nested MoE layers. In some cases, the system trains the computer vision neural network from scratch. In some other cases, the system “fine-tunes” the computer vision neural network starting from a pre-trained computer vision neural network, e.g., a neural network that does not include routing function and includes only the largest expert for each nested MoE layer but otherwise has the same architecture as the computer vision neural network. The system can generally train the neural network on any appropriate objective and on any appropriate training data set that is appropriate for the computer vision task the neural network is configured to perform. For example, example data sets for image and video classification are described below with reference to FIGS.6 and 7. FIG.6 is an example 600 of the performance of the described techniques on an image classification task. In particular, FIG.6 shows three graphs a)-c) that each depict the performance of a corresponding model size (S / 16, B / 16, and L / 16) on the ImageNet-21k data set. More specifically, FIG.6 shows the performance of a variant of the described techniques (“MoNE”) relative to three existing techniques (MatViT, ViT, and Mixture-of- Depth) that do not have nested experts but that have similar parameter sizes to the neural network used for the MoNE results. As can be seen from graphs a) and b), on all of the model sizes, the described techniques outperform the existing techniques in terms of accuracy for a given number of FLOPs and achieve comparable accuracy to the ViT approach despite requiring significantly fewer FLOPs. FIG.7 is an example 700 of the performance of the described techniques on a video classification task. In particular, FIG.7 shows two graphs, graph a) showing performance on the Kinetics-400 data set and graph b) showing performance on the Something- Something-v2 data set. More specifically, FIG.7 shows the performance of the described techniques (“MoNE”) relative to two existing techniques (MatViViT and ViViT) that do not have nested experts but that have similar parameter sizes to the neural network used for the MoNE results. As can be seen from graphs a) and b), on both of the data sets, the described techniques outperform the existing techniques in terms of accuracy for a given number of FLOPs and achieve comparable accuracy to the ViViT approach despite requiring significantly fewer FLOPs. This specification uses the term “configured” in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions. Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly- embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus. The term “data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them. A computer program, which may also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it can be deployed in any form, including as a stand alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, sub programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communication network. In this specification, the term “database” is used broadly to refer to any collection of data: the data does not need to be structured in any particular way, or structured at all, and it can be stored on storage devices in one or more locations. Thus, for example, the index database can include multiple collections of data, each of which may be organized and accessed differently. Similarly, in this specification the term “engine” is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers. The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers. Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read only memory or a random access memory or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few. Computer readable media suitable for storing computer program instructions and data include all forms of non volatile memory, media and memory devices, including by way of example semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto optical disks; and CD ROM and DVD-ROM disks. To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return. Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and compute-intensive parts of machine learning training or production, i.e., inference, workloads. Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework or a Jax framework. Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back end, middleware, or front end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet. The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at the server from the device. While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination. Similarly, while operations are correspond toed in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products. Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes correspond toed in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous. Aspects of the present disclosure may be as set out in the following clauses: Clause 1. A method performed by one or more computers, the method comprising: obtaining a visual input; generating, from the visual input, a sequence of tokens; processing the sequence of tokens using a computer vision neural network to generate a network output for the visual input, wherein the computer vision neural network includes a sequence of layers that each update the sequence of tokens, and wherein the sequence of layers includes one or more MoE layers that each include: (i) a routing layer; and (ii) a plurality of nested expert layer blocks each having respective parameters, wherein, for each nested expert layer other than a largest nested expert layer of the plurality of nested expert layer blocks, the parameters of the nested expert layer are a proper subset of the parameters of one or more other nested expert layer blocks in the plurality of nested experts, and wherein processing the network input comprises, for each MoE layer: receiving the sequence of tokens; and updating the sequence of tokens, comprising, for each token: processing the token using the routing layer to assign the token to one of the nested expert layer blocks; and processing the token using the assigned nested expert layer block to generate an initial updated token. Clause 2. The method of clause 1, wherein the visual input is an image. Clause 3. The method of clause 1, wherein the visual input is a video. Clause 4. The method of any preceding clause, wherein processing the sequence of tokens using the computer vision neural network further comprises: processing the sequence of tokens after being updated by a last layer in the sequence to generate the network output. Clause 5. The method of clause 4, wherein the network output comprises an embedding of the visual input in an embedding space, and wherein processing the sequence of tokens after being updated by a last layer in the sequence to generate the network output comprises: processing the sequence of tokens after being updated by a last layer in the sequence to generate the embedding. Clause 6. The method of clause 5, wherein processing the sequence of tokens after being updated by a last layer in the sequence to generate the embedding comprises: applying global average pooling to the sequence of tokens after being updated by a last layer in the sequence. Clause 7. The method of clause 4, wherein the network output comprises a classification output for a classification task. Clause 8. The method of clause 8, wherein processing the sequence of tokens after being updated by a last layer in the sequence to generate the network output comprises: processing the sequence of tokens after being updated by a last layer in the sequence to generate an embedding of the visual input; and processing the embedding to generate the classification output. Clause 9. The method of clause 8, wherein processing the embedding to generate the classification output comprises processing the embedding using one or more output layers to generate the classification output. Clause 10. The method of clause 8, wherein processing the embedding to generate the classification output comprises processing the embedding and respective embeddings for each of a plurality of object classes to generate the classification output. Clause 11. The method of any one of clauses 7-10, wherein the network output further comprises an object detection output. Clause 12. The method of clause 4, wherein the network output is a segmentation output. Clause 13. The method of any preceding clause, wherein each nested expert layer block has a different subset of a shared set of parameters. Clause 14. The method of any preceding clause, wherein: each nested expert layer block comprises a respective instance of each of one or more linear neural network layers that each have a corresponding weight matrix, and for each nested expert layer block and for each of the one or more linear neural network layers, the respective instance of the linear neural network layer has a different slice of the corresponding weight matrix for the linear neural network layer. Clause 15. The method of clause 14, wherein for each nested expert layer block and for each of the one or more linear neural network layers, the respective instance of the linear neural network layer is configured to apply the slice of the corresponding weight matrix for the linear neural network layer to an input slice derived from a given token to generate an output of the respective instance of the linear neural network layer for the given token. Clause 16. The method of clause 15, wherein each nested expert layer block comprises a respective instance of a self-attention layer block, wherein the self-attention block comprises, for each attention head in a set of one or more attention heads: a set of input linear neural network layers that process an input token to generate a query, key, and value from the input token; and a self-attention mechanism that applies self-attention across queries, keys, and values for each of the tokens in the sequence to generate an initial attended token for each of the tokens in the sequence, wherein: (i) each nested expert layer block comprises a respective instance of the set of input linear neural network layers and (ii) the self-attention mechanism is shared across all of the nested expert layer blocks. Clause 17. The method of clause 16, wherein the set of attention heads comprises a plurality of attention heads, the self-attention layer block comprises an output linear neural network layer that, for a given token, processes a combined token generated from the initial attended tokens for the given token to generate an attended token, and each nested expert layer block comprises a respective instance of the output linear neural network layer. Clause 18. The method of clause 16 or 17, wherein each nested expert layer block is configured to, prior to providing a query, key, and value for a given token as input to the shared self-attention mechanism, pad the query, key, and value to each have a fixed dimensionality that is shared across the nested expert layer blocks. Clause 19. The method of any one of clauses 16-18, wherein the self-attention layer block further comprises a multi-layer perceptron (MLP) layer that comprises one or more linear neural network layers and wherein each nested expert layer block comprises a respective instance of the MLP layer. Clause 20. The method of any preceding clause, wherein processing the token using the routing layer to assign the token to one of the nested expert layer blocks comprises: processing the token to generate a respective score for each of the nested expert layer blocks; and assigning the token to one of the nested expert layer blocks based on the respective scores. Clause 21. The method of clause 20, wherein assigning the token to one of the nested expert layer blocks based on the respective scores comprises: assigning the token to one of the nested expert layer blocks based on the respective scores and on respective capacity bounds for each of the nested expert layer blocks. Clause 22. The method of clause 21, wherein assigning the token to one of the nested expert layer blocks comprises: assigning the token to one of nested expert layer blocks based on the respective scores and on respective capacity bounds for each of the nested expert layer blocks through expert preferred routing. Clause 23. The method of clause 21 or clause 22, further comprising: determining the respective capacity bounds for each of the nested expert layer blocks based on a respective model dimension for each of the nested expert layer blocks and an inference compute requirement for the computer vision neural network. Clause 24. The method of clause 23, wherein the inference compute requirement is a requirement on a number of FLOPs that is required to generate a network output for a visual input using the computer vision neural network. Clause 25. A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one more computers to perform the operations of the respective method of any one of clauses 1- 24. Clause 26. One or more computer storage media storing instructions that when executed by one or more computers cause the one more computers to perform the operations of the respective method of any one of clauses 1-24. What is claimed is:

Claims

CLAIMS 1. A method performed by one or more computers, the method comprising: obtaining a visual input; generating, from the visual input, a sequence of tokens; processing the sequence of tokens using a computer vision neural network to generate a network output for the visual input, wherein the computer vision neural network includes a sequence of layers that each update the sequence of tokens, and wherein the sequence of layers includes one or more MoE layers that each include: (i) a routing layer; and (ii) a plurality of nested expert layer blocks each having respective parameters, wherein, for each nested expert layer other than a largest nested expert layer of the plurality of nested expert layer blocks, the parameters of the nested expert layer are a proper subset of the parameters of one or more other nested expert layer blocks in the plurality of nested experts, and wherein processing the sequence of tokens comprises, for each MoE layer: receiving the sequence of tokens; and updating the sequence of tokens, comprising, for each token: processing the token using the routing layer to assign the token to one of the nested expert layer blocks; and processing the token using the assigned nested expert layer block to generate an initial updated token.

2. The method of claim 1, wherein the visual input is an image.

3. The method of claim 1, wherein the visual input is a video.

4. The method of any preceding claim, wherein processing the sequence of tokens using the computer vision neural network further comprises: processing the sequence of tokens after being updated by a last layer in the sequence to generate the network output.

5. The method of claim 4, wherein the network output comprises an embedding of the visual input in an embedding space, and wherein processing the sequence of tokens afterbeing updated by a last layer in the sequence to generate the network output comprises: processing the sequence of tokens after being updated by a last layer in the sequence to generate the embedding.

6. The method of claim 5, wherein processing the sequence of tokens after being updated by a last layer in the sequence to generate the embedding comprises: applying global average pooling to the sequence of tokens after being updated by a last layer in the sequence.

7. The method of claim 4, wherein the network output comprises a classification output for a classification task.

8. The method of claim 8, wherein processing the sequence of tokens after being updated by a last layer in the sequence to generate the network output comprises: processing the sequence of tokens after being updated by a last layer in the sequence to generate an embedding of the visual input; and processing the embedding to generate the classification output.

9. The method of claim 8, wherein processing the embedding to generate the classification output comprises processing the embedding using one or more output layers to generate the classification output.

10. The method of claim 8, wherein processing the embedding to generate the classification output comprises processing the embedding and respective embeddings for each of a plurality of object classes to generate the classification output.

11. The method of any one of claims 7-10, wherein the network output further comprises an object detection output.

12. The method of claim 4, wherein the network output is a segmentation output.

13. The method of any preceding claim, wherein each nested expert layer block has a different subset of a shared set of parameters.

14. The method of any preceding claim, wherein: each nested expert layer block comprises a respective instance of each of one or more linear neural network layers that each have a corresponding weight matrix, and for each nested expert layer block and for each of the one or more linear neural network layers, the respective instance of the linear neural network layer has a different slice of the corresponding weight matrix for the linear neural network layer.

15. The method of claim 14, wherein for each nested expert layer block and for each of the one or more linear neural network layers, the respective instance of the linear neural network layer is configured to apply the slice of the corresponding weight matrix for the linear neural network layer to an input slice derived from a given token to generate an output of the respective instance of the linear neural network layer for the given token.

16. The method of claim 15, wherein each nested expert layer block comprises a respective instance of a self-attention layer block, wherein the self-attention block comprises, for each attention head in a set of one or more attention heads: a set of input linear neural network layers that process an input token to generate a query, key, and value from the input token; and a self-attention mechanism that applies self-attention across queries, keys, and values for each of the tokens in the sequence to generate an initial attended token for each of the tokens in the sequence, wherein: (i) each nested expert layer block comprises a respective instance of the set of input linear neural network layers and (ii) the self-attention mechanism is shared across all of the nested expert layer blocks.

17. The method of claim 16, wherein the set of attention heads comprises a plurality of attention heads, the self-attention layer block comprises an output linear neural network layer that, for a given token, processes a combined token generated from the initial attended tokens for the given token to generate an attended token, and each nested expert layer block comprises a respective instance of the output linear neural network layer.

18. The method of claim 16 or 17, wherein each nested expert layer block is configured to, prior to providing a query, key, and value for a given token as input to the shared self-attention mechanism, pad the query, key, and value to each have a fixed dimensionality that is shared across the nested expert layer blocks.

19. The method of any one of claims 16-18, wherein the self-attention layer block further comprises a multi-layer perceptron (MLP) layer that comprises one or more linear neural network layers and wherein each nested expert layer block comprises a respective instance of the MLP layer.

20. The method of any preceding claim, wherein processing the token using the routing layer to assign the token to one of the nested expert layer blocks comprises: processing the token to generate a respective score for each of the nested expert layer blocks; and assigning the token to one of the nested expert layer blocks based on the respective scores.

21. The method of claim 20, wherein assigning the token to one of the nested expert layer blocks based on the respective scores comprises: assigning the token to one of the nested expert layer blocks based on the respective scores and on respective capacity bounds for each of the nested expert layer blocks.

22. The method of claim 21, wherein assigning the token to one of the nested expert layer blocks comprises: assigning the token to one of nested expert layer blocks based on the respective scores and on respective capacity bounds for each of the nested expert layer blocks through expert preferred routing.

23. The method of claim 21 or claim 22, further comprising: determining the respective capacity bounds for each of the nested expert layer blocks based on a respective model dimension for each of the nested expert layer blocks and an inference compute requirement for the computer vision neural network.

24. The method of claim 23, wherein the inference compute requirement is a requirement on a number of FLOPs that is required to generate a network output for a visualinput using the computer vision neural network.

25. A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one more computers to perform the operations of the respective method of any one of claims 1-24.

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

Images