WO2026015992A1

WO2026015992A1 - Operating mixture-of-experts (moe) architectures using neural processing units

Info

Publication number: WO2026015992A1
Application number: PCT/CN2024/105375
Authority: WO
Inventors: Changzhou HE; Xun JIANG; Liang Zhang; Sha HE; Xiaolin Gu; Yifei Zhang; Veluppillai Arulesan; Xiaojun SANG; Qinyu WANG; Kanghwan JANG; Jian Wang; Milan OLJACA; Jeffrey Baginsky Gehlhaar
Original assignee: Qualcomm Inc
Current assignee: Qualcomm Inc
Priority date: 2024-07-13
Filing date: 2024-07-13
Publication date: 2026-01-22
Anticipated expiration: 2027-01-13

Abstract

A processor-implemented method includes receiving a machine learning (ML) model. The ML model includes multiple expert models. The ML model is adapted to include an expert placeholder. The adapted ML model dynamically selects one or more expert models based on an input. The selected one or more expert models is included in the expert placeholder. The adapted ML model processes the input based on the dynamically selected one or more expert models to generate an output.

Description

OPERATING MIXTURE-OF-EXPERTS (MoE) ARCHITECTURES USING NEURAL PROCESSING UNITS

FIELD OF THE DISCLOSURE

Aspects of the present disclosure generally relate to artificial neural network models, and more specifically to mixture-of-experts (MoE) architectures using neural processing units.

BACKGROUND

Artificial neural networks may comprise interconnected groups of artificial neurons (e.g., neuron models) . The artificial neural network (ANN) may be a computational device or be represented as a method to be performed by a computational device. Various ANN model structures are available for consideration. Convolutional neural networks (CNNs) are a type of feed-forward ANN. Convolutional neural networks may include collections of neurons that each have a receptive field and that collectively tile an input space.

A transformer ANN structure makes use of attention mechanisms that may enable the model to process input sequences in a parallel and efficient manner. An attention mechanism allows the model to focus on different parts of the input sequence at different times. A transformer ANN structure may be of particular use for tasks that involve sequence modeling, such as text generation by a large language model (LLM) . However, training and operating such large scale models (e.g., including one billion parameters) is challenging and time consuming, particularly with resource-constrained devices such as mobile devices (e.g., smartphones, tablet computing devices) , edge processing devices, and Internet-of-Things (IoT) devices, for instance.

SUMMARY

The present disclosure is set forth in the independent claims, respectively. Some aspects of the disclosure are described in the dependent claims.

In some aspects of the present disclosure, a processor-implemented method includes receiving a machine learning (ML) model including multiple expert models. The processor-implemented method also includes adapting the ML model to include an expert placeholder. The processor-implemented method additionally includes dynamically selecting, by the adapted ML model, one or more expert models based on an input. The selected one or more expert models is included in the expert placeholder. The processor-implemented method further includes processing, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output.

Some aspects of the present disclosure are directed to an apparatus. The apparatus includes means for receiving a machine learning (ML) model including multiple expert models. The apparatus also includes means for adapting the ML model to include an expert placeholder. The apparatus additionally includes means for dynamically selecting, by the adapted ML model, one or more expert models based on an input. The selected one or more expert models is included in the expert placeholder. The apparatus further includes means for processing, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output.

Some aspects of the present disclosure are directed to an apparatus. The apparatus has at least one memory and one or more processors coupled to the at least one memory. The processor (s) is configured to receive a machine learning (ML) model including multiple expert models. The processor (s) is also configured to adapt the ML model to include an expert placeholder. The processor (s) is additionally configured to dynamically select, by the adapted ML model, one or more expert models based on an input. The selected one or more expert models is included in the expert placeholder. The processor (s) is further configured to process, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output.

Additional features and advantages of the disclosure will be described below. It should be appreciated by those skilled in the art that this disclosure may be readily utilized as a basis for modifying or designing other structures for carrying out the same purposes of the present disclosure. It should also be realized by those skilled in the art that such equivalent constructions do not depart from the teachings of the disclosure as set forth in the appended claims. The novel features, which are believed to be characteristic of the disclosure, both as to its organization and method of operation, together with further objects and advantages, will be better understood from the following description when considered in connection with the accompanying figures. It is to be expressly understood, however, that each of the figures is provided for the purpose of illustration and description only and is not intended as a definition of the limits of the present disclosure.

BRIEF DESCRIPTION OF THE DRAWINGS

The features, nature, and advantages of the present disclosure will become more apparent from the detailed description set forth below when taken in conjunction with the drawings in which like reference characters identify correspondingly throughout.

FIGURE 1 illustrates an example implementation of a neural network using a system-on-a-chip (SOC) , including a general-purpose processor in accordance with certain aspects of the present disclosure.

FIGURE 2 is an illustrative block diagram of an example machine learning (ML) model represented by an artificial neural network (ANN) , in accordance with various aspects of the present disclosure.

FIGURE 3 is a block diagram illustrating an exemplary software architecture that may modularize artificial intelligence (AI) functions, in accordance with various aspects of the present disclosure.

FIGURES 4A-4B are block diagrams illustrating an example switch transformer encoder block, in accordance with various aspects of the present disclosure.

FIGURE 5 is a block diagram illustrating an example mixture-of-experts (MoE) architecture, in accordance with various aspects of the present disclosure.

FIGURE 6 is a block diagram illustrating another example MoE architecture, in accordance with various aspects of the present disclosure.

