Machine learning model architecture combining mixture of experts and model ensembling

Abstract

Some aspects of the present disclosure provide techniques and apparatus for machine learning. In one aspect, base model output data is generated, where generating includes processing input data using at least a portion of a base model of the machine learning model architecture, and the base model output data is processed using a routing model of the machine learning model architecture to determine a selected expert model among a plurality of expert models to use for processing the base model output data. Expert model output data is generated, where generating the expert model output data includes processing the base model output data using the selected expert model, and final output data from the machine learning model architecture is generated, where generating the final output data includes processing the base model output data and the expert model output data using an ensemble model of the machine learning model architecture.

Description

[Technical field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS

[0001] This application claims priority to U.S. Provisional Patent Application No. 63 / 317,236, filed March 7, 2022, which claims priority to U.S. Provisional Patent Application No. 18 / 154,429, filed January 13, 2023, the entire contents of which are incorporated herein by reference.

[0002] introduction Aspects of the present disclosure relate to a machine learning model architecture that combines mixtures of experts and model ensembling.

[0003]

[0003] Mixture of experts (MoE) is a form of conditional computational architecture for machine learning models that allows processing relevant parts of the model architecture while bypassing unnecessary parts of the model architecture. In particular, sparse MoE models dynamically select which parameters of the network to use, conditional on the input. Compared to dense models, sparse MoE models can significantly expand the number of parameters and improve performance, while maintaining a similar computational cost per input.

[0004]

[0004] However, while there are several approaches for training very large (e.g., with over 100 billion parameters) MoE models, there is currently no effective method for training small and efficient MoE models. Therefore, MoE models are generally not suitable for use on resource-constrained devices, including mobile devices, edge processing devices, always-on devices, Internet of things (IoT) devices, etc.

[0005]

[0005] Therefore, an efficient and high-performance MoE framework that can run on a wider variety of devices, including resource-constrained devices, is desirable. Summary of the Invention

[0006]

[0006] Some aspects provide a computer-implemented method, the method including: generating base model output data, the method including processing input data using at least a portion of a base model of the machine learning model architecture; processing the base model output data using a routing model of the machine learning model architecture to determine a selected expert model of a plurality of expert models to use for processing the base model output data; generating expert model output data, the method including processing the base model output data using the selected expert model; and generating final output data from the machine learning model architecture, the method including processing the base model output data and the expert model output data using an ensemble model of the machine learning model architecture.

[0007]

[0007] Some aspects provide a computer-implemented method for training a machine learning model architecture, the method including: training a base model including a plurality of layers using a training dataset; performing clustering on features output from an intermediate layer of the plurality of layers to generate a plurality of training data subsets; training each respective expert model of the plurality of expert models on a respective training data subset of the plurality of training data subsets; training a router model to route training data samples between the plurality of expert models; and training an ensemble model to generate machine learning model architecture output data based on base model output data generated by the base model and expert model output data generated by one or more of the plurality of expert models.

[0008]

[0008] Another aspect provides a processing system configured to perform the aforementioned methods as well as methods described herein; a non-transitory computer readable medium comprising instructions which, when executed by one or more processors of the processing system, cause the processing system to perform the aforementioned methods as well as methods described herein; a computer program product embodied on the computer readable storage medium comprising code for performing the aforementioned methods as well as methods further described herein; and the processing system comprising means for performing the aforementioned methods as well as methods further described herein.

[0009] The following description and the annexed drawings set forth in detail certain illustrative features of the one or more aspects. [Brief description of the drawings]

[0010]

[0010] The accompanying drawings depict several aspects of the present disclosure and therefore should not be considered as limiting the scope of the present disclosure. [Figure 1]

[0011] 1 depicts an exemplary mixed-experts model architecture using model ensembling. [Diagram 2]

[0012] 2 illustrates an example of processing input data using the architecture described with respect to FIG. 1; [Diagram 3]

[0013] 2 illustrates another example of processing input data using the architecture described with respect to FIG. 1; [Figure 4]

[0014] 2 illustrates another example of processing input data using the architecture described with respect to FIG. 1; [Diagram 5]

[0015] FIG. 1 is a flow diagram depicting an example method for processing data using the model architecture. [Figure 6]

[0016] FIG. 1 is a flow diagram depicting an exemplary inference method. [Figure 7]

[0017] FIG. 1 is a flow diagram depicting an exemplary training method. [Figure 8]

[0018] 1 depicts an exemplary processing system configured to perform various methods described herein.

[0011]

[0019] For ease of understanding, wherever possible, like reference numerals have been used to designate like elements common to the figures, and it is contemplated that elements and features of one embodiment may be beneficially incorporated in other embodiments without further detail. DETAILED DESCRIPTION OF THE PREFERRED EMBODIMENTS

[0012]

[0020] Aspects of the present disclosure provide apparatus, methods, processing systems, and non-transitory computer-readable media for training and inference using a mixture of experts model architecture that includes model ensembling.

[0013]

[0021] At least some conventional convolutional neural networks systematically perform all interneuron connections for any given input. However, this conventional processing technique is often inefficient because not all features are relevant in all situations. For example, if an image does not contain a human, features specialized for distinguishing baby humans from adult humans may not be relevant to the model's inference (e.g., classification) and therefore may be turned off or not used. Similarly, images have various levels of complexity, ranging from simple cases such as a single object on a white background to images with much more clutter and difficult camera angles. Conditional computation and early termination have been considered to give networks the ability to exploit these natural phenomena.

[0014]

[0022] In conditional computation, subparts of the network are dynamically turned on or off based on the input data. In early termination, a similar conditional behavior is applied, but across the depth / complexity dimension. In the former case, predictions can be terminated early in the network, thus preventing unnecessary computation, while in the latter case, full depth execution is used.

[0015]

[0023] Mixture of Experts (MoE) models can be used to provide conditional computation in Transformer models. As a result, models with a large number of parameters but many sparse routing decisions can be used for natural language processing applications. In computer vision, similar models have emerged, introducing conditional routers in Transformers applied to patches of images. Conditional computation has shown great promise for such large-scale Transformer models, but has not been used in standard convolutional networks or smaller datasets. However, many networks in commercial applications run on resource-constrained devices, such as (but not limited to) mobile phones and low-computation IoT platforms. For this reason, lightweight convolutional networks such as MobileNetV3 are still very popular in practice, as such networks can indeed be run on resource-constrained devices. Some attempts to integrate MoE models into convolutional networks have suffered from various issues, including training instability and lack of clarity on the benefits brought by expert models in terms of increased training costs.

[0016]

[0024] Aspects described herein provide techniques for single-gate MoE models to enable conditional computation in various types of machine learning models, including traditional convolutional neural networks, as well as techniques for both small and large datasets. In particular, aspects described herein improve upon at least some traditional MoE frameworks with an additional regularization technique in which selected experts are ensembled with a shared knowledge branch shared across experts. In some aspects described herein, this shared knowledge branch can be leveraged as an early termination point in the network, as well as by feeding early features of the shared knowledge branch as input to the expert.

[0017]

[0025] Some aspects described herein provide a stable, asynchronous training pipeline for single-gate MoE models that reach similar or better accuracy than other approaches (such as joint expectation-maximization training schemes) while significantly reducing the complexity and cost during training. In particular, the single-gate design of the hierarchical classification used in some aspects described herein allows for more efficient training (e.g., since experts can be trained more independently) and reduces the number of memory accesses to load the computational path for any given sample.

[0018]

[0026] Thus, aspects described herein improve upon the state of the art in many ways. For example, some architectures described herein allow for ensembling predictions of a shared base model with predictions of an expert model, which can result in significant performance improvements over at least some conventional architectures. Furthermore, some architectures described herein can allow for early routing of samples to experts, which obviates the need to fully run the base model and saves computational costs. In addition, some architectures described herein can allow for early termination for easy samples that can be confidently classified without the need for an expert. The routing function for assigning experts to inputs, in some aspects described herein, implements per-sample routing that outperforms class-based routing approaches. In addition, the training schemes described herein for gates and experts can address the instability problem of at least some conventional training schemes for MoE models. Furthermore, aspects described herein generally outperform state-of-the-art methods in high-efficiency computational domains. Beneficially, the efficient MoE architecture described herein can perform many types of machine learning tasks (e.g., image classification and recognition) with reduced computational complexity and cost, enabling the disclosed architecture to be used even on resource-constrained devices (e.g., mobile platforms).

[0019] An introduction to mixed-expert models.

[0027] Conventional MoE architectures can be deployed in the context of image classification. In general, MoE architectures use a set of K experts, (e k ) 1...K The execution of the experts is conditioned by the output of the router r, which outputs a probability distribution over the set of experts. Thus, the total likelihood of the model for the training dataset D to be maximized can be expressed as

[0020]

number

[0021]

[0028] Traditional MoE architectures generally rely on the assumption that each expert specializes in a portion of the data and that a mapping from samples to the most relevant experts can be easily learned. However, at least some traditional MoE architectures suffer from two common problems: (1) experts may become over-specialized and begin to overfit to their respective subsets, which is particularly dangerous when targeting small data sets, and (2) learning the routing without supervision is non-trivial, and the router / learning can easily collapse effectively using only a few experts. Such router collapse is particularly harmful when using small neural network architectures, since router collapse reduces the parameter space considerably.

[0022]

[0029] Aspects described herein modify the standard MoE architecture to mitigate these problems. First, the aspects described herein introduce a shared knowledge branch (which may be referred to in some aspects as a "base model"). In some aspects, the base model is trained on the entire dataset D and used as a form of expert regularization through ensembling, as described in more detail below. While some previous work has ensembled specialized experts together in MoE models to increase accuracy, the aspects described herein ensemble one expert with a non-specialist branch (from the base model), which can result in higher prediction accuracy. Furthermore, some aspects described herein provide a more efficient training scheme that keeps the router and experts independent, which avoids router collapse (e.g., one expert becomes the preferred expert simply by gaining training earlier than other experts).

[0023]

[0030] In various aspects described herein, the base model may be a simple (and possibly lightweight) neural network model trained on the entire dataset D, which is run (at least in part) on all input samples. By ensembling the base model with a selected expert model, accuracy is greatly improved. Furthermore, in some aspects, the base model allows for early termination in inference time when the base model predictions are reliable, thus avoiding redundant expert computations for some (simpler) samples. In some aspects, outputs from early layers of the base model may be reused as inputs to the router and / or expert model, allowing for computational reuse and improved efficiency of the overall architecture.

[0024]

[0031] In some aspects described herein, the router model r is a simple linear layer that takes the pre-logits of the base model as input. In other aspects, the router may correspond to other architectures, such as a multilayer perceptron (MLP), a convolutional neural network, a transformer, etc. During training, the router model outputs a probability distribution r(k|x) over the experts, which allows for direct backpropagation through its weights. During inference, in some aspects, the most likely expert model, e.g., r test (k|x)=[|argmax k’ r(k'|x)|] is selected for execution. In other embodiments, rather than selecting the best (single) expert model, several expert models may be selected.

[0025]

[0032] In various aspects described herein, the expert model is parameterized as a neural network whose inputs are the intermediate feature maps generated by the base model. This can provide two advantages: (1) early features of the expert models are shared and frozen, thereby reducing the number of trainable parameters in each and reducing the risk of overfitting the experts (especially when training on small datasets), and (2) this allows the expert model to reuse computations from the base model during inference, further improving efficiency.

[0026]

[0033] In some aspects described herein, the ensemble model is a (potentially shallow) neural network model. In some aspects, a single ensemble model is used to combine the output of the base model with the output of the selected expert model. In other aspects, there may be one ensemble model for each expert model, where each expert-specific ensemble model combines the output of the base model with the output of the corresponding expert. This combination of ensembling of base and expert models may be referred to as an "ensemble model" in some aspects. For example, e'k(x)=d k (e k (x),φ(x)) is the ensemble model d k , which ensembles the k-th expert and the base model φ.

[0027]

[0034] According to various aspects, the architectures and techniques described herein may generally include one or more of a shared knowledge branch (e.g., a base model), one or more gating models that allow early termination from the base model, a set of expert models specialized for particular portions of the dataset, a router model for expert selection, and / or an ensemble model that combines predictions of a generic base model with predictions of selected expert models.

[0028] Exemplary Mixture of Experts Model Architecture with Model Ensembling

[0035] 1 depicts an example Mixture of Experts (MoE) model with model ensembling architecture 100. In some aspects, architecture 100 corresponds to, is part of, and / or is implemented by a machine learning system.

[0029]

[0036] The architecture 100 includes a base model 110 that includes several layers 112, 114, and 116 (labeled layers 1-n in FIG. 1). Although three separate layers 112, 114, and 116 are depicted for conceptual clarity, in aspects there may be any number of layers in the base model 110. In general, each of the layers of the base model 110 may be one or more layers of a neural network model, such as, but not limited to, a convolutional layer, a normalization layer (e.g., batch normalization), a pooling layer (e.g., average, max, or min pooling), an activation layer, an attention layer, etc. The base model 110 may form part of a so-called general purpose or shared knowledge branch of the architecture 100. The architecture of the base model 110 may vary depending on the particular implementation. In some aspects, the base model 110 is a (deep) neural network. In one aspect, the base model 110 is a transformer neural network.

[0030]

[0037] Generally, in some aspects, the base model 110 may be trained on the entire training dataset. That is, while other components of the architecture 100 (e.g., expert models) may each be trained on a subset of exemplars available in the training dataset, the base model 110 may be trained on all such exemplars. As shown, input data 102 may be provided as input to the base model 110 to generate an output (e.g., output data 104). As described further below, if the base model 110 is confident in its output (i.e., if the output of the base model 110 is associated or scored with sufficient confidence), the architecture 100 may optionally bypass the specialized expert models 140, 142, and / or 144 for accurate prediction, thereby reducing computational costs.

[0031]

[0038] Further, in some aspects, features generated and / or output by one or more layers of the base model 110 (e.g., from the final layer of the base model 110, from the penultimate layer of the base model 110, or from an intermediate / inner layer) may be used as input to a router model 130 (also referred to in some aspects as a routing model) for expert selection. In some aspects, predictions of the base model 110 may be combined with predictions of one or more selected expert models 140, 142, and / or 144 using an ensemble model 150 to improve task performance (e.g., classification, recognition, etc.). In general, the specific architecture of the expert models may vary depending on the particular implementation. For example, in some aspects, the expert models 140, 142, and / or 144 may comprise convolutional neural networks, transformer models, etc.

[0032]

[0039] As shown, the architecture 100 further includes several early termination gating models 120, 122, and 124, which allow the architecture 100 to exploit the fact that not all samples of the input data 102 involve equal effort or computational expense for processing. For example, the early termination gating models 120 and 122 allow early termination from the base model 110 based on factors such as, but not limited to, the difficulty of the samples of the input data 102 and / or the prediction confidence of the base model 110. The early termination gating model 124 may allow bypassing of the ensemble model 150 and / or the router model 130, as well as the expert models 140, 142, and 144. In general, the early termination gating models 120, 122, and 124 can be implemented using a variety of techniques, including simple threshold or learned gating models. Various (non-limiting) early termination scenarios are described with respect to FIGS. 2-4.

[0033]

[0040] In some aspects, an early termination gating model (e.g., early termination gating models 120, 122, and / or 124) may be implemented using a single linear layer that takes as input a flat (e.g., one-dimensional) vector, such as (but not limited to) a pre-logit. In alternative aspects, the early termination gating model may include additional layers, such as one or more convolutional layers, and may take as input a multi-dimensional (e.g., 2D) feature map. In such aspects, the additional layers may reduce the multi-dimensional feature map to a flat vector for input to the linear layer. In some aspects, the early termination gating model includes an MLP.

[0034]

[0041] In some aspects, the early termination gate models 120 and 122 are configured to make a binary decision indicating whether to proceed to the router model 130 (e.g., provide the input of the termination gate model as an input to the router model 130). Specifically, as shown, the early termination gate model 120 can evaluate the intermediate feature map 118A (also referred to in some aspects as "intermediate activation data") generated by the layer 112 of the base model 110 to determine whether to provide the intermediate feature map 118A as an input to a subsequent layer of the base model 110 (e.g., layer 114) or as an input to the router model 130. Similarly, the early termination gate model 122 can evaluate the intermediate feature map 118B generated by the layer 114 of the base model 110 to determine whether to provide the intermediate feature map 118B as an input to a subsequent layer of the base model 110 (e.g., layer 116) or as an input to the router model 130.

[0035]

[0042] In some such cases, a small transformation model or layer (e.g., a small convolutional neural network) may be used to convert the intermediate feature maps output by layers of the base model 110 and provided as input to the early-terminating gating model into the appropriate “shape” or format of input used by the router model 130.

[0036]

[0043] In some aspects, an early termination gate (e.g., early termination gate models 120 and 122) that controls the termination from the intermediate layers of the base model 110 may be referred to as a "base model early termination gate." An early termination gate (e.g., early termination gate model 124) that controls the bypass of the router model 130, the expert models 140-144, and the ensemble model 150 may be referred to as an "expert model bypass gate."

[0037]

[0044] For conceptual clarity, three early exit gate models 120, 122, and 124 are depicted, however, in aspects, there may be any number and variety of early exit gates in architecture 100. In some aspects, each layer of base model 110 has a corresponding early exit gate model. In other aspects, only a subset of layers of base model 110 may have exit gate models.

[0038]

[0045] As shown, the architecture 100 further includes a router model 130. In some aspects, the router model r is parameterized by a neural network and takes as input features either the output of the final layer of the base model 110 (e.g., when coming from the early-termination gate model 124, referred to as a "base feature map") or an intermediate feature map (e.g., 118A, 118B) from one of the early-termination gate models 120, 122. In some aspects, although the illustrated example depicts the termination gate model 124 as receiving input from the last layer of the base model 110, the last layer of the base model 110 may actually correspond to the penultimate layer of the base model. In some such aspects, if the early-termination gate model 124 decides not to use the router model 130, these features may be processed using a final layer or component (e.g., a fully connected layer). Based on these input features, the router model 130 may assign the input data 102 to one or more of the expert models 140, 142, and / or 144 based on the probabilities generated by the router model 130. In some aspects, routing an input sample to a single expert model (e.g., one of the expert models 140, 142, and 144) reduces computational costs and communication overhead. In some aspects, the router model 130 relies on a per-sample routing mechanism that allows the expert models 140, 142, and 144 to specialize on different partitions of the input space, rather than a per-class routing mechanism as used in at least some conventional MoE architectures.

[0039]

[0046] In some aspects, the router model 130 may be implemented using a single linear layer that takes flat (e.g., one-dimensional) vectors, such as prelogits, as input. In alternative aspects, the router model 130 may include additional layers, such as one or more convolutional layers, and may take multi-dimensional (e.g., 2D) feature maps as input. In some such aspects, the additional layers may reduce the multi-dimensional feature maps to flat vectors for input to the linear layer.

[0040]

[0047] In some examples, the router model 130 may implement a thresholding technique to select an expert model (e.g., expert model 140, 142, or 144) based on a simple thresholding rule on the output probability of the router model 130. Intuitively, if two expert models have equally high weights (e.g., equally high probability of being generated by the router model 130), both of these expert models may be relevant to the current sample. However, if all experts have low weights, the sample may be interpreted as difficult or hard to route and thus may be misclassified by any selected expert. In some such cases, early termination may be a better route than going through the architecture 100 to avoid unnecessary computation.

[0041]

[0048] Formally, the overall “anytime” model of architecture 100 can be expressed as the following dynamically weighted ensemble of p at-τ It can be defined as:

[0042]

number

[0043] where [|·|] is the indicator function and τ ∈ [0,1] is the thresholding parameter.

[0044]

[0049] In some aspects, the aforementioned anytime models are only capable of early termination of "difficult" samples (e.g., samples that are difficult to route). However, from a computational standpoint, it may also be beneficial to early terminate "easy" samples (e.g., samples that have already been accurately and reliably classified by the base model 110) since such samples do not need to be refined by one of the expert models 140, 142, and 144.

[0045]

[0050] In some aspects, the router model 130 does not provide a simple signal for separating such "easy" samples, whereas the base model 110 does, and such early termination may be equivalent to simply using the predictions of the base model 110 as outputs, without using any expert models. Thus, in some aspects, a threshold may be used on those outputs of the base model 110 to identify input data 102 (e.g., images) that the base model 110 trusts.

[0046]

[0051] In some aspects, the threshold is defined manually. In some aspects, the threshold may be set based on a holdout validation set, but there are different approaches to thresholding that allow the model to fit within a given computational budget while still producing accurate outputs over a range of input sample difficulties. This mitigates validation set overfitting issues that can arise if the threshold is always set based on validation set performance.

[0047]

[0052] Thus, in the illustrated architecture 100, early-termination gating models 120, 122, and 124 can be used to implement base model thresholding.

[0048]

[0053] As shown, the architecture 100 further includes expert models 140, 142, and 144 (labeled expert models 1 through k in FIG. 1). In some aspects, the expert models 140, 142, and 144 arek :x→y, where

[0049]

number

[0050] The expert models 140, 142, and 144 may be implemented using a neural network model such that x is the predicted output vector. In some aspects, the expert models 140, 142, and 144 receive the original input data (e.g., the input data 102) as input. In other aspects, feature maps (e.g., 118A, 118B) from the last processed layer of the base model 110 (e.g., as determined by the early-termination gating models 120, 122, and / or 124 in this example) are provided as input to one or more selected expert models.

[0051]

[0054] Compared to at least some conventional MoE architectures (wherein the input to the expert model is typically taken from the output of a previous layer in the overall architecture), in architecture 100, expert models 140, 142, and 144 may operate on raw input features or on intermediate feature maps branched from intermediate layers in base model 110. This beneficially simplifies the training scheme and model stability, as described further below with respect to an exemplary training scheme.

[0052]

[0055] The architecture 100 further includes one or more ensemble models 150. In general, the expert models 140, 142, and 144 in the architecture 100 may be trained and specialized for specific portions of the dataset and therefore may perform better on these portions of the input data. However, despite initializing the experts using pre-trained weights taken from the entire dataset, in some aspects the expert models 140, 142, and 144 may perform significantly worse than the base model 110 on portions of the dataset for which each expert model is not specialized. This may be due to the fact that the individual expert models have modified parameters during training that were useful for out-of-domain samples (samples that are no longer routed to a particular expert). This may pose a risk in some aspects, as relying exclusively on the expert model selection of the router model 130 may significantly degrade the performance of the entire architecture 100 in the event of a routing error. In some aspects, to improve or overcome this issue, the architecture 100 may implement model ensembling.

[0053]

[0056] At least some conventional MoE architectures may use ensembling at the expert level, where predictions of one or more expert models (e.g., top-t experts) are combined to improve performance. However, this conventional ensembling of experts significantly increases computational costs while providing little benefit. Thus, as shown, architecture 100 ensembles predictions of selected expert models (e.g., from expert models 140, 142, or 144 in this example), which may be referred to as a "specific knowledge branch," with predictions of base model 110, which may be referred to as a general knowledge branch or "shared knowledge branch" as described above.

[0054]

[0057] The ensemble model 150 ultimately generates the output data 104 in a scenario in which the expert model is not bypassed, as will be further described with respect to the example of Figures 2-4. In some aspects, the output data 104 may correspond to the output of the base model 110 if the early-termination gating model 124 decides to bypass the expert model.

[0055]

[0058] While FIG. 1 generally provides an overview of the structure and data flow of architecture 100, it should be noted that for clarity, additional details (such as data transformation or layering) may be omitted or added based on implementation without departing from the scope of the overall framework depicted in FIG. 1.

[0056]

[0059] Figure 2 depicts an example 200 of processing input data 102 using architecture 100 described with respect to Figure 1. In some aspects, example 200 may correspond to a workflow performed by a machine learning system, such as a machine learning system that corresponds to, includes, and / or implements architecture 100, as described above.

[0057]

[0060] In the example 200, the input data 102 is processed by each layer in the base model 110. After each respective layer of the base model 110 is processed, the output (e.g., intermediate feature maps) is provided to a corresponding early termination gate associated with each layer (e.g., early termination gate model 120 for layer 112, early termination gate model 122 for layer 114, etc.) to determine whether to terminate the base model 110 early. As described, in some aspects, each early termination gate model 120 and 122 may generally output a binary classification (based on the input intermediate features) indicating whether to perform early termination. In this example, none of the early termination gate models (e.g., early termination gate models 120, 122) results in early termination from the base model 110. This may be because, for example, predictions based on the intermediate feature maps (as evaluated by the early termination gate models 120 and 122) are not sufficiently reliable to warrant early termination (e.g., above a threshold). In this example, when a base model early termination gate (e.g., early termination gate model 120 or 122) determines not to terminate, the input feature map (provided to the early termination gate model) is provided to the next or subsequent layer of the base model 110. In some aspects, the early termination gate model itself provides features to the next layer of the base model 110. In other aspects, the base model early termination gate may provide a signal (e.g., a binary signal) that then causes the intermediate feature map from the previous layer to be provided to the next layer in the base model 110.

[0058]

[0061] As depicted, the base model 110 is fully processed in this example, and the output from the base model 110 is processed by the early-termination gating model 124. As used herein, the base model 110 may be referred to as "fully processed" to indicate that the data has passed through at least the penultimate layer of the base model 110, in which case the final processing (e.g., the final fully connected layer or classification layer) may be implemented as a separate component. In this example, the output from the base model 110 provided to the early-termination gating model 124 may be referred to as a "base feature map" 202, which may be a feature map from the penultimate or final layer of the base model 110, such as before the final fully connected or classification layer of the base model 110.

[0059]

[0062] As described above, the early termination gating model 124 may generally be used to determine whether to pass the base feature map 202 to the router model 130 or to an output (e.g., to a final layer of the base model 110). In the illustrated example, the early termination gating model 124 determines to provide the output from the base model 110 to the router model 130 and the ensemble model 150. In some aspects, the output provided by the early termination gating model 124 is the base feature map 202, as described above. In other aspects, the output of the early termination gating model 124 may be a control signal (e.g., a binary signal) that causes the base feature map 202 to be selectively provided to the ensemble model 150 and the router model 130.

[0060]

[0063] In this example, the router model 130 receives the output from the base model 110 (e.g., the final feature map generated by the base model 110) as an input from the early-termination gate model 124. The router model 130 processes the output feature map to select one of the expert models 140, 142, and 144. In the illustrated example, the expert model 142 is selected to process the base feature map 202. The output of the expert model 142 is then provided to the ensemble model 150. In this example, the output of the expert model 142 may be another feature map, shown in FIG. 2 as the expert feature map 204 (e.g., rather than a classification or other discrete output).

[0061]

[0064] Further, in the depicted example, the ensemble model 150 processes the output from the base model 110 (e.g., the base feature map 202 in this example) and the output of the selected expert model 142 (e.g., the expert feature map 204 in this example) to generate the output data 104. The specific content and nature of the output data 104 may vary depending on the particular implementation and task. For example, the output data 104 may correspond to discrete task outputs, such as, but not limited to, classification, object detection, image segmentation, ranges of probabilities associated with ranges of classes for classification, etc. In some aspects, the input data 102 includes image data. In some aspects, as described above, the output data 104 may be processed using a final layer or component (e.g., a fully connected layer or classifier) ​​to generate a desired output prediction. In some aspects, the output data 104 includes image classification, object detection output, image segmentation output, etc.

[0062]

[0065] Figure 3 depicts another example 300 of processing input data 102 using architecture 100 described with respect to Figure 1. In some aspects, example 300 may correspond to a workflow performed by a machine learning system, such as a machine learning system that corresponds to, includes, and / or implements architecture 100, as described above.

[0063]

[0066] In the example 300, the input data 102 is processed by only a subset of layers in the base model 110. Unlike the example of Fig. 2, in the depicted example of Fig. 3, one of the early termination gating models (here, the early termination gating model 122) decides to terminate processing by the base model 110. This may be because, for example, the early termination gating model 122 determines that the intermediate feature map 118B has a high probability of generating a reliable inference by one of the expert models 140, 142, and 144 (learned through training).

[0064]

[0067] Thus, the intermediate feature map 118B may be used and / or converted into an intermediate base feature map 302, which is provided as an input to the router model 130 and the ensemble model 150. That is, the intermediate feature map 118B (generated by the layer 114 of the base model 110) may be referred to as the intermediate base feature map 302 to indicate that it is an output of the base model 110 in the depicted workflow (e.g., because the early termination gating model 122 has determined to terminate processing by the base model 110).

[0065]

[0068] In this example, the router model 130 receives the output from the inner layer of the base model 110 (e.g., intermediate feature map 118A / intermediate base feature map 302 generated by the base model 110) as an input from the early-termination gate model 122. The router model 130 processes the intermediate feature map to select one of the expert models 140, 142, and 144. In the illustrated example, the expert model 140 is selected to process the intermediate base feature map 302, and the output of the expert model 140 (expert feature map 304) is then provided to the ensemble model 150. In this example, the output of the expert model 140 is another feature map (e.g., rather than a classification or other discrete output), shown in FIG. 3 as expert feature map 304.

[0066]

[0069] The ensemble model 150 then processes the output from the base model 110 (in this example, the intermediate base feature map 302) and the output of the expert model 140 (in this example, the expert feature map 304) to generate output data 104. The output data 104 can be, for example, a discrete output such as a classification, or a range of probabilities associated with a range of classes for classification, etc.

[0067]

[0070] Figure 4 depicts another example 400 of processing input data 102 using architecture 100 described with respect to Figure 1. In some aspects, example 400 may correspond to a workflow performed by a machine learning system, such as a machine learning system that corresponds to, includes, and / or implements architecture 100, as described above.

[0068]

[0071] In the example 400, the input data 102 is again processed by each layer in the base model 110. Similar to the example of FIG. 2, in this example, none of the early-terminating gating models (e.g., early-terminating gating models 120, 122) terminate early from the base model 110, and the base feature map 402 is then provided to the early-terminating gating model 124.

[0069]

[0072] Unlike the example of Fig. 2, in the illustrated example of Fig. 4, the early-termination gating model 124 decides to terminate without utilizing the router model 130, the expert models 140, 142, and 144, or the ensemble model 150. This may be because, for example, the early-termination gating model 124 determines that the predictions made by the base model 110 are sufficiently reliable that no further processing is necessary. This early termination thus allows to avoid significant additional processing by the router model 130, one or more selected expert models, and the ensemble model 150 without further processing.

[0070]

[0073] Thus, the early termination gating model 124 provides the base model prediction 404 as the output data 104. In some aspects, as described above, the early termination gating model 124 may receive the base feature map 402 and / or the base model prediction 404 from the base model 110 and output the base feature map 402 to the router model 130 and the ensemble model 150 when the early termination gating model 124 decides not to terminate early (as in the examples of FIGS. 2 and 3), or provide the base model prediction 404 as an output when the early termination gating model 124 decides to terminate early (as in the example of FIG. 4).

[0071]

[0074] In other aspects, a final task output layer (e.g., a classification layer) not depicted in Figure 4 may process the base feature map 402 to generate the output data 104. In other aspects, the early termination gating model 124 can provide a signal, such as (but not limited to) a binary signal, that causes the base feature map 402 to be provided to the router model 130 and the ensemble model 150, as in the examples of Figures 2 and 3, or provides the base model prediction 404 as the output data 104 when the early termination gating model 124 decides to terminate early, as in the example of Figure 4.

[0072] Training Procedure

[0075] In some aspects, an asynchronous and easily parallelizable training scheme may be used in which the router model (e.g., router model 130 of FIGS. 1-4) and the expert model (e.g., expert models 140, 142, and 144 of FIGS. 1-4) are trained independently according to the objective of Equation 1. To initialize the router model, embeddings of pre-trained base models may be clustered using techniques such as (but not limited to) K-means.

[0073]

[0076] The training aspects described herein address two training pitfalls that are more specific to neural networks. The first problem arises from a calibration problem in convolutional neural networks. For example, jointly training an ensemble model (e.g., ensemble model 150 in Figs. 1-4) with an expert model (e.g., expert models 140, 142, and 144) may lead to the ensemble model heavily favoring the base model and preventing the expert model from specializing. This may be due to the base model being overly reliable for many training samples. This may be especially evident in small data sets where the base model is already perfectly adapted to the training set. To avoid this problem, in some aspects, the ensemble model may be trained after fully training the expert model.

[0074]

[0077] The second problem arises from the trade-off between the specialization of the expert model and the potential routing error by the router model (e.g., router model 130), which is related to the so-called catastrophic forgetting behavior. In other words, because the expert model is trained on a subset of the data (e.g., the clusters discussed above), the expert model may "forget" classes of data that the expert model never sees during training. While ensemble models mitigate this problem to some extent in small datasets, it may be beneficial to route additional negative samples, albeit with lower weights, to the expert model during training. More specifically, during training of the expert model, the router model r in Equation 1 is subjected to the transformation Γ:

[0075]

number

[0076] The router model outputs are "smoothed" using, where Γ is a smoothing function, γ is a hyperparameter, and x corresponds to the output of the router model and is fed into the smoothing function. For example, in one example, γ=0.05 may be used. In some aspects, in addition to or instead of smoothing the router model outputs during training, the router model outputs are used as sampling probabilities when forming training batches for each expert model.

[0077]

[0078] Thus, the training pipeline can be constructed as follows:

[0078]

[0079] First, train a base model φ (e.g., base model 110) on the dataset D, or use an off-the-shelf pre-trained base model.

[0079]

[0080] Second, define an initial router model r0 (e.g., router model 130) by clustering the input training data. As mentioned above, K-means may be used with hard cluster assignments for the pre-logits of the base model.

[0080]

[0081] Third, the router model r is matched to r0 by minimizing the Kullback-Leibler (KL) divergence. That is, the router model can be trained to minimize the KL divergence between the output of the router model and the set of training subsets generated by clustering the input data. In some aspects, the router model can be further trained based on the expert model with the lowest task loss (e.g., the most accurate expert model for each data sample). For example, the system can identify the most accurate expert model (with the lowest task-specific loss) and calculate the cross-entropy loss between the router output (e.g., the selected expert) and the determined best expert. This can be further used to refine the router.

[0081]

[0082] Fourth, we use the initial gate Γ○r0(k|·) smoothed according to Eq. 1 to estimate the k (e.g., expert models 140, 142, and 143). Then, train ensemble model d (e.g., ensemble model 150) using the same objective. In an alternative configuration where multiple ensemble models are used, such as one ensemble model per expert model, this step trains each ensemble model d using the same objective. k will be trained.

[0082]

[0083] Note that the third and fourth steps are independent of each other: thus, the router model and the expert model can be trained in parallel, which provides increased efficiency compared to at least some conventional end-to-end methods.

[0083] A joint training scheme for routers and experts.

[0084] The architecture 100 of FIGS. 1 to 4 includes a router model r (e.g., router model 130), an expert model {e k} 1...K (e.g., expert models 140, 142, and 144), early termination gates (e.g., early termination gate models 120, 122, and 124), and ensemble model d (e.g., ensemble model 150) (or, in other aspects, the respective ensemble modules {d k} 1...K ) An exemplary objective is to maximize the total likelihood from Equation 1.

[0084]

[0085] In some aspects, the training scheme can be extended to jointly train the router model and the expert model. The difficulty in jointly training these models lies in the "chicken and egg" dependency between the router model and the expert model. One technique to address this problem is to use the Expectation Maximization (EM) algorithm to treat the router model weights as latent variables. The derivation of the EM algorithm in the current setting results in the following Equations 4 and 5. Specifically, Equation 4 depicts the "E-step" (e.g., expectation step), and Equation 5 depicts the "M-step" (e.g., maximization step). In Equation 4, for every (x,y), the posterior probability is updated as follows, where x is the input data, y is the associated ground truth for the input data, q(k|x) is the posterior value to be updated, r(k│x) is the output of the router model (e.g., the probability that input x is routed to expert model k), and r(k'│x) corresponds to the probability that the input is routed to expert model k';

[0085]

number

[0086] is the output (using ensembling) of expert model k for input x,

[0087]

number

[0088] is the output (using ensembling) of expert model k' for this input.

[0089]

number

[0090] In addition, in the M step, the system calculates r and all

[0091]

number

[0092] can be trained with the objective (to minimize) where x is the input data with associated ground truth y, r(k|x) is the output of the router model (e.g., the probability that sample x is routed to expert k), q(k|x) is the posterior probability that the router model is trained to match (e.g., given by equation 4 above), e'k(y|x) is the output of expert model k, and KL refers to Kullback-Leibler divergence,

[0093]

number

[0094] is the expectation function).

[0095]

number

[0096]

[0087] In other words, the EM algorithm alternates between an E step of computing updated predictive router model weights based on the performance of the current expert model and an M step of separately training the expert model according to this new assignment indicated by the right-hand term (e.g., the posterior probability given by Equation 4), while also conforming the router model to the new assignment indicated by the left-hand term.

[0097] Exemplary Method for Processing Data Using a Model Architecture

[0088] Figure 5 depicts an example method 500 of processing data using a model architecture as described above with respect to Figures 1-4. In some aspects, method 500 is performed by a system, such as a machine learning system, that corresponds to, includes, and / or implements architecture 100 and / or performs the workflow depicted in examples 200, 300, and / or 400.

[0098]

[0089] At block 505, input data (e.g., input data 102 of FIG. 1) is accessed for processing. As discussed above, the particular form and content of the input data 102 may vary depending on the particular implementation. In at least one aspect, the input data 102 includes image data.

[0099] At block 510, the system generates an intermediate feature map (e.g., intermediate feature map 118 of FIG. 1) based on the input data. For example, as described above, the system may process the input data using a first layer of a base model (e.g., base model 110 of FIG. 1) to generate the intermediate feature map.

[0100]

[0091] In block 515, the system determines whether to terminate the base model early based on the intermediate feature maps (generated in block 510). For example, as described above, the system can process the intermediate feature maps using an early termination gating model (e.g., early termination gating model 120 of FIG. 1) to determine whether to terminate the base model early (e.g., provide the intermediate feature maps to a router model) or continue processing using the base model.

[0101]

[0092] If at block 515 the system determines to terminate the base model early, method 500 continues to block 530, which is described in more detail below. If the system determines not to terminate the base model early, method 500 continues to block 520, where the system determines whether at least one additional layer remains in the base model. If so, method 500 returns to block 510, where the system generates another intermediate feature map (e.g., intermediate feature map 118B) by processing the previously generated intermediate feature map using the next or subsequent layer of the base model.

[0102]

[0093] If, at block 520, the system determines that there are no remaining layers in the base model, method 500 proceeds to block 525, where the system determines whether to bypass the expert model based on the output of the final layer of the base model (e.g., base feature map 202 of FIG. 2). For example, as described above, the system can process the base feature map using an early-termination gate or bypass model (e.g., early-termination gate model 124) to determine whether to bypass the expert model.

[0103]

[0094] If at block 525, the system determines to bypass the expert model, the method 500 continues to block 545, which is described in more detail below. If the system determines not to bypass the expert model, the method 500 continues to block 530.

[0104] At block 530, the system routes the feature map (generated using one or more layers of the base model) to one or more expert models. For example, as described above, the system may process the feature map (generated by the base model) using a routing model (e.g., router model 130) to select one or more expert models (e.g., expert models 140, 142, or 144).

[0105] In one aspect, the particular feature maps routed to the expert model may vary depending on the particular execution path followed for the input data. For example, if the system decides (at block 515) to terminate the base model early, the intermediate feature maps (produced by the intermediate layers of the base model) may be routed to the expert model. If the system processes the base model completely, the base model output features may be routed to the expert model.

[0106]

[0097] In block 535, the system generates one or more expert feature maps (e.g., expert feature maps 204, 304) by processing the routed feature maps (e.g., base feature maps or intermediate feature maps) using the expert model selected in block 530.

[0107]

[0098] In block 540, the system can then ensemble the generated expert feature map with the base model feature map (generated by the final layer of the base model). For example, as described above, an ensemble model (such as ensemble model 150 of FIG. 1) can be used to combine or aggregate the expert feature map and the base feature map.

[0108]

[0099] In block 545, the system generates a final output from the model architecture. In general, the specific operations used to generate the final output may vary depending on the specific implementation. For example, in some aspects, the features themselves (e.g., the ensemble features generated in block 540, or the base model output features generated by the base model) are provided directly as the architecture output.

[0109]

[0100] In some aspects, the final features (e.g., base model features or ensembled features) are processed using a final layer or component (e.g., a fully connected layer or a classifier layer) to generate a model architecture output (e.g., to classify the input data, to provide image segmentation for the input data, etc.).

[0110]

[0101] In general, the method 500 can be used to perform inference using the model architecture, as well as training. For example, during inference, the final output (produced in block 545) can be returned as an output of the architecture. In some aspects, during training, the output can be compared to the labels or ground truth to generate one or more loss terms that can be used to refine various model components, as described above.

[0111]

[0102] In at least one aspect, rather than training the entire architecture end-to-end, the system can train parts of the architecture separately. For example, the system can train the base model using all of the data samples. In addition, a router model can be trained or generated (e.g., using clustering) based on the input data and / or based on the features generated by the base model for each input sample. In some aspects, the system can then use the first trained router to route each training sample to one or more of the expert models (e.g., to generate a subset of the training data). Each subset of the training data can then be used to train a corresponding expert model. Finally, an ensemble model can be trained based on the output of the expert model and the base model.

[0112] Methods of illustrative reasoning

[0103] FIG. 6 illustrates an example inference method 600 that may be implemented using a model architecture such as that described above with respect to FIGS. 1-5.

[0113]

[0104] In block 602, base model output data is generated, where generating includes processing input data using at least a portion of a base model of the machine learning model architecture.

[0114]

[0105] In some aspects, processing the input data using at least a portion of the base model includes processing intermediate feature data from at least one layer of the base model using a gate model associated with the at least one layer to determine whether to terminate the base model early.

[0115]

[0106] In some aspects, the method 600 further includes determining not to terminate the base model early based on output from the gate model, and processing the input data using at least a portion of the base model includes processing all layers of the base model to generate base model output data.

[0116]

[0107] In some aspects, the base model comprises a deep convolutional neural network or a transformer neural network.

[0117]

[0108] In block 604, the base model output data is processed using a routing model of the machine learning model architecture to determine a selected expert model from among a plurality of expert models to use for processing the base model output data.

[0118]

[0109] In some aspects, the method 600 further includes determining to terminate the base model early based on an output from the gate model, and the base model output data processed by the routing model includes an output from the gate model.

[0119]

[0110] In some aspects, the gate model includes a multi-layer perceptron model.

[0120]

[0111] In some aspects, the routing model includes a linear layer, a convolutional neural network model, or a multi-layer perceptron model.

[0121]

[0112] At block 606, expert model output data is generated, where generating the expert model output data includes processing the base model output data with the selected expert model.

[0122]

[0113] In some embodiments, the selected expert model includes a convolutional neural network model or a transformer model.

[0123]

[0114] In block 608, final output data from the machine learning model architecture is generated, and generating the final output data includes processing the base model output data and the expert model output data using an ensemble model of the machine learning model architecture.

[0124]

[0115] In some aspects, the method 600 further includes generating second base model output data, the generating including processing the second input data using at least a portion of the base model; processing the second base model output data using a second gate model associated with a final layer of the base model to determine whether to process the second base model output data using a routing model; determining not to process the second base model output data using the routing model; and generating second final output data from the machine learning model architecture based on the second base model output data.

[0125]

[0116] In some aspects, generating second final output data from the machine learning model architecture includes processing the second base model output data using a fully connected layer.

[0126]

[0117] It should be noted that FIG. 6 is merely one example of a method, and that other methods including fewer, additional, or alternative steps are possible consistent with this disclosure.

[0127] Exemplary Training Methods

[0118] FIG. 7 illustrates an example training method 700 that may be used to train a machine learning model architecture such as those described above with respect to FIGS. 1-5.

[0128]

[0119] In block 702, a base model including multiple layers is trained using a training dataset.

[0129]

[0120] At block 704, clustering is performed on the features output from the intermediate layers of the multiple layers to generate multiple training data subsets.

[0130]

[0121] In some aspects, performing clustering on features output from the intermediate layer of the base model to generate multiple training data subsets includes performing K-means clustering on features output from the intermediate layer of the base model.

[0131]

[0122] In some aspects, K-means clustering generates a number of clusters, K, equal to the number of expert models in the plurality of expert models.

[0132]

[0123] At block 706, each respective expert model of the plurality of expert models is trained on a respective training data subset of the plurality of training data subsets.

[0133]

[0124] At block 708, a router model is trained to route training data samples among multiple expert models.

[0134]

[0125] In some aspects, training a router model to route training data samples among a plurality of expert models includes minimizing a first loss component based on a plurality of training data subsets generated by clustering and output from the router model.

[0135]

[0126] In some aspects, training the router model to route training data samples between the multiple expert models further includes minimizing a second loss component based on an output from the router model and an output from an expert model among the multiple expert models having the lowest task loss.

[0136]

[0127] In some aspects, the first loss component includes a Kullback-Leibler information loss.

[0137]

[0128] In some aspects, the second loss component includes a cross-entropy loss.

[0138]

[0129] In block 710, the ensemble model is trained to generate machine learning model architecture output data based on the base model output data generated by the base model and the expert model output data generated by one or more of the multiple expert models.

[0139]

[0130] In some aspects, the method 700 further includes training one or more gate models to minimize task loss based on the machine learning model architecture output data.

[0140]

[0131] It should be noted that FIG. 7 is merely one example of a method, and that other methods including fewer, additional, or alternative steps are possible consistent with this disclosure.

[0141] Exemplary Processing Device FIG. 8 illustrates an exemplary processing system configured to perform various methods described herein, for example, with respect to FIGS.

[0142]

[0133] The processing system 800 includes a central processing unit (CPU) 802, which may be a multi-core CPU in some examples. Instructions executed in the CPU 802 may be loaded from a program memory associated with the CPU 802 or from a partition of the memory 824, for example.

[0143]

[0134] The processing system 800 also includes additional processing components tailored to specific functions, such as a graphics processing unit (GPU) 804, a digital signal processor (DSP) 806, a neural processing unit (NPU) 808, a multimedia processing unit 810, and wireless connectivity components 812.

[0144]

[0135] NPUs, such as NPU 808, are generally special-purpose circuits configured to implement control and arithmetic logic for running machine learning algorithms, such as algorithms for processing artificial neural networks (ANNs), deep neural networks (DNNs), random forests (RFs), etc. NPUs may alternatively be referred to as neural signal processors (NSPs), tensor processing units (TPUs), neural network processors (NNPs), intelligence processing units (IPUs), vision processing units (VPUs), or graph processing units.

[0145]

[0136] NPUs, such as NPU 808, are configured to accelerate the performance of common machine learning tasks, such as image classification, machine translation, object detection, and various other predictive models. In some examples, multiple NPUs can be instantiated on a single chip, such as a system on a chip (SoC), while in other examples, an NPU can be part of a dedicated neural network accelerator.

[0146]

[0137] NPUs can be optimized for training or inference, or in some cases can be configured to balance performance between both. For NPUs capable of performing both training and inference, the two tasks can still generally be performed independently.

[0147]

[0138] NPUs designed to accelerate training are generally configured to accelerate the optimization of new models, which is a highly computationally intensive operation that involves inputting an existing dataset (often labeled or tagged), iterating over the dataset, and then adjusting model parameters, such as weights and biases, to improve model performance. In general, optimization based on misprediction involves backpropagating through layers of the model to determine gradients to reduce prediction error. In some cases, NPUs may be configured to implement the federated learning methods described herein.

[0148]

[0139] NPUs designed to accelerate inference are generally configured to operate on complete models. Thus, such NPUs can be configured to input new data and rapidly process the data through already trained models to generate model outputs (e.g., inferences).

[0149]

[0140] In one implementation, the NPU 808 is part of one or more of the CPU 802, the GPU 804, and / or the DSP 806.

[0150] In some examples, the wireless connectivity component 812 may include sub-components for, for example, third generation (3G) connectivity, fourth generation (4G) connectivity (e.g., 4G LTE), fifth generation connectivity (e.g., 5G or NR), Wi-Fi connectivity, Bluetooth connectivity, and other wireless data transmission standards. The wireless connectivity component 812 is further coupled to one or more antennas 814. In some examples, the wireless connectivity component 812 enables performing federated learning in accordance with the methods described herein over a variety of wireless data connections, including cellular connections.

[0151]

[0142] The processing system 800 may also include one or more sensor processing units 816 associated with any type of sensor, one or more image signal processors (ISPs) 818 associated with any type of image sensor, and / or a navigation processor 820, which may include satellite-based positioning system components (e.g., GPS or GLONASS), as well as inertial positioning system components.

[0152]

[0143] Processing system 800 may also include one or more input and / or output devices 822, such as a screen, a touch-sensitive surface (including a touch-sensitive display), physical buttons, a speaker, a microphone, etc.

[0153]

[0144] In some examples, one or more of the processors of processing system 800 may be based on an ARM or RISC-V instruction set.

[0154] Processing system 800 also includes memory 824, which represents one or more static and / or dynamic memories, such as dynamic random access memories, flash-based static memories, etc. In this example, memory 824 includes computer-executable components that may be executed by one or more of the aforementioned processors of processing system 800.

[0155]

[0146] In particular, in this example, memory 824 includes a model training component 824A, a model inference component 824B, a base model 824C, an early-termination gating model 824D, an expert model 824E, an ensemble model 824F, a clustering component 824G, and training data 824H. The depicted components, and other components not depicted, may be configured to perform various aspects of the methods described herein.

[0156]

[0147] The processing system 800 further comprises a model training circuit 826, a base model circuit 828, an expert model circuit 830, a clustering circuit 832, a model inference circuit 834, an early termination gating circuit 836, and an ensemble circuit 838. The depicted circuits, and other circuits not depicted, may be configured to implement various aspects of the techniques described herein.

[0157]

[0148] Although depicted as separate components and circuits in FIG. 8 for clarity, the model training circuit 826, the base model circuit 828, the expert model circuit 830, the clustering circuit 832, the model inference circuit 834, the early termination gate circuit 836, and the ensemble circuit 838 may be collectively or individually implemented in other processing devices of the processing system 800, such as within the CPU 802, GPU 804, DSP 806, NPU 808, etc.

[0158] For example, the model training component 824A and / or the model training circuitry 826 may be used to train one or more components of the model architecture (e.g., the base model 824C, the expert model 824E, the clustering component 824G, the early termination gating model 824D, and / or the ensemble model 824F) using the training data 824H, as described above. The model inference component 824B and / or the model inference circuitry 834 may be used to generate output predictions or features using the model architecture (e.g., the base model 824C, the expert model 824E, the clustering component 824G, the early termination gating model 824D, and / or the ensemble model 824F), as described above.

[0159]

[0150] The base model 824C and / or the base model circuitry 828 may correspond to the base model 110 of Figure 1 and may be used to perform initial (and possibly final) processing of input data, as described above. The early termination gate model 824D and / or the early termination gate circuitry 836 may correspond to the early termination gate models 120, 122, and / or 124 of Figure 1 and may be used to evaluate base model features (e.g., intermediate features or final base model feature outputs) to determine whether to terminate the base model early and / or whether to bypass the expert model, as described above.

[0160]

[0151] The expert model 824E and / or expert model circuitry 830 may correspond to the expert models 140, 142, and / or 144 of Fig. 1 and may be used to provide specialized or additional processing of the input data and / or intermediate or base model features, as described above. The ensemble model 824F and / or ensemble circuitry 838 may correspond to the ensemble model 150 of Fig. 1 and may be used to ensemble, combine, or otherwise aggregate the output of one or more expert models 824E and the output of the base model 824C, as described above.

[0161]

[0152] The clustering component 824G and / or the clustering circuitry 832 may correspond to the router model 130 of FIG. 1 and may be used to cluster the training data 824H and / or the output of the base model 824C to generate subsets used to train the expert models 824E, as described above, and / or may be used to route data samples to one or more expert models 824E.

[0162]

[0153] Generally, the processing system 800 and / or its components may be configured to perform the methods described herein.

[0163]

[0154] In particular, aspects of the processing system 800 may be omitted or added in other cases. For example, the multimedia processing unit 810, the wireless connectivity component 812, the sensor processing unit 816, the ISP 818, and / or the navigation processor 820 may be omitted in other cases. Furthermore, aspects of the processing system 800 may be distributed among multiple devices. For example, in some cases, training of a model architecture such as the architecture 100 of FIG. 1 may be performed on one device and inference may be performed on a separate device, such as a resource-constrained device.

[0164] Example clauses

[0155] Example implementations are described in the following numbered clauses.

[0165]

[0156] Clause 1: A computer-implemented method comprising: generating base model output data, the generating step including processing input data using at least a portion of a base model of the machine learning model architecture; processing the base model output data using a routing model of the machine learning model architecture to determine a selected expert model among a plurality of expert models to use for processing the base model output data; generating expert model output data, the generating step including processing the base model output data using the selected expert model; and generating final output data from the machine learning model architecture, the generating step including processing the base model output data and the expert model output data using an ensemble model of the machine learning model architecture.

[0166]

[0157] Clause 2: The method of claim 1, wherein processing the input data using at least a portion of the base model includes processing intermediate feature data from at least one layer of the base model using a gate model associated with the at least one layer to determine whether to terminate the base model early.

[0167]

[0158] Clause 3: The method described in clause 2, further comprising determining to terminate the base model early based on an output from the gate model, wherein the base model output data processed by the routing model includes an output from the gate model.

[0168]

[0159] Clause 4: The method of clause 2, further comprising determining not to terminate the base model early based on output from the gate model, and wherein processing the input data using at least a portion of the base model includes processing all layers of the base model to generate base model output data.

[0169]

[0160] Clause 5: The method of any one of clauses 1 to 4, wherein the base model comprises a deep convolutional neural network or a transformer neural network.

[0170]

[0161] Clause 6: The method of any one of clauses 1 to 5, wherein the router model includes a linear layer, a convolutional neural network model, or a multi-layer perceptron model.

[0171]

[0162] Clause 7: The method of any one of clauses 1 to 6, wherein the selected expert model comprises a convolutional neural network model or a transformer model.

[0172]

[0163] Clause 8: The method of any one of clauses 2 to 7, wherein the gate model includes a multi-layer perceptron model.

[0173]

[0164] Clause 9: The method of any one of clauses 1 to 8, wherein the input data comprises image data and the final output data comprises image classification.

[0174]

[0165] Clause 10: The method of any one of clauses 1 to 9, wherein the input data comprises image data and the final output data comprises an object detection output.

[0175]

[0166] Clause 11: The method of claim 10, wherein the object detection output includes a bounding box.

[0176]

[0167] Clause 12: The method of claim 1, wherein the input data comprises image data and the final output data comprises image segmentation output.

[0177]

[0168] Clause 13: A computer-implemented method for training a machine learning model architecture, the method comprising: training a base model using a training dataset; performing clustering on features output from an intermediate layer of the base model to generate a plurality of training data subsets; training each expert model of a plurality of expert models on one training data subset of the plurality of training data subsets; training a router model to route each training data sample to one expert model of the plurality of expert models; and training an ensemble model to generate machine learning model architecture output data based on the base model output data and the expert model output data.

[0178]

[0169] Clause 14: The method of claim 13, wherein training a router model to route each training data sample to one expert model among a plurality of expert models includes minimizing a first loss component based on a plurality of training data subsets generated by clustering and output from the router model.

[0179]

[0170] Clause 15: The method described in clause 14, wherein training a router model to route each training data sample to one expert model among a plurality of expert models includes minimizing a second loss component based on an output from the router model and an output from an expert model among the plurality of expert models having the lowest task loss.

[0180]

[0171] Clause 16: The method of clause 14 or 15, wherein the first loss component comprises a Kullback-Leibler information loss.

[0181]

[0172] Clause 17: The method of any one of clauses 14 to 16, wherein the second loss component comprises a cross-entropy loss.

[0182]

[0173] Clause 18: A method described in any one of clauses 13 to 17, wherein performing clustering on features output from an intermediate layer of the base model to generate multiple training data subsets includes performing K-means clustering on features output from the intermediate layer of the base model.

[0183]

[0174] Clause 19: The method of clause 18, wherein the K-means clustering generates a number K of clusters equal to a number of expert models in the plurality of expert models.

[0184]

[0175] Clause 20: A method according to any one of clauses 13 to 19, further comprising training one or more gate models to minimize task loss based on the machine learning model architecture output data.

[0185]

[0176] Clause 21: A method according to any one of clauses 13 to 20, wherein the machine learning model architecture comprises: a base model including a plurality of layers and configured for generating based model output data based on at least a portion of the base model; one or more gate models configured to process output data from layers among the plurality of layers of the base model into the base model and determine whether to terminate the base model early; a router model configured to receive the base model output data and select an expert model among the plurality of expert models for processing the base model output data; and an ensemble model configured to process the base model output data and the expert model output data from the selected expert model to generate the machine learning model architecture output data.

[0186]

[0177] Clause 22: A method comprising: generating base model output data by processing input data using at least a portion of a base model of the machine learning model architecture; processing the base model output data using a routing model of the machine learning model architecture to determine a selected expert model among a plurality of expert models to use for processing the base model output data; generating expert model output data by processing the base model output data using the selected expert model; and generating final output data from the machine learning model architecture by processing the base model output data and the expert model output data using an ensemble model of the machine learning model architecture.

[0187]

[0178] Clause 23: The method of clause 22, wherein processing the input data using at least a portion of the base model includes processing intermediate feature data from at least one layer of the base model using a gate model associated with the at least one layer to determine whether to terminate the base model early.

[0188]

[0179] Clause 24: The method described in clause 22 or 23, further comprising determining to terminate the base model early based on an output from the gate model, wherein the base model output data processed by the routing model includes an output from the gate model.

[0189]

[0180] Clause 25: A method according to any one of clauses 22 to 24, further comprising determining not to terminate the base model early based on output from the gate model, and wherein processing the input data using at least a portion of the base model includes processing all layers of the base model to generate base model output data.

[0190]

[0181] Clause 26: The method of any one of clauses 22 to 25, wherein the gate model includes a multi-layer perceptron model.

[0191]

[0182] Clause 27: A method according to any one of clauses 22 to 26, wherein the base model comprises a deep convolutional neural network or a transformer neural network.

[0192]

[0183] Clause 28: A method according to any one of clauses 22 to 27, wherein the routing model comprises a linear layer, a convolutional neural network model, or a multi-layer perceptron model.

[0193]

[0184] Clause 29: The method of any one of clauses 22 to 28, wherein the selected expert model comprises a convolutional neural network model or a transformer model.

[0194]

[0185] Clause 30: A method according to any one of clauses 22 to 29, comprising: generating second base model output data, the generating comprising processing second input data using at least a portion of the base model; processing the second base model output data using a second gate model associated with a final layer of the base model to determine whether to process the second base model output data using a routing model; determining not to process the second base model output data using the routing model; and generating second final output data from the machine learning model architecture based on the second base model output data.

[0195]

[0186] Clause 31: A method described in any one of clauses 22 to 30, wherein generating second final output data from the machine learning model architecture includes processing second base model output data using a fully connected layer.

[0196]

[0187] Clause 32: A method comprising: training a base model including a plurality of layers using a training dataset; performing clustering on features output from an intermediate layer of the plurality of layers to generate a plurality of training data subsets; training each respective expert model of the plurality of expert models on a respective training data subset of the plurality of training data subsets; training a router model to route training data samples between the plurality of expert models; and training an ensemble model to generate machine learning model architecture output data based on base model output data generated by the base model and expert model output data generated by one or more of the plurality of expert models.

[0197]

[0188] Clause 33: The method described in clause 32, wherein training a router model to route training data samples between a plurality of expert models includes minimizing a first loss component based on a plurality of training data subsets generated by clustering and output from the router model.

[0198]

[0189] Clause 34: The method of clause 32 or 33, wherein training a router model to route training data samples between a plurality of expert models further includes minimizing a second loss component based on an output from the router model and an output from an expert model among the plurality of expert models having the lowest task loss.

[0199]

[0190] Clause 35: The method of any one of clauses 32 to 34, wherein the first loss component comprises a Kullback-Leibler information loss.

[0200]

[0191] Clause 36: The method of any one of clauses 32 to 35, wherein the second loss component comprises a cross-entropy loss.

[0201]

[0192] Clause 37: A method described in any one of clauses 32 to 36, wherein performing clustering on features output from an intermediate layer of the base model to generate multiple training data subsets includes performing K-means clustering on features output from the intermediate layer of the base model.

[0202]

[0193] Clause 38: The method of any one of clauses 32 to 37, wherein the K-means clustering generates a number K of clusters equal to a number of expert models in the plurality of expert models.

[0203]

[0194] Clause 39: A method according to any one of clauses 32 to 38, further comprising training one or more gate models to minimize task loss based on the machine learning model architecture output data.

[0204]

[0195] Clause 40: A processing system comprising a memory containing computer-executable instructions and one or more processors configured to execute the computer-executable instructions and cause the processing system to perform a method according to any one of clauses 1 to 39.

[0205]

[0196] Clause 41: A processing system comprising means for carrying out the method according to any one of clauses 1 to 39.

[0206]

[0197] Clause 42: A non-transitory computer-readable medium comprising computer-executable instructions that, when executed by one or more processors of a processing system, cause the processing system to perform a method according to any one of clauses 1 to 39.

[0207]

[0198] Clause 43: A computer program product embodied on a computer-readable storage medium comprising code for implementing the method according to any one of clauses 1 to 39.

[0208] Additional Considerations

[0199] The foregoing description is provided to enable any person skilled in the art to practice the various aspects described herein. The embodiments described herein are not intended to limit the scope, applicability, or aspects described in the claims. Various modifications to these aspects will be readily apparent to those skilled in the art, and the general principles defined herein may be applied to other aspects. For example, changes may be made in the function and arrangement of the elements described without departing from the scope of the disclosure. Various embodiments may omit, substitute, or add various procedures or components as appropriate. For example, the methods described may be performed in an order different from that described, and various steps may be added, omitted, or combined. Also, features described with respect to some embodiments may be combined in some other embodiments. For example, an apparatus may be implemented or a method may be practiced using any number of the aspects described herein. In addition, the scope of the disclosure is intended to encompass such apparatus or methods practiced using other structures, functions, or structures and functions in addition to or other than the various aspects of the disclosure described herein. It should be understood that any aspect of the disclosure disclosed herein may be embodied by one or more elements of a claim.

[0209]

[0200] As used herein, the word "exemplary" means "serving as an example, instance, or illustration." Any aspect described herein as "exemplary" is not necessarily to be construed as preferred or advantageous over other aspects.

[0210]

[0201] As used herein, a phrase referring to "at least one of" a list of items refers to any combination of those items, including single members. By way of example, "at least one of a, b, or c" is intended to encompass a, b, c, ab, ac, bc, and abc, as well as any combination having multiples of the same element (e.g., aa, aaa, aab, aac, abb, acc, bb, bbb, bbc, cc, and ccc, or any other permutation of a, b, and c).

[0211]

[0202] As used herein, the term "determining" encompasses a wide variety of actions. For example, "determining" can include calculating, computing, processing, deriving, investigating, searching (e.g., searching a table, database, or another data structure), ascertaining, and the like. "Determining" can also include receiving (e.g., receiving information), accessing (e.g., accessing data in a memory), and the like. "Determining" can also include resolving, selecting, choosing, establishing, and the like.

[0212]

[0203] The methods disclosed herein include one or more steps or actions for achieving the method. The steps and / or actions of the methods may be interchanged with one another without departing from the scope of the claims. In other words, unless a specific order of steps or actions is specified, the order of specific steps and / or actions and / or the use of those steps and / or actions may be modified without departing from the scope of the claims. Furthermore, the various operations of the methods described above may be performed by any suitable means capable of performing the corresponding functions. Those means may include various hardware and / or software component(s) including, but not limited to, circuits, application specific integrated circuits (ASICs), or processors, and / or various hardware and / or software modules. In general, when operations are illustrated in a figure, those operations may have corresponding equivalent means-plus-function components numbered similarly.

[0213]

[0204] The following claims are not intended to be limited to the embodiments set forth herein, but are to be accorded the full scope consistent with the language of the claims. Within the claims, reference to an element by the singular is not intended to mean "only one" unless so expressly stated, but rather "one or more." Unless otherwise expressly stated, the term "several" refers to one or more. No element of a claim is to be construed under the provisions of 35 U.S.C. 112(f) unless the element is expressly recited using the phrase "means for" or, in the case of a method claim, unless the element is recited using the phrase "step for." All structural and functional equivalents of the elements of the various embodiments described throughout this disclosure that are known or later become known to those skilled in the art are expressly incorporated herein by reference and are intended to be encompassed by the claims. Furthermore, nothing disclosed herein is intended to be made public, regardless of whether such disclosure is expressly set forth in the claims.

Claims

1. generating base model output data, the generating comprising processing the input data with at least a portion of a base model of the machine learning model architecture; processing the base model output data using a routing model of the machine learning model architecture to determine a selected expert model from a plurality of expert models to use for processing the base model output data; generating expert model output data, the generating comprising processing the base model output data with the selected expert model; generating final output data from the machine learning model architecture, the final output data including processing the base model output data and the expert model output data with an ensemble model of the machine learning model architecture; A computer-implemented method comprising:

2. 2. The computer-implemented method of claim 1, wherein processing the input data with the at least a portion of the base model comprises processing intermediate feature data from at least one layer of the base model with a gate model associated with the at least one layer to determine whether to early terminate the base model.

3. 3. The computer-implemented method of claim 2, further comprising determining to terminate the base model early based on an output from the gate model, wherein the base model output data processed by the routing model includes the output from the gate model.

4. 3. The computer-implemented method of claim 2, further comprising determining not to terminate the base model early based on an output from the gate model, and wherein processing the input data using the at least a portion of the base model comprises processing all layers of the base model to generate the base model output data.

5. The computer-implemented method of claim 2 , wherein the gate model comprises a multi-layer perceptron model.

6. 2. The computer-implemented method of claim 1 , wherein the base model comprises a deep convolutional neural network or a transformer neural network.

7. 2. The computer-implemented method of claim 1, wherein the routing model comprises a linear layer, a convolutional neural network model, or a multi-layer perceptron model.

8. 2. The computer-implemented method of claim 1, wherein the selected expert model comprises a convolutional neural network model or a transformer model.

9. generating second base model output data, the generating comprising processing second input data with at least a portion of the base model; processing the second base model output data with a second gate model associated with a final layer of the base model to determine whether to process the second base model output data using the routing model; determining not to process the second base model output data using the routing model; and generating second final output data from the machine learning model architecture based on the second base model output data; The computer-implemented method of claim 1 , further comprising:

10. 10. The computer-implemented method of claim 9, wherein generating the second final output data from the machine learning model architecture comprises processing the second base model output data using a fully connected layer.

11. 1. A computer-implemented method for training a machine learning model architecture, comprising: Training a base model including a plurality of layers using a training dataset; performing clustering on features output from an intermediate layer of the plurality of layers to generate a plurality of training data subsets; training each respective expert model of a plurality of expert models on a respective training data subset of the plurality of training data subsets; training a router model to route training data samples among the plurality of expert models; training an ensemble model to generate machine learning model architecture output data based on base model output data generated by the base model and expert model output data generated by one or more of the plurality of expert models; A computer-implemented method comprising:

12. 12. The computer-implemented method of claim 11, wherein training the router model to route the training data samples among the plurality of expert models comprises minimizing a first loss component based on the plurality of training data subsets generated by the clustering and output from the router model.

13. 13. The computer-implemented method of claim 12, wherein training the router model to route the training data samples among the plurality of expert models further comprises minimizing a second loss component based on the output from the router model and an output from the expert model of the plurality of expert models having a lowest task loss.

14. 14. The computer-implemented method of claim 13, wherein the first loss component comprises a Kullback-Leibler information loss.

15. The computer-implemented method of claim 14 , wherein the second loss component comprises a cross-entropy loss.

16. 12. The computer-implemented method of claim 11 , wherein performing the clustering on the features output from the intermediate layer of the base model to generate the plurality of training data subsets comprises performing K-means clustering on the features output from the intermediate layer of the base model.

17. 17. The computer-implemented method of claim 16, wherein the K-means clustering generates a number of clusters, K, equal to a number of expert models in the plurality of expert models.

18. 12. The computer-implemented method of claim 11, further comprising training one or more gating models to minimize task loss based on the machine learning model architecture output data.

19. a memory containing computer executable instructions; one or more processors configured to execute the computer-executable instructions and cause the processing system to perform operations; a processing system comprising: generating base model output data, the generating comprising processing the input data with at least a portion of a base model of the machine learning model architecture; processing the base model output data using a routing model of the machine learning model architecture to determine a selected expert model from a plurality of expert models to use for processing the base model output data; generating expert model output data, the generating comprising processing the base model output data with the selected expert model; generating final output data from the machine learning model architecture, the final output data including processing the base model output data and the expert model output data with an ensemble model of the machine learning model architecture; a processing system comprising:

20. 20. The processing system of claim 19, wherein processing the input data with the at least a portion of the base model comprises processing intermediate feature data from at least one layer of the base model with a gate model associated with the at least one layer to determine whether to early terminate the base model.

21. 21. The processing system of claim 20, wherein the operations further comprise determining to early terminate the base model based on an output from the gate model, and wherein the base model output data processed by the routing model comprises the output from the gate model.

22. 21. The processing system of claim 20, wherein the operations further comprise determining not to terminate the base model early based on an output from the gate model, and wherein processing the input data using the at least a portion of the base model comprises processing all layers of the base model to generate the base model output data.

23. The processing system of claim 20 , wherein the gate model comprises a multi-layer perceptron model.

24. 20. The processing system of claim 19, wherein the base model comprises a deep convolutional neural network or a transformer neural network.

25. 20. The processing system of claim 19, wherein the routing model comprises a linear layer, a convolutional neural network model, or a multi-layer perceptron model.

26. 20. The processing system of claim 19, wherein the selected expert model comprises a convolutional neural network model or a transformer model.

27. The operation, generating second base model output data, the generating comprising processing second input data with at least a portion of the base model; processing the second base model output data with a second gate model associated with a final layer of the base model to determine whether to process the second base model output data using the routing model; determining not to process the second base model output data using the routing model; and generating second final output data from the machine learning model architecture based on the second base model output data; 20. The processing system of claim 19, further comprising:

28. 30. The processing system of claim 27, wherein generating the second final output data from the machine learning model architecture comprises processing the second base model output data using a fully connected layer.

29. means for generating base model output data, the means being configured to process the input data with at least a portion of a base model of the machine learning model architecture; means for processing the base model output data using a routing model of the machine learning model architecture to determine a selected expert model from a plurality of expert models to use for processing the base model output data; means for generating expert model output data, the means being configured to process the base model output data with the selected expert model; a means for generating final output data from the machine learning model architecture, the means being configured to process the base model output data and the expert model output data with an ensemble model of the machine learning model architecture; A processing system comprising:

30. means for generating second base model output data, the means being configured to process second input data using at least a portion of the base model; means for processing the second base model output data using a gate model associated with a final layer of the base model to determine whether to process the second base model output data using the routing model; means for determining not to process the second base model output data using the routing model; means for generating second final output data from the machine learning model architecture based on the second base model output data; 30. The processing system of claim 29, further comprising:

Images