FIGURE 7 is a flow diagram illustrating a processor-implemented method for operating an MoE architecture, in accordance with various aspects of the present disclosure.

DETAILED DESCRIPTION

The detailed description set forth below, in connection with the appended drawings, is intended as a description of various configurations and is not intended to represent the only configurations in which the concepts described may be practiced. The detailed description includes specific details for the purpose of providing a thorough understanding of the various concepts. However, it will be apparent to those skilled in the art that these concepts may be practiced without these specific details. In some instances, well-known structures and components are shown in block diagram form in order to avoid obscuring such concepts.

Based on the teachings, one skilled in the art should appreciate that the scope of the disclosure is intended to cover any aspect of the disclosure, whether implemented independently of or combined with any other aspect of the disclosure. For example, an apparatus may be implemented or a method may be practiced using any number of the aspects set forth. In addition, the scope of the disclosure is intended to cover such an apparatus or method practiced using other structure, functionality, or structure and functionality in addition to or other than the various aspects of the disclosure set forth. It should be understood that any aspect of the disclosure disclosed may be embodied by one or more elements of a claim.

The word “exemplary” is used to mean “serving as an example, instance, or illustration. ” Any aspect described as “exemplary” is not necessarily to be construed as preferred or advantageous over other aspects.

Although particular aspects are described, many variations and permutations of these aspects fall within the scope of the disclosure. Although some benefits and advantages of the preferred aspects are mentioned, the scope of the disclosure is not intended to be limited to particular benefits, uses or objectives. Rather, aspects of the disclosure are intended to be broadly applicable to different technologies, system configurations, networks, and protocols, some of which are illustrated by way of example in the figures and in the following description of the preferred aspects. The detailed description and drawings are merely illustrative of the disclosure rather than limiting, the scope of the disclosure being defined by the appended claims and equivalents thereof.

Mixture-of-experts (MoE) may be considered a form of conditional computing architecture for machine learning models that enables processing relevant portions of a model architecture while bypassing unnecessary portions of the model architecture. In particular, sparse MoE models may dynamically select which parameters of the network to use, conditioned on the input. Compared to dense models, sparse MoE models can vastly expand their number of parameters and improve performance, while keeping the computation costs per input similar. As a result, MoE technology has increasingly been employed in large language models (LLM) or large vision models (LVM) , because MoE may greatly reduce the cost of model training and accelerate inference speed. Moreover, MoE may enable models to be pretrained with substantially fewer computations.

Although MoE may greatly reduce the cost of model training and accelerate inference speed, MoE may not directly operate on some neural processing units. One reason for such inoperability may be due to the conditional control of expert selection. That is, in some cases, only experts with high scores from a router, such as top-half, will be activated dynamically. However, some neural processing units may be configured to run all experts, including expert models that are not activated (e.g., not selected) each time because those expert models are included within the machine learning model.

Accordingly, to address these and other problems, aspects of the present disclosure are directed to enabling MoE in machine learning models. In accordance with aspects of the present disclosure, a set of experts may be selected in each block of the machine learning model. An expert placeholder may be included in each block. In some aspects, weights of the experts in each feed-forward network (FFN) layer may be static constants. The model may dynamically select the set of experts weights. For instance, in some aspects the top-K experts may be dynamically selected.

In some aspects, the dynamic selection may be implemented using a gather operation, for example. A gather operation may receive inputs from two or more input streams (e.g., example experts) and may concatenate entries (which may be referred to as tuples) that share a same key value. In one example, a gather-operator (or similar operator) may take the top-K [index, values] and all experts’ static constant weights as inputs and may output the selected expert weights.

Particular aspects of the subject matter described in this disclosure can be implemented to realize one or more of the following potential advantages. In some examples, the adapting of a machine learning model to include an expert placeholder and dynamically selecting, by an adapted machine learning (ML) model, one or more expert models based on an input, the selected one or more expert models being included in the expert placeholder may increase processing efficiency and reduce model training costs.

FIGURE 1 illustrates an example implementation of a system-on-a-chip (SOC) 100, which may include a central processing unit (CPU) 102 or a multi-core CPU, as well as a graphics processing unit (GPU) 104, and/or a neural processing unit (NPU) 108 configured for operating a mixture-of-experts model. Variables (e.g., neural signals and synaptic weights) , system parameters associated with a computational device (e.g., neural network with weights) , delays, frequency bin information, and task information may be stored in a memory block associated with the NPU 108, in a memory block associated with the CPU 102, in a memory block associated with the GPU 104, in a memory block associated with a digital signal processor (DSP) 106, in a memory block 118, or may be distributed across multiple blocks. Instructions executed at the CPU 102 may be loaded from a program memory associated with the CPU 102 or may be loaded from a memory block 118.

The SOC 100 may also include additional processing blocks tailored to specific functions, such as a GPU 104, a DSP 106, a connectivity block 110, which may include fifth generation (5G) connectivity, fourth generation long term evolution (4G LTE) connectivity, Wi-Fi connectivity, USB connectivity, Bluetooth connectivity, and the like, and a multimedia processor 112 that may, for example, detect and recognize gestures. In one implementation, the NPU 108 is implemented in the CPU 102, DSP 106, and/or GPU 104. The SOC 100 may also include a sensor processor 114, image signal processors (ISPs) 116, and/or navigation module 120, which may include a global positioning system.

The SOC 100 may be based on an ARM, RISC-V (RISC-five) , or any reduced instruction set computing (RISC) architecture. In aspects of the present disclosure, the instructions loaded into the general-purpose processor 102 may include code to receive a machine learning (ML) model. The ML model includes multiple expert models. The instructions loaded into the general-purpose processor 102 may also include code to adapt the ML model to include an expert placeholder. The instructions loaded into the general-purpose processor 102 may also include code to dynamically select, by the adapted ML model, one or more expert models based on an input. The selected one or more expert models is included in the expert placeholder. The instructions loaded into the general-purpose processor 102 may further include code to process, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output.

Neural networks may be designed with a variety of connectivity patterns. In feed-forward networks, information is passed from lower to higher layers, with each neuron in a given layer communicating to neurons in higher layers. A hierarchical representation may be built up in successive layers of a feed-forward network, as described above. Neural networks may also have recurrent or feedback (also called top-down) connections. In a recurrent connection, the output from a neuron in a given layer may be communicated to another neuron in the same layer. A recurrent architecture may be helpful in recognizing patterns that span more than one of the input data chunks that are delivered to the neural network in a sequence. A connection from a neuron in a given layer to a neuron in a lower layer is called a feedback (or top-down) connection. A network with many feedback connections may be helpful when the recognition of a high-level concept may aid in discriminating the particular low-level features of an input.

FIGURE 2 is an illustrative block diagram of an example machine learning (ML) model represented by an artificial neural network (ANN) 200. The ANN 200 may receive input data 206, which may include one or more bits of data 202, pre-processed data output from a pre-processor 204 (optional) , or some combination thereof. Here, data 202 may include training data, verification data, application-related data, or the like, based, for example, on the stage of deployment of the ANN 200. The pre-processor 204 may be included within the ANN 200 in some other implementations. The pre-processor 204 may, for example, process all or a portion of the data 202, which may result in some of the data 202 being changed, replaced, deleted, etc. In some implementations, the pre-processor 204 may add additional data to the data 202.

The ANN 200 includes at least one first layer 208 of artificial neurons 210 to process input data 206 and provide resulting first layer data via connections or “edges” such as the edges 212 to at least a portion of at least one second layer 214. The second layer 214 processes data received via the edges 212 and provides second layer output data via the edges 216 to at least a portion of at least one third layer 218. The third layer 218 processes data received via the edges 216 and provides third layer output data via the edges 220 to at least a portion of a final layer 222 including one or more neurons to provide output data 224. All or part of the output data 224 may be further processed in some manner by an optional post-processor 226. Thus, in certain examples, the ANN 200 may provide output data 228 that is based on output data 224, post-processed data output from the post-processor 226, or some combination thereof.

The post-processor 226 may be included within the ANN 200 in some other implementations. The post-processor 226 may, for example, process all or a portion of the output data 224, which may result in the output data 228 being different, at least in part, to the output data 224, as result of data being changed, replaced, deleted, etc. In some implementations, the post-processor 226 may be configured to add additional data to the output data 224. In this example, the second layer 214 and third layer 218 represent intermediate or hidden layers arranged in a hierarchical or other like structure. Although not explicitly shown, there may be one or more further intermediate layers between the second layer 214 and the third layer 218.

The structure and training of artificial neurons 210 in the various layers may be tailored to specific requirements of an application. Within a given layer such as the first layer 208, second layer 214, or third layer 218 of the ANN 200, some or all of the neurons may be configured to process information provided to the layer and output corresponding transformed information from the layer. For example, transformed information from a layer may represent a weighted sum of the input information associated with or otherwise based on a non-linear activation function or other activation function used to “activate” artificial neurons of a next layer. Artificial neurons in such a layer may be activated by or be responsive to parameters such as the previously described weights and biases of the ANN 200. The weights and biases of the ANN 200 may be adjusted during a training process or during operation of the ANN 200. The weights of the various artificial neurons may control a strength of connections between layers or artificial neurons, while the biases may control a direction of connections between the layers or artificial neurons. An activation function may select or determine whether an artificial neuron transmits its output to the next layer or not in response to its received data.

Different activation functions may model different types of non-linear relationships. By introducing non-linearity into an ML model, an activation function allows the configuration for the ML model to change in response to identifying or detecting complex patterns and relationships in the input data 206. Some non-exhaustive example activation functions include a sigmoid based activation function, a hyperbolic tangent (tanh) based activation function, a convolutional activation function, up-sampling, pooling, and a rectified linear unit (ReLU) based activation function.

Training of an ML model, such as the ANN 200, may be conducted using training data. Training data may include one or more datasets the ANN 200 may use to identify patterns or relationships. Training data may represent various types of information, including written, visual, audio, environmental context, operational properties, etc. During training, the parameters (such as the weights and biases) of artificial neurons 210 may be changed, such as to minimize or otherwise reduce a loss function or a cost function. A training process may repeat multiple times to fine-tune the ANN 200 with each iteration.

Various ANN model structures are available for consideration. For example, in a feed-forward ANN structure, each artificial neuron 210 in the second layer 214 receives information from the previous layer (such as, one or more artificial neurons 210 in the first layer 208) and produces information for the next layer (such as, one or more artificial neurons 210 in the third layer 218) . In a convolutional ANN structure, some layers may be organized into filters that extract features from data, such as the training data or the input data. In a recurrent ANN structure, some layers may have connections that allow for processing of data across time, such as for processing information having a temporal structure, such as time series data forecasting.

A transformer ANN structure makes use of attention mechanisms that may enable the model to process input sequences in a parallel and efficient manner. An attention mechanism allows the model to focus on different parts of the input sequence at different times. Attention mechanisms may be implemented using a series of layers known as attention layers to compute weighted sums of input features based on a similarity between different elements of the input sequence. A transformer ANN structure may include a series of feed-forward ANN layers whose configurations may change in response to identifying non-linear relationships between the input and output sequences, which may also be referred to as a process of “learning” by the ANN layers. The output of a transformer ANN structure may be obtained by applying a linear transformation to the output of a final attention layer. A transformer ANN structure may be of particular use for tasks that involve sequence modeling, or other like processing, such as text generation. A large language model may be a particularly useful implementation of a transformer ANN structure.

FIGURE 3 is a block diagram illustrating an exemplary software architecture 300 that may modularize artificial intelligence (AI) functions. Using the architecture 300, applications may be designed that may cause various processing blocks of an SOC 320 (for example a CPU 322, a DSP 324, a GPU 326 and/or an NPU 328) (which may be similar to SOC 100 of FIGURE 1) to operate a mixture-of-experts model for an AI application 302, according to aspects of the present disclosure. The architecture 300 may, for example, be included in a computational device, such as a smartphone.

The AI application 302 may be configured to call functions defined in a user space 304 that may, for example, provide for text, video, and/or sound generation. The AI application 302 may make a request to compiled program code associated with a library defined in an AI function application programming interface (API) 306. This request may ultimately rely on the output of a deep neural network configured to provide an inference response based on input, for example.

The run-time engine 308, which may be compiled code of a runtime framework, may be further accessible to the AI application 302. The AI application 302 may cause the run-time engine 308, for example, to request an inference at a particular time interval or triggered by an event detected by the user interface of the AI application 302. When caused to provide an inference response, the run-time engine 308 may in turn send a signal to an operating system in an operating system (OS) space 310, such as a Kernel 312, running on the SOC 320. In some examples, the Kernel 312 may be a LINUX Kernel. The operating system, in turn, may cause non-contiguous attention masks to be processed on the CPU 322, the DSP 324, the GPU 326, the NPU 328, or some combination thereof. The CPU 322 may be accessed directly by the operating system, and other processing blocks may be accessed through a driver, such as a driver 314, 316, or 318 for, respectively, the DSP 324, the GPU 326, or the NPU 328. In the exemplary example, the deep neural network may be configured to run on a combination of processing blocks, such as the CPU 322, the DSP 324, and the GPU 326, or may be run on the NPU 328.

As described, aspects of the present disclosure are directed to enabling MoE in machine learning models. In accordance with aspects of the present disclosure, a set of experts may be selected in each block of the machine learning model. An expert placeholder may be included in each block. In some aspects, weights of the experts in each feed forward network (FFN) layer may be static constants. The model may dynamically select the set of experts weights. For instance, in some aspects, the top-K experts may be dynamically selected.

In some aspects, the dynamic selection may be implemented using a gather operation, for example. A gather operation may receive inputs from two or more input streams (e.g., expert models) and may concatenate entries (which may be referred to as tuples) that share a same key value. In one example, a gather-operator (or similar operator) may take top-K [index, values] and all experts’ static constant weights as inputs and may output the selected-expert’s weight.

The techniques of the present disclosure may beneficially provide flexibility and may be implemented in a source model (e.g., open neural network exchange (ONNX) /TensorFlow) or at a model-converter level.

FIGURES 4A-4B are block diagrams illustrating an example switch transformer encoder block 400, in accordance with various aspects of the present disclosure. As shown in FIGURE 4A, the example switch transformer encoder block 400 includes alternating layers of self-attention blocks 402 and switching feed-forward network (FFN) blocks 404. Normalization blocks 406 (e.g., 406a, 406b) may be applied after every block (e.g., 402, 404) of the example transformer encoder block 400.

FIGURE 4B is a block diagram illustrating the example transformer encoder block 400 with an expanded view of the switching FFN block 404, in accordance with aspects of the present disclosure. As shown in FIGURE 4B, the switching FFN blocks 404 may include one or more router models 420a, 420b. The switching FFN blocks 404 may be configured to implement a mixture-of-experts (MoE) . Router models 420a, 420b may receive the output of the normalization block 406a and may determine to select one or more expert models FFN1-FFN4 to process an input token. The self-attention blocks 402 and switching FFN blocks 404 may each include a residual connection.

Each layer of the example switch transformer encoder block 400 may operate independently on tokens in a sequence. In the example of FIGURE 4B, inputs 422a, 422b (e.g., “More” and “Parameters” ) may be received by the example switch transformer encoder block 400. The inputs 422a, 422b may be tokenized and provided to the self-attention block 402. The tokens (x1 = “More” and x2 = “Parameters” ) may be routed (solid lines) across the FFN expert models FFN1-FFN4. In some aspects, the routers 420a, 420b may independently route each token to the FFN expert models FFN1-FFN4. In turn, the switch FFN block 404 may return the output of the selected FFN expert model FFN1-FFN4 multiplied by the router gate value (e.g., 0.65 or 0.8) and used to generate the model outputs 424a, 424b (e.g., y1, y2) . The model outputs 424a, 424b may, for instance comprise a following token for subsequent processing or an inference.

FIGURE 5 is a block diagram illustrating an example MoE architecture 500, in accordance with various aspects of the present disclosure. Referring to FIGURE 5, the example MoE architecture 500 may include a self-attention block 402 and a switching FFN block 502. Although the MoE architecture in the example of FIGURE 5 is a transformer-based architecture, this is merely an example and not limiting. In addition, while one switching transformer encoder block 400 is shown, this is merely an example and not limiting. Rather, any number of switching transformer encoder block 400 may be included in the MoE architecture 500.

The switching FFN block 502 may include a router model 504 and an expert placeholder 506. The example MoE architecture 500 may also include a set of multiple expert models 508a for up projection and a set of multiple expert models 508b for down projection. Each of the expert models 508a, 508b may be trained on a subset of a training dataset. In some aspects, the expert models 508a, 508b may be trained on a different subset of the of the training dataset. The weights of the expert models 508a, 508b may be static during inference.

In operation, the self-attention block 402 may receive an input token along with a positional embedding and may generate a set of attention scores. The attention scores may be provided to the normalization layer (e.g., 406a) to perform a normalization process. For instance, the normalization layer (e.g., 406a) may apply a softmax function to generate a set of logits that may, in turn, be provided to the router model 504. The router model 504 may apply a selection function (e.g., Top-K) to dynamically determine a subset of the expert models of the multiple expert models 508a (e.g., up projection) to select based on the input. The router model 504 may also apply a selection function (e.g., Top-K) to dynamically determine a subset of the expert models of the multiple expert models 508b (e.g., down projection) to select based on the input.

The weights of the selected expert models 508a (e.g., up projection) and expert models 508b (e.g., down projection) may each be collected using a gather operation (e.g., 510a, 510b) . Then, the collected expert model weights may be inserted into the expert placeholder 506. The selected expert model weights may be multiplied by the respective router gate value (e.g., 0.85 or 0.8) and output to the normalization layer 406b to generate an output token.

In some aspects, the example MoE architecture 500 may also include one or more shared expert models 512, which may be included with the dynamically selected expert models to further improve model performance (e.g., increase model accuracy) . Moreover, the MoE architecture 500 may beneficially enable operation on an NPU without splitting the machine learning model graph into numerous smaller model graphs, which may reduce model efficiency.

FIGURE 6 is a block diagram illustrating another example MoE architecture 600, in accordance with various aspects of the present disclosure. Referring to FIGURE 6, the example MoE architecture 600 includes the switching transformer encoder block 400. Although, one switching transformer encoder block 400 is shown, this is merely an example. It should be understood that the MoE architecture may include many switching transformer encoder blocks 400. However, the switching transformer encoder block 400 may be split such that a first part includes the self-attention block 402 and the router models 420a, 420b (illustrated with the circled 1) and a second part includes multiple expert models FFN1-FFN4 (illustrated with the circled 2) .

In part 1, the self-attention block 402 and router models 420a, 420b may receive as inputs the previous self-attention block 402 and router models 420a, 420b hidden states, the K value, top-K weights, and the hidden state of each expert model. part 1 may output the hidden state and the top-k weights. Accordingly, part 1 may absorb “Top-K, Top-K weights” to a post-router model and may absorb a “weighted sum” to a pre-attention block.

In part 2, each switching FFN block (e.g., 404) may be split into multiple expert models. In one example, each switching FFN block (e.g., 404) may be split into eight expert models such that the whole MoE model (e.g., including 32 switching transformer blocks (e.g., 400) ) may include two hundred eighty-eight expert models (e.g., 32 blocks *8 split experts in each block + 32 blocks *1 part-1 model in each block) . As such, control flow or graph selection may be determined by the router and may be conducted using the CPU (e.g., 102, 422) . The example MoE architecture 600 of FIGURE 6 may, however, introduce frequent model switching as well as memory read/write (e.g., remote procedure call (RPC) memory read/write) operations for input/output tensors across different graphs in comparison to the example MoE architecture 500 of FIGURE 5.

In some aspects, the expert model weights may be used as inputs to reduce a number of graphs. For instance, a selected expert model weights and may be employed to reduce the number of graphs from 288 to 64. That is, by using the weights as inputs, each switching FFN block (e.g., 404) may still be split into eight expert models, but the eight expert models in each switching FFN block (e.g., 404) may share an expert-placeholder model. As a result, the number of graphs may be reduced from 288 to 64 (e.g., 32 blocks *1 expert-placeholder in each block + 32 blocks *1 part-1 model in each block) .

Each such expert model may be represented by a different model graph. The switching transformer block (e.g., 400) may be transformed in a manner similar to that shown in FIGURE 5 (e.g., 500) . That is, each switching transformer block (e.g., 500) may include an expert-placeholder block (e.g., 400) , because the expert models may have the same architecture. However, an expert’s weights may not be stored inside the expert model. Rather, the weights of each expert model may be stored as const tensors outside the expert model (e.g., may be stored in memory 118) , and may be dynamically loaded into the expert placeholder model as one of the expert-placeholder’s input tensors according to router’s output (e.g., the selected expert model) .

FIGURE 7 is a flow diagram illustrating a processor-implemented method 700 for operating a mixture-of-experts (MoE) architecture, in accordance with various aspects of the present disclosure. The processor-implemented method 700 may be performed by one or more processors such as the NPU (e.g., 108, 328) , GPU (e.g., 104, 326) , and/or other processing unit (e.g., DSP 106, 324 or CPU 102, 322) , for example.

As shown in FIGURE 7, at block 702, the processor receives a machine learning (ML) model including multiple expert models. As shown in FIGURES 4A-4B, a switch transformer encoder block 400 may be received. The switch transformer encoder block 400 may be received includes alternating layers of self-attention blocks 402 and switching feed-forward network (FFN) blocks 404. Normalization blocks 406 (e.g., 406a, 406b) may be applied after every block (e.g., 402, 404) of the example transformer encoder block 400. The switching FFN blocks 404 may include one or more router models 420a, 420b. The switching FFN blocks 404 may be configured to implement a mixture-of-experts (MoE) .

At block 704, the processor adapts the ML model to include an expert placeholder. For example, as described with reference to FIGURE 5, The switching FFN block 502 may be adapted to include a router model 504 and an expert placeholder 506. The example MoE architecture 500 may also include a set of multiple expert models 508a for up projection and a set of multiple expert models 508b for down projection.

At block 706, the processor dynamically selects, by the adapted ML model, one or more expert models based on an input, the selected one or more expert models being included in the expert placeholder. For instance, as described with reference to FIGURE 5, the router model 504 may apply a selection function (e.g., Top-K) to dynamically determine a subset of the expert models of the multiple expert models 508a (e.g., up projection) to select based on the input. The router model 504 may also apply a selection function (e.g., Top-K) to dynamically determine a subset of the expert models of the multiple expert models 508b (e.g., down projection) to select based on the input.

At block 708, the processor processes, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output. As described, for example, with reference to FIGURE 5, the weights of the selected expert models 508a (e.g., up projection) and expert models 508b (e.g., down projection) may each be collected using a gather operation (e.g., 510a, 510b) . Then, the collected expert model weights may be inserted into the expert placeholder 506. The selected expert model weights may be multiplied by the respective router gate value (e.g., 0.85 or 0.8) and output to the normalization layer 406b to generate an output token.

Example Aspects

Aspect 1: An apparatus, comprising:

at least one memory; and

at least one processor coupled to the at least one memory, the at least one processor configured:

to receive a machine learning (ML) model including multiple expert models;

to adapt the ML model to include an expert placeholder;

to dynamically select, by the adapted ML model, one or more expert models based on an input, the selected one or more expert models being included in the expert placeholder; and

to process, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output.

Aspect 2: The apparatus of Aspect 1, in which the dynamically selected one or more expert models indicate expert model weights that are applied to the input to compute an output.

Aspect 3: The apparatus of Aspect 1 or 2, in which the multiple expert models are trained on different subsets of a common training dataset.

Aspect 4: The apparatus of any preceding Aspect, in which the at least one processor is further configured to select the one or more expert models using a selection function based on the input at inference time.

Aspect 5: The apparatus of any preceding Aspect, in which the at least one processor comprises one or more neural processing units.

Aspect 6: The apparatus of any preceding Aspect, in which the ML model comprises a transformer-based model.

Aspect 7: A processor-implemented method performed by at least one processor, comprising:

receiving a machine learning (ML) model including multiple expert models;

adapting the ML model to include an expert placeholder;

dynamically selecting, by the adapted ML model, one or more expert models based on an input, the selected one or more expert models being included in the expert placeholder; and

processing, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output.

Aspect 8: The processor-implemented method of Aspect 7, in which the dynamically selected one or more expert models indicate expert model weights that are applied to the input to compute an output.

Aspect 9: The processor-implemented method of Aspect 7 or 8, in which the multiple expert models are trained on different subsets of a common training dataset.

Aspect 10: The processor-implemented method of any of Aspects 7-9, in which the one or more expert models are selected using a selection function based on the input at inference time.

Aspect 11: The processor-implemented method of any of Aspects 7-10, in which the processing is performed by one or more neural processing units.

Aspect 12: The processor-implemented method of any of Aspects 7-11, in which the ML model comprises a transformer-based model.

Aspect 13: An apparatus, comprising:

means for receiving a machine learning (ML) model including multiple expert models;

means for adapting the ML model to include an expert placeholder;

means for dynamically selecting, by the adapted ML model, one or more expert models based on an input, the selected one or more expert models being included in the expert placeholder; and

means for processing, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output.

Aspect 14: The apparatus of Aspect 13, in which the dynamically selected one or more expert models indicate expert model weights that are applied to the input to compute an output.

Aspect 15: The apparatus of Aspect 13 or 14, in which the multiple expert models are trained on different subsets of a common training dataset.

Aspect 16: The apparatus of any of Aspects 13-15, further comprising means for selecting the one or more expert models using a selection function based on the input at inference time.

Aspect 17: The apparatus of any of Aspects 13-16, in which the processing is performed by one or more neural processing units.

Aspect 18: The apparatus of any of Aspects 13-17, in which the ML model comprises a transformer-based model.

In one aspect, the receiving means, adapting means, dynamically selecting means and/or processing means may be the NPU (108, 328) , the GPU (104, 326) program memory associated with the NPU (108, 328) or the GPU (104, 326) , the dedicated memory block 118, and/or the CPU (102, 322) configured to perform the functions recited. In another configuration, the aforementioned means may be any module or any apparatus configured to perform the functions recited by the aforementioned means.

The various operations of methods described above may be performed by any suitable means capable of performing the corresponding functions. The means may include various hardware and/or software component (s) and/or module (s) , including, but not limited to, a circuit, an application specific integrated circuit (ASIC) , or processor. Generally, where there are operations illustrated in the figures, those operations may have corresponding counterpart means-plus-function components with similar numbering.

As used, the term “determining” encompasses a wide variety of actions. For example, “determining” may include calculating, computing, processing, deriving, investigating, looking up (e.g., looking up in a table, a database, or another data structure) , ascertaining and the like. Additionally, “determining” may include receiving (e.g., receiving information) , accessing (e.g., accessing data in a memory) and the like. Furthermore, “determining” may include resolving, selecting, choosing, establishing, and the like.

As used, a phrase referring to “at least one of” a list of items refers to any combination of those items, including single members. As an example, “at least one of: a, b, or c” is intended to cover: a, b, c, a-b, a-c, b-c, and a-b-c.

The various illustrative logical blocks, modules and circuits described in connection with the present disclosure may be implemented or performed with a general-purpose processor, a digital signal processor (DSP) , an application specific integrated circuit (ASIC) , a field programmable gate array signal (FPGA) or other programmable logic device (PLD) , discrete gate or transistor logic, discrete hardware components or any combination thereof designed to perform the functions described. A general-purpose processor may be a microprocessor, but in the alternative, the processor may be any commercially available processor, controller, microcontroller, or state machine. A processor may also be implemented as a combination of computing devices, e.g., a combination of a DSP and a microprocessor, a plurality of microprocessors, one or more microprocessors in conjunction with a DSP core, or any other such configuration.

The steps of a method or algorithm described in connection with the present disclosure may be embodied directly in hardware, in a software module executed by a processor, or in a combination of the two. A software module may reside in any form of storage medium that is known in the art. Some examples of storage media that may be used include random access memory (RAM) , read only memory (ROM) , flash memory, erasable programmable read-only memory (EPROM) , electrically erasable programmable read-only memory (EEPROM) , registers, a hard disk, a removable disk, a CD-ROM and so forth. A software module may comprise a single instruction, or many instructions, and may be distributed over several different code segments, among different programs, and across multiple storage media. A storage medium may be coupled to a processor such that the processor can read information from, and write information to, the storage medium. In the alternative, the storage medium may be integral to the processor.

The methods disclosed comprise one or more steps or actions for achieving the described method. The method steps and/or actions 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 and/or use of specific steps and/or actions may be modified without departing from the scope of the claims.

The functions described may be implemented in hardware, software, firmware, or any combination thereof. If implemented in hardware, an example hardware configuration may comprise a processing system in a device. The processing system may be implemented with a bus architecture. The bus may include any number of interconnecting buses and bridges depending on the specific application of the processing system and the overall design constraints. The bus may link together various circuits including a processor, machine-readable media, and a bus interface. The bus interface may be used to connect a network adapter, among other things, to the processing system via the bus. The network adapter may be used to implement signal processing functions. For certain aspects, a user interface (e.g., keypad, display, mouse, joystick, etc. ) may also be connected to the bus. The bus may also link various other circuits such as timing sources, peripherals, voltage regulators, power management circuits, and the like, which are well known in the art, and therefore, will not be described any further.

The processor may be responsible for managing the bus and general processing, including the execution of software stored on the machine-readable media. The processor may be implemented with one or more general-purpose and/or special-purpose processors. Examples include microprocessors, microcontrollers, DSP processors, and other circuitry that can execute software. Software shall be construed broadly to mean instructions, data, or any combination thereof, whether referred to as software, firmware, middleware, microcode, hardware description language, or otherwise. Machine-readable media may include, by way of example, random access memory (RAM) , flash memory, read only memory (ROM) , programmable read-only memory (PROM) , erasable programmable read-only memory (EPROM) , electrically erasable programmable Read-only memory (EEPROM) , registers, magnetic disks, optical disks, hard drives, or any other suitable storage medium, or any combination thereof. The machine-readable media may be embodied in a computer-program product. The computer-program product may comprise packaging materials.

In a hardware implementation, the machine-readable media may be part of the processing system separate from the processor. However, as those skilled in the art will readily appreciate, the machine-readable media, or any portion thereof, may be external to the processing system. By way of example, the machine-readable media may include a transmission line, a carrier wave modulated by data, and/or a computer product separate from the device, all which may be accessed by the processor through the bus interface. Alternatively, or in addition, the machine-readable media, or any portion thereof, may be integrated into the processor, such as the case may be with cache and/or general register files. Although the various components discussed may be described as having a specific location, such as a local component, they may also be configured in various ways, such as certain components being configured as part of a distributed computing system.

The processing system may be configured as a general-purpose processing system with one or more microprocessors providing the processor functionality and external memory providing at least a portion of the machine-readable media, all linked together with other supporting circuitry through an external bus architecture. Alternatively, the processing system may comprise one or more neuromorphic processors for implementing the neuron models and models of neural systems described. As another alternative, the processing system may be implemented with an application specific integrated circuit (ASIC) with the processor, the bus interface, the user interface, supporting circuitry, and at least a portion of the machine-readable media integrated into a single chip, or with one or more field programmable gate arrays (FPGAs) , programmable logic devices (PLDs) , controllers, state machines, gated logic, discrete hardware components, or any other suitable circuitry, or any combination of circuits that can perform the various functionality described throughout this disclosure. Those skilled in the art will recognize how best to implement the described functionality for the processing system depending on the particular application and the overall design constraints imposed on the overall system.

The machine-readable media may comprise a number of software modules. The software modules include instructions that, when executed by the processor, cause the processing system to perform various functions. The software modules may include a transmission module and a receiving module. Each software module may reside in a single storage device or be distributed across multiple storage devices. By way of example, a software module may be loaded into RAM from a hard drive when a triggering event occurs. During execution of the software module, the processor may load some of the instructions into cache to increase access speed. One or more cache lines may then be loaded into a general register file for execution by the processor. When referring to the functionality of a software module below, it will be understood that such functionality is implemented by the processor when executing instructions from that software module. Furthermore, it should be appreciated that aspects of the present disclosure result in improvements to the functioning of the processor, computer, machine, or other system implementing such aspects.

If implemented in software, the functions may be stored or transmitted over as one or more instructions or code on a computer-readable medium. Computer-readable media include both computer storage media and communication media including any medium that facilitates transfer of a computer program from one place to another. A storage medium may be any available medium that can be accessed by a computer. By way of example, and not limitation, such computer-readable media can comprise RAM, ROM, EEPROM, CD-ROM or other optical disk storage, magnetic disk storage or other magnetic storage devices, or any other medium that can be used to carry or store desired program code in the form of instructions or data structures and that can be accessed by a computer. Additionally, any connection is properly termed a computer-readable medium. For example, if the software is transmitted from a website, server, or other remote source using a coaxial cable, fiber optic cable, twisted pair, digital subscriber line (DSL) , or wireless technologies such as infrared (IR) , radio, and microwave, then the coaxial cable, fiber optic cable, twisted pair, DSL, or wireless technologies such as infrared, radio, and microwave are included in the definition of medium. Disk and disc, as used, include compact disc (CD) , laser disc, optical disc, digital versatile disc (DVD) , floppy disk, and Blu-disc where disks usually reproduce data magnetically, while discs reproduce data optically with lasers. Thus, in some aspects, computer-readable media may comprise non-transitory computer-readable media (e.g., tangible media) . In addition, for other aspects computer-readable media may comprise transitory computer-readable media (e.g., a signal) . Combinations of the above should also be included within the scope of computer-readable media.

Thus, certain aspects may comprise a computer program product for performing the operations presented. For example, such a computer program product may comprise a computer-readable medium having instructions stored (and/or encoded) thereon, the instructions being executable by one or more processors to perform the operations described. For certain aspects, the computer program product may include packaging material.

Further, it should be appreciated that modules and/or other appropriate means for performing the methods and techniques described can be downloaded and/or otherwise obtained by a user terminal and/or base station as applicable. For example, such a device can be coupled to a server to facilitate the transfer of means for performing the methods described. Alternatively, various methods described can be provided via storage means (e.g., RAM, ROM, a physical storage medium such as a compact disc (CD) or floppy disk, etc. ) , such that a user terminal and/or base station can obtain the various methods upon coupling or providing the storage means to the device. Moreover, any other suitable technique for providing the methods and techniques described to a device can be utilized.

It is to be understood that the claims are not limited to the precise configuration and components illustrated above. Various modifications, changes, and variations may be made in the arrangement, operation, and details of the methods and apparatus described above without departing from the scope of the claims.

Claims

An apparatus, comprising:

at least one memory; and

at least one processor coupled to the at least one memory, the at least one processor configured:

to receive a machine learning (ML) model including multiple expert models;

to adapt the ML model to include an expert placeholder;

to dynamically select, by the adapted ML model, one or more expert models based on an input, the selected one or more expert models being included in the expert placeholder; and

to process, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output.
The apparatus of claim 1, in which the dynamically selected one or more expert models indicate expert model weights that are applied to the input to compute an output.
The apparatus of claim 1, in which the multiple expert models are trained on different subsets of a common training dataset.
The apparatus of claim 1, in which the at least one processor is further configured to select the one or more expert models using a selection function based on the input at inference time.
The apparatus of claim 1, in which the at least one processor comprises one or more neural processing units.
The apparatus of claim 1, in which the ML model comprises a transformer-based model.
A processor-implemented method performed by at least one processor, comprising:

receiving a machine learning (ML) model including multiple expert models;

adapting the ML model to include an expert placeholder;

dynamically selecting, by the adapted ML model, one or more expert models based on an input, the selected one or more expert models being included in the expert placeholder; and

processing, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output.
The processor-implemented method of claim 7, in which the dynamically selected one or more expert models indicate expert model weights that are applied to the input to compute an output.
The processor-implemented method of claim 7, in which the multiple expert models are trained on different subsets of a common training dataset.
The processor-implemented method of claim 7, in which the one or more expert models are selected using a selection function based on the input at inference time.
The processor-implemented method of claim 7, in which the processing is performed by one or more neural processing units.
The processor-implemented method of claim 7, in which the ML model comprises a transformer-based model.
An apparatus, comprising:

means for receiving a machine learning (ML) model including multiple expert models;

means for adapting the ML model to include an expert placeholder;

means for dynamically selecting, by the adapted ML model, one or more expert models based on an input, the selected one or more expert models being included in the expert placeholder; and

means for processing, by the adapted ML model, the input based on the dynamically selected one or more expert models to generate an output.
The apparatus of claim 13, in which the dynamically selected one or more expert models indicate expert model weights that are applied to the input to compute an output.
The apparatus of claim 13, in which the multiple expert models are trained on different subsets of a common training dataset.
The apparatus of claim 13, further comprising means for selecting the one or more expert models using a selection function based on the input at inference time.
The apparatus of claim 13, in which the processing is performed by one or more neural processing units.
The apparatus of claim 13, in which the ML model comprises a transformer-based model.