Neural network with regularized attention layer

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for processing network inputs using a neural network that includes one or more regularized attention layers. In one aspect, a method includes receiving a layer input to a regularized attention layer, where the layer input to the regularized attention layer includes a set of input inserts; and applying a regularized attention operation to the input embedded set to generate an output embedded set, including: transforming intermediate attention scores using a set of shaping constants to generate a transformed set of attention scores, in which: values of the shaping constants are initialized prior to training of the neural network, and values of the shaping constants are initialized prior to training of the neural network; the neural network is not adjusted during the training period of the neural network; and the values of the shaping constants are selected to regularize the output embedded set.

Description

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS

[0002] This application claims priority to U.S. Provisional Application No. 63 / 411,007, filed on September 28, 2022. The disclosure of the prior application is considered part of the disclosure of the present application and is incorporated by reference into the disclosure of the present application. Background Art

[0003] This specification relates to processing data using machine learning models.

[0004] A machine learning model receives an input and generates an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models and generate an output based on the received input and the values ​​of the model's parameters.

[0005] Some machine learning models are deep models that use multiple layers of models to generate outputs for received inputs. For example, a deep neural network is a deep machine learning model that includes an output layer and one or more hidden layers, each of which applies a nonlinear transformation to the received input to generate an output. Summary of the invention

[0006] This specification describes a system implemented as a computer program on one or more computers in one or more locations that processes an input using a neural network including one or more regularized attention layers to generate a network output.

[0007] Throughout this specification, "embedding" may refer to an ordered collection of values, such as a vector, matrix, or other numerical tensor.

[0008] Throughout the specification, the "rank" of a matrix may refer to the dimensionality of the vector space generated by the columns (or rows) of the matrix. For example, a matrix with a rank equal to 1 (i.e., a "rank-1" matrix) is a matrix having columns (or rows) that generate a one-dimensional vector space, e.g., such that each column (or row) is a scalar multiple of every other column (or row). Similarly, the rank of an embedded set may refer to the dimensionality of the vector space generated by the embedded set.

[0009] Throughout this specification, "regularizing" data generated by a neural network (e.g., data generated by one or more layers of a neural network) may refer to modifying the data so as to reduce or prevent numerical problems during training or inference. For example, regularizing data generated by a neural network may include one or more of: (i) modifying the data to maintain or reduce the norm of the data (or portion of the data), (ii) modifying the data to maintain or increase the norm of the gradient of an objective function used for training of the neural network, or (iii) modifying the data to maintain or increase the rank of the data.

[0010] According to one aspect, a method performed by one or more computers is provided, the method comprising: receiving a network input; processing the network input using a neural network comprising a plurality of neural network layers arranged as a directed graph to generate a network output for the network input, wherein the plurality of neural network layers comprises one or more regularized attention layers, and wherein, for each regularized attention layer, processing the network input comprises: receiving a layer input to the regularized attention layer, wherein the layer input to the regularized attention layer comprises a set of input embeddings; and applying a regularized attention operation to the set of input embeddings to generate a set of output embeddings, comprising: processing the set of input embeddings according to values ​​of a set of regularized attention layer parameters to generate: (i) a set of value embeddings comprising a corresponding value embedding for each input embedding, and (ii) a set of intermediate attention scores; using a shaping constant (shaping) to generate the set of input embeddings. constant) to generate a set of transformed attention scores, wherein: the values ​​of the shaping constants are initialized prior to training of the neural network and are not adjusted during training of the neural network; and the values ​​of the shaping constants are selected to regularize the set of output embeddings; and a set of output embeddings is generated using (i) the set of value embeddings and (ii) the set of transformed attention scores; and a layer output is provided for the attention layer based on the set of output embeddings.

[0011] In some embodiments, the value of the shaping constant is selected to maintain or increase the rank of the set of output embeddings.

[0012] In some embodiments, the value of the shaping constant is selected to increase the likelihood that the rank of the set of output embeddings exceeds a threshold.

[0013] In some embodiments, the values ​​of the set of shaping constants are derived from a shaping matrix by the following operations, which include: determining a decomposition of the shaping matrix into the product of: (i) a diagonal matrix and (ii) a block matrix, wherein the block matrix has row sums equal to 1; and applying logarithms to the block matrix, the shaping matrix being based on the result of applying the logarithms to the block matrix.

[0014] In some embodiments, the shaped matrix is ​​derived from at least one base matrix, wherein the values ​​of off-diagonal items of the base matrix decay exponentially based on the distance from the diagonal of the base matrix.

[0015] In some implementations, entries on each diagonal of the basis matrix have the same value.

[0016] In some implementations, each diagonal entry of the basis matrix has a value of 1.

[0017] In some embodiments, the shaped matrix is ​​derived from at least one base matrix, wherein diagonal entries of the base matrix each have a same first value and non-diagonal entries of the base matrix each have a same second value.

[0018] In some embodiments, the first value is 1 and the second value is strictly less than 1.

[0019] In some embodiments, transforming the intermediate attention scores using a set of shaping constants to generate a set of transformed attention scores includes, for each intermediate attention score: generating a corresponding transformed attention score by combining the intermediate attention score with the corresponding shaping constant.

[0020] In some embodiments, for each intermediate attention score, generating a corresponding transformed attention score includes: adding the intermediate attention score to the corresponding shaping constant.

[0021] In some embodiments, processing a set of input embeddings to generate a set of intermediate attention scores includes: processing the set of input embeddings to generate, for each input embedding: (i) a corresponding query embedding, and (ii) a corresponding key embedding; and generating each intermediate attention score based on a measure of similarity between the corresponding query embedding and the corresponding key embedding.

[0022] In some embodiments, generating a set of output embeddings using (i) a set of value embeddings and (ii) a set of transformed attention scores includes: generating a set of final attention scores by applying a causal masking operation to the set of transformed attention scores and then performing a nonlinear transformation; and generating a set of output embeddings using (i) a set of value embeddings and (ii) a set of transformed attention scores.

[0023] In some implementations, the non-linear transformation is a soft-max transformation.

[0024] In some embodiments, generating a set of output embeddings using (i) a set of value embeddings and (ii) a set of transformed attention scores includes: generating each output embedding based on a linear combination of the set of value embeddings, wherein coefficients of the linear combination are defined by corresponding transformed attention scores from the set of transformed attention scores.

[0025] In some implementations, generating the set of output embeddings includes applying an embedding-specific rescaling to the set of output embeddings based on a diagonal matrix.

[0026] In some implementations, prior to training of the neural network, the values ​​of the regularized attention layer parameters are initialized so that the value of each of the intermediate attention scores is zero.

[0027] In some embodiments, prior to training of the neural network, the values ​​of the regularized attention layer parameters are initialized to encourage the value of each of the intermediate attention scores to be close to zero.

[0028] In some embodiments, prior to training of the neural network, the values ​​of the regularized attention layer parameters selected are: (i) initialized to zero, or (ii) selected from a tolerance range around zero, or (iii) sampled from a probability distribution with a standard deviation selected from a tolerance range around zero.

[0029] In some implementations, the neural network includes a plurality of regularized attention layers associated with the ranking, wherein each of the plurality of regularized attention layers is associated with a different set of shaping constants.

[0030] In some embodiments, the neural network is configured to autoregressively generate the output sequence.

[0031] In some implementations, the neural network does not include skip connections associated with a regularized attention layer in the neural network.

[0032] In some implementations, the neural network does not include a normalization layer associated with a regularized attention layer in the neural network.

[0033] In some embodiments, regularizing the set of output embeddings includes one or more of: (i) modifying the set of output embeddings to maintain or reduce the norm of the output embeddings included in the set of output embeddings, (ii) modifying the set of output embeddings to maintain or increase the norm of the gradient of an objective function used for training the neural network, or (iii) modifying the set of output embeddings to maintain or increase the rank of the set of output embeddings.

[0034] According to another aspect, a system is provided, comprising: one or more computers; and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform the operations of the methods described herein.

[0035] One or more non-transitory computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of the methods described herein.

[0036] Particular embodiments of the subject matter described in this specification can be implemented so as to realize one or more of the following advantages.

[0037] Many deep neural networks that include attention layers are difficult to train quickly and do not generalize well to unseen data after training. For example, many tasks require neural networks to have very specific architectural elements, such as skip connections, normalization layers, etc., in order to be trained to perform the task well. However, the need to include these specific architectural elements can mask hidden problems in the neural network architecture, can make it difficult to design new neural network architectures, and more generally, can make it difficult to train neural networks that do not have these elements but may show improved performance on these tasks. The present specification describes techniques for modifying the attention layers of neural networks to eliminate the need for these elements and allow neural networks to be trained efficiently and quickly even without including these elements. In addition, applying the described techniques results in the neural network generalizing better to unseen data after training, resulting in improved inference performance. As a specific example, a neural network that cannot otherwise be trained in a reasonable amount of time or a reasonable amount of computation due to the lack of one or more specific elements, such as skip connections or batch normalization, can instead be trained to exceed the performance of a traditional neural network with the specific elements (if the trained neural network includes the described regularized attention layer).

[0038] During training, neural networks that include attention layers may suffer from regularization problems such as "rank collapse," a situation in which the rank of a set of embeddings operated by the attention layer is reduced to a small value, e.g., 1, 2, or 3. For example, rank collapse may occur if each embedding in the set of embeddings operated by the attention layer becomes aligned (approximately or exactly) in the same direction. Rank collapse can significantly hinder the training of a neural network, e.g., by zeroing the gradients of certain parameters in the attention layer. Traditionally, as described above, regularization problems such as rank collapse are addressed using architectural elements such as skip connections and normalization layers. The present specification describes techniques for modifying the attention operations performed by an attention layer in a manner that can reduce the likelihood of regularization problems without requiring the use of skip connections and normalization layers.

[0039] The techniques described in this specification can reduce the use of computational resources (e.g., memory and computing power) by a neural network by eliminating the need to include architectural elements such as skip connections and normalization layers in the neural network. For example, implementing skip connections that "jump" blocks in a neural network may require storing the inputs of the blocks while generating the outputs of the blocks, e.g., so that the block inputs can be combined (e.g., added) with the block outputs. Thus, removing skip connections from a neural network reduces the memory footprint of the neural network by reducing temporary storage of intermediate outputs. Implementing a normalization layer may require performing computationally intensive operations by aggregating data to generate normalization constants, and removing the normalization layer thus eliminates part of the computational footprint of the neural network. By reducing the use of computational resources (as described above), the techniques described in this specification can enable certain neural networks to be implemented on a single target device, rather than in a distributed manner, e.g., across multiple devices.

[0040] Specifically, neural networks with attention layers (e.g., large-scale transformer neural networks) are often deployed on hardware accelerators (e.g., graphics processing units, GPUs) with limited on-chip memory. However, as described above, such neural networks often have significant memory requirements that may exceed the memory available on the hardware accelerator, for example, because of the inclusion of skip connections and normalization layers. The neural networks described in this specification may have lower memory requirements, for example, by eliminating the need for skip connections and normalization layers using regularized attention operations, and may therefore be more easily implemented on hardware accelerators. In addition, the neural networks described in this specification may reduce requirements for memory bandwidth, for example, writing / reading data to / from disk during training and inference, thereby further reducing consumption of computing resources.

[0041] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the following description. Other features, aspects, and advantages of the subject matter will become apparent from the description, drawings, and claims. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] Figure 1 An example neural network system is shown.

[0043] Figure 2 is a flowchart of an example process for processing a layer input through a regularized attention layer to generate a layer output.

[0044] Figure 3 is a flow chart of an example process for generating a set of shaping constants for regulating a regularized attention layer of a neural network.

[0045] Figure 4is a flowchart of an example process for training a neural network that includes one or more regularized attention layers.

[0046] Figure 5 Shown is the evolution of the normalized kernel matrix of an attention layer in a neural network.

[0047] Figure 6 The training loss is shown for various possible neural network architectures over a range of training steps.

[0048] Figure 7 Shown is the sensitivity of the training performance of a neural network with a regularized attention layer to the initial values ​​of certain parameters of the regularized attention layer.

[0049] The same reference numbers and names in different drawings indicate the same elements. DETAILED DESCRIPTION

[0050] Figure 1 An example neural network system 100 is shown. Neural network system 100 is an example of a system implemented as a computer program on one or more computers in one or more locations in which the systems, components, and techniques described below are implemented.

[0051] The neural network system 100 includes a neural network 104 configured to process a network input 102 according to the values ​​of a set of neural network parameters of the neural network 104 to generate a corresponding network output. The neural network 104 includes one or more regularized attention layers 108 that perform an attention operation conditioned on a set of "shaping constants." The values ​​of the shaping constants are selected to regularize the network outputs of the regularized attention layers, e.g., to reduce the likelihood of regularization problems, such as rank collapse among the set of output embeddings generated by the regularized attention layers, as will be described in more detail below.

[0052] Neural network 104 may be configured to perform any suitable neural network tasks, and in particular, may be configured to process any suitable network input to generate any suitable network output. Several examples of neural network tasks that may be performed by neural network 104 are described next.

[0053] In some embodiments, the neural network can be configured to receive an input image and process the input image to generate a network output of the input image, i.e., perform an image processing task. In this specification, processing the input image refers to processing the intensity values ​​of the pixels of the image using the neural network. The image can be, for example, an image captured by a camera, a point cloud image captured by a lidar or other laser sensor, a hyperspectral image, a medical image captured by a medical imaging device, or any other suitable type of data that can be represented in an image format.

[0054] For example, the task may be image classification, and the output generated by the neural network for a given image may be a score for each of a set of object categories, each score representing an estimated likelihood that the image contains an object belonging to that category.

[0055] As another example, the task may be image embedding generation, and the output generated by the neural network may be a numerical embedding of the input image.

[0056] As yet another example, the task may be object detection, and the output generated by the neural network may identify locations in the input image, e.g., bounding boxes or other geometric regions within the image, where particular types of objects are depicted.

[0057] As yet another example, the task may be image segmentation, and the output generated by the neural network may define, for each pixel of the input image, to which of a number of classes the pixel belongs.

[0058] In some embodiments, the input to the neural network is an Internet resource (e.g., a web page), a document or a document portion, or features extracted from an Internet resource, document or a document portion, and the neural network performs the task of classifying the resource or document. For example, the output generated by the neural network for a given Internet resource, document or document portion may be a score for each of a set of topics, each score representing an estimated likelihood that the Internet resource, document or document portion is about the topic.

[0059] In some implementations, the inputs to the neural network are features of the impression context of a particular advertisement, and the output generated by the neural network may be a score representing an estimated likelihood that the particular advertisement will be clicked.

[0060] In some implementations, the input to the neural network is features of personalized recommendations for the user, e.g., features describing the context of the recommendation, e.g., features describing actions previously taken by the user, and the output generated by the neural network can be a score for each of a set of content items, each score representing an estimated likelihood that the user will respond favorably to the recommended content item.

[0061] In some implementations, the input to the neural network is a sequence of text in one language, and the output generated by the neural network may be a score for each of a set of text segments in another language, each score representing an estimated likelihood that the text segment in the other language is a correct translation of the input text into the other language.

[0062] In some implementations, the neural network may perform an audio processing task. For example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network may be a score for each of a set of text segments, each score representing an estimated likelihood that the text segment is a correct transcript of the utterance. As another example, the task may be a keyword spotting task, where, where the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network may indicate whether a particular word or phrase (a "hot word") was spoken in the utterance. As another example, the input to the neural network may be a sequence representing a spoken utterance, and the output generated by the neural network may identify the natural language in which the utterance was spoken.

[0063] In some embodiments, the neural network can perform natural language processing or understanding tasks, such as entailment tasks, paraphrase tasks, text similarity tasks, sentiment tasks, sentence completion tasks, grammar tasks, etc., which operate on text sequences of some natural language.

[0064] In some implementations, a neural network can perform a text-to-speech task where the input is text in a natural language or features of text in a natural language, and the network output is a spectrogram or other data defining the audio of the text spoken in the natural language.

[0065] In some embodiments, a neural network can perform a genomics task where the input is a sequence representing a fragment of a DNA sequence or other molecular sequence, and the output is an embedding of the fragment for use in a downstream task, e.g., by utilizing unsupervised learning techniques on a dataset of DNA sequence fragments, or is the output of a downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting the functional effects of non-coding variants, and the like.

[0066] In some embodiments, the neural network can be configured to generate network outputs autoregressively. That is, the neural network can be configured to generate network outputs that include a corresponding output element for each location in the output location sequence. The neural network can generate one output element at a time based on the order of the locations in the output sequence. For each location, the neural network can generate an output element at that location based on the output element at the previous location in the output sequence, but not based on the output element at the future location in the output sequence (i.e., the output element has not yet been generated). For example, the neural network can autoregressively generate pixel values ​​in an image, or audio samples in an audio waveform, or characters in a text sequence, etc.

[0067] In some embodiments, the neural network can perform a combination of multiple separate machine learning tasks, i.e., the neural network is configured to perform multiple different separate machine learning tasks, such as two or more of the machine learning tasks described above. For example, the neural network can be configured to perform multiple separate natural language understanding tasks. Optionally, the network input can include identifiers of separate tasks to be performed on the network input. As another example, the neural network can be configured to perform multiple separate image processing or computer vision tasks, i.e., to generate outputs of multiple different separate image processing tasks in parallel by processing a single input image.

[0068] In some implementations, a neural network can perform an agent control task where the input is an observation describing a state of an environment and the output defines an action for the agent to perform in response to the observation. For example, the output can define a score distribution over a set of possible actions that the agent can perform, and the score distribution can be used to select an action for the agent to perform.

[0069] In some embodiments, the environment is a real-world environment, the agent is a mechanical agent that interacts with the real-world environment, such as a robot or an autonomous or semi-autonomous land, air, or sea vehicle that operates or navigates in the environment, and the action is an action taken by the mechanical agent to perform a task in the real-world environment. For example, the agent can be a robot that interacts with the environment to complete a specific task, such as locating an object of interest in the environment or moving an object of interest to a specified location in the environment or navigating to a specified destination in the environment.

[0070] In these embodiments, observations may include, for example, one or more of images, object positioning data, and sensor data to capture observations as the agent interacts with the environment, for example, from images, distance or positioning sensors, or sensor data from actuators. For example, in the case of a robot, observations may include data describing the current state of the robot, for example, one or more of the following: joint positioning, joint velocity, joint force, torque or acceleration, for example, gravity-compensated torque feedback, and the global or relative pose of an item held by the robot. In the case of a robot or other mechanical agent or vehicle, observations may similarly include one or more of the positioning, linear or angular velocity, force, torque or acceleration, and global or relative pose of one or more parts of the agent. Observations may be defined in 1, 2, or 3 dimensions and may be absolute and / or relative observations. Observations may also include, for example, sensed electronic signals, such as motor current or temperature signals; and / or image or video data, for example, from a camera or lidar sensor, for example, data from a sensor of the agent or data from a sensor located separately from the agent in the environment.

[0071] In these embodiments, the action can be a control signal to control a robot or other mechanical agent, such as a torque or higher-level control command to a robot joint, or an autonomous or semi-autonomous land, air, or sea vehicle, such as a torque to a control surface or other control element (e.g., a steering control element of a vehicle), or a higher-level control command. The control signal can include, for example, positioning, velocity, or force / torque / acceleration data of one or more joints of a robot or part of another mechanical agent. The control signal can also or alternatively include electronic control data, such as motor control data, or more generally, data for controlling one or more electronic devices within an environment, the control of which has an effect on the observed state of the environment. For example, in the case of an autonomous or semi-autonomous land, air, or sea vehicle, the control signal can define actions to control navigation (e.g., steering) and motion (e.g., braking and / or acceleration of the vehicle).

[0072] In some embodiments, the environment is a simulation of the real-world environment described above, and the agent is implemented as one or more computers that interact with the simulated environment. For example, the simulated environment can be a simulation of a robot or a vehicle, and the reinforcement learning system can be trained on the simulation and then, once trained, used in the real world.

[0073] In some embodiments, the environment is a real-world manufacturing environment for manufacturing products such as chemical, biological or mechanical products or food. As used herein, "manufacturing" products also include refining starting materials to produce products, or processing starting materials such as removing pollutants to generate clean or recycled products. A manufacturing plant may include multiple manufacturing units, such as containers for chemical or biological substances, or machines for processing solids or other materials, for example, robots. Manufacturing units are configured so that intermediate versions or components of a product can be moved between manufacturing units during product manufacturing, for example, via pipelines or mechanical transport. As used herein, the manufacture of a product also includes manufacturing food by a kitchen robot.

[0074] The agent may include an electronic agent configured to control a manufacturing unit or a machine, such as a robot, that operates to manufacture a product. That is, the agent may include a control system configured to control the manufacture of a chemical, biological, or mechanical product. For example, the control system may be configured to control one or more manufacturing units or machines, or to control the movement of intermediate versions or components of a product between manufacturing units or machines.

[0075] As one example, the tasks performed by the agent may include tasks for manufacturing a product or an intermediate version or component thereof. As another example, the tasks performed by the agent may include tasks for controlling (e.g., minimizing) the use of resources, such as tasks for controlling the consumption of electricity or water or any materials or consumables used in the manufacturing process.

[0076] Actions may include control actions to control the use of a machine or manufacturing cell for processing solid or liquid materials to manufacture a product or an intermediate product or component thereof, or to control the movement of an intermediate version or component of a product within a manufacturing environment, for example, between manufacturing cells or machines. In general, these actions may be any actions that have an effect on the observed state of the environment, for example, actions configured to adjust any of the sensed parameters described below. These may include actions that adjust the physical or chemical conditions of a manufacturing cell, or actions that control the movement of a mechanical part of a machine or a robotic joint. These actions may include actions that impose operating conditions on a manufacturing cell or machine, or actions that result in changes to settings to regulate, control, or turn on or off the operation of a manufacturing cell or machine.

[0077] The reward or reward may be related to a performance measure of the task. For example, in the case of a task to manufacture a product, the measure may include a measure of the number of products manufactured, the quality of the product, the speed at which the product is produced, or a measure of the physical cost of performing the manufacturing task, such as a measure of the amount of energy, material, or other resource used to perform the task. In the case of a task to control resource usage, the measure may include any measure of resource usage.

[0078] In general, observations of environmental states may include any electronic signals that represent the function of an electronic and / or mechanical item of equipment. For example, representations of environmental states may be derived from observations made by sensors that sense the state of a manufacturing environment, such as sensors that sense the state or configuration of a manufacturing cell or machine, or sensors that sense the movement of materials between manufacturing cells or machines. As some examples, such sensors may be configured to sense mechanical motion or force, pressure, temperature; electrical conditions such as current, voltage, frequency, impedance; quantity, level, flow / movement rate, or flow / movement path of one or more materials; physical or chemical conditions, such as physical state, shape, or configuration, or chemical state, such as pH; configuration of a cell or machine, such as a mechanical configuration of a cell or machine, or valve configuration; image or video sensors for capturing images or video observations of a manufacturing cell or machine or movement; or any other suitable type of sensor. In the case of a machine, such as a robot, observations from sensors may include observations of positioning, linear or angular velocity, force, torque, or acceleration, or posture of one or more parts of the machine, such as data describing the current state of the machine or robot or an item held or handled by the machine or robot. Observations may also include, for example, sensed electronic signals, such as motor current or temperature signals, or image or video data, such as from a camera or lidar sensor. Sensors such as these may be part of an agent in the environment or located separately from the agent.

[0079] In some embodiments, the environment is a real-world environment of a service facility including multiple electronic equipment items, such as a server farm or a data center, such as a telecommunications data center, or a computer data center for storing or processing data, or any service facility. The service facility may also include auxiliary control equipment for controlling the operating environment of the equipment item, such as environmental control equipment, such as temperature control, such as cooling equipment, or air flow control or air conditioning equipment. The task may include a task to control (e.g., minimize) resource usage, such as a task to control power consumption or water consumption. The agent may include an electronic agent configured to control the operation of the equipment item, or to control the operation of auxiliary, such as environmental control equipment.

[0080] In general, these actions may be any action that has an effect on the observed environmental state, such as actions configured to adjust any of the sensed parameters described below. These may include actions that control or impose operating conditions on items of equipment or auxiliary control equipment, such as actions that result in changing settings to adjust, control, or turn on or off the operation of equipment or auxiliary control equipment.

[0081] Typically, observations of environmental conditions may include any electronic signal that represents the function of a facility or equipment in the facility. For example, representations of environmental conditions may be derived from observations made by any sensor that senses the physical environmental conditions of the facility or any sensor that senses the state of one or more equipment items or one or more auxiliary control equipment items. These include sensors configured to sense electrical conditions such as current, voltage, power, or energy; the temperature of the facility; the flow, temperature, or pressure of a fluid within the facility or within the facility's cooling system; or physical facility configurations, such as whether a vent is open.

[0082] The reward or reward may be related to a performance metric of the task. For example, in the case of a task to control (eg, minimize) resource usage, such as a task to control the usage of electricity or water, the metric may include any metric of resource usage.

[0083] In some embodiments, the environment is a real-world environment of a power generation facility, for example, a renewable power generation facility, such as a solar farm or a wind farm. The task may include a control task for controlling the power generated by the facility, for example, controlling the delivery of power to a distribution network, for example, to meet demand or reduce the risk of mismatch between grid elements, or to maximize the power generated by the facility. The agent may include an electronic agent configured to control the power generation of the facility or to couple the generated power into the grid. These actions may include actions to control the electrical or mechanical configuration of the generator, such as the electrical or mechanical configuration of one or more renewable power generation elements, for example, controlling the configuration of a wind turbine or solar panel or panel or mirror, or the electrical or mechanical configuration of a rotating generator. Mechanical control actions may, for example, include actions to control the conversion of energy input to electrical energy output, for example, the conversion efficiency or the degree of coupling of energy input to electrical energy output. Electrical control actions may, for example, include actions to control one or more of the voltage, current, frequency, or phase of the generated power.

[0084] The reward or reward may be related to a performance metric of the task. For example, in the case of a task to control the delivery of power to a distribution network, the metric may involve a measure of the power delivered, or a measure of an electrical mismatch (such as a voltage, current, frequency, or phase mismatch) between a generating facility and the grid, or a measure of power or energy lost in the generating facility. In the case of a task to maximize the power delivered to the distribution network, the metric may involve a measure of power or energy delivered to the grid, or a measure of power or energy lost in the generating facility.

[0085] Typically, observations of environmental conditions may include any electronic signal representing the electrical or mechanical function of power generation equipment in a power generation facility. For example, representations of environmental conditions may be derived from observations made by any sensor that senses the physical or electrical state of equipment that is generating electricity in a power generation facility, or the physical environment of such equipment, or the condition of auxiliary equipment that supports the power generation equipment. Such sensors may include sensors configured to sense electrical conditions of equipment (such as current, voltage, power, or energy); temperature or cooling of the physical environment; fluid flow; or the physical configuration of the equipment; and observations of electrical conditions of the power grid, such as from local or remote sensors. Observations of environmental conditions may also include one or more predictions about future operating conditions of power generation equipment, such as predictions of future wind levels or solar irradiance or predictions of future electrical conditions of the power grid.

[0086] As another example, the environment can be a chemical synthesis or protein folding environment, such that each state is a corresponding state of a protein chain or one or more intermediates or precursor chemicals, and the agent is a computer system for determining how to fold the protein chain or synthesize the chemical. In this example, these actions are possible folding actions for folding the protein chain or actions for assembling precursor chemicals / intermediates, and the results to be achieved can include, for example, folding the protein so that the protein is stable, and so that it realizes a specific biological function, or providing an efficient synthesis pathway for the chemical. As another example, the agent can be a mechanical agent that performs or controls protein folding actions or chemical synthesis steps automatically selected by the system without human interaction. Observations can include direct or indirect observations of the states of proteins or chemicals / intermediates / precursors, and / or can be derived from simulations.

[0087] In a similar manner, the environment can be a drug design environment, such that each state is a corresponding state of a potential pharmaceutical drug, and the agent is a computer system for determining pharmaceutical drug elements and / or pharmaceutical drug synthesis pathways. For example, in a simulation, the drug / synthesis can be designed based on a reward derived from the drug's goal. As another example, the agent can be a mechanical agent that performs or controls the synthesis of the drug.

[0088] In some further applications, the environment is a real-world environment, and the agent manages the distribution of tasks across computing resources, for example on a mobile device and / or in a data center. In these embodiments, an action may include allocating a task to a specific computing resource.

[0089] As a further example, an action may include presenting an advertisement, an observation may include an ad impression or click-through count or rate, and a reward may describe a previous selection of an item or content by one or more users.

[0090] In some cases, the observations may include text or verbal instructions provided to the agent by a third party (e.g., an operator of the agent). For example, the agent may be an autonomous vehicle, and the user of the autonomous vehicle may provide text or verbal instructions to the agent (e.g., to navigate to a particular location).

[0091] As another example, the environment may be an electrical, mechanical, or electromechanical design environment, e.g., an environment in which the design of an electrical, mechanical, or electromechanical entity is simulated. A simulation environment may be a simulation of a real-world environment in which the entity is intended to work. The task may be to design an entity. Observations may include observations describing the entity, i.e., observations of the mechanical shape or electrical, mechanical, or electromechanical configuration of the entity, or observations of parameters or properties of the entity. Actions may include actions to modify the entity, e.g., to modify one or more observations. A reward or reward may include a measure of one or more properties of the entity design. For example, the reward or reward may relate to one or more physical characteristics of the entity, such as weight or strength, or to one or more electrical characteristics of the entity, such as a measure of the efficiency of performing a specific function for which the entity is designed. The design process may include outputting a design for manufacture, e.g., in the form of computer executable instructions for manufacturing the entity. The process may include making the entity according to the design. Thus, the design of the entity may be optimized, e.g., by reinforcement learning, and then used to manufacture an optimized design output of the entity, e.g., as computer executable instructions; the entity with the optimized design may then be manufactured.

[0092] As previously described, the environment can be a simulated environment. Typically, in the case of a simulated environment, observations can include simulated versions of one or more previously described observations or observation types, and actions can include simulated versions of one or more previously described actions or action types. For example, the simulated environment can be a motion simulation environment, such as a driving simulation or a flight simulation, and the agent can be a simulated vehicle navigating in the motion simulation. In these embodiments, the action can be a control input to control a simulated user or a simulated vehicle. Typically, the agent can be implemented as one or more computers that interact with the simulated environment.

[0093] The simulation environment can be a simulation of a specific real-world environment and agent. For example, the system can be used to select actions of an agent in a simulation environment during training or evaluation of the system, and after training or evaluation or both are completed, it can be deployed to control a real-world agent in a specific real-world environment that is the object of simulation. This can avoid unnecessary wear and damage to the real-world environment or real-world agent, and can allow the control neural network to be trained and evaluated in situations that rarely occur in the real-world environment or are difficult or unsafe to reproduce. For example, the system can be partially trained using a simulation of a mechanical agent in a simulation of a specific real-world environment, and then the system can be deployed to control a real mechanical agent in a specific real-world environment. Therefore, in such a case, the observation of the simulation environment involves the real-world environment, and the selected action in the simulation environment is related to the action to be performed by the mechanical agent in the real-world environment.

[0094] Optionally, in any of the above embodiments, the observation at any given time step may include data from a previous time step that may be helpful in describing the environment, such as an action performed at a previous time step, a reward received at a previous time step, or both.

[0095] Neural network 104 may be implemented in any suitable location, for example, on a user device, on an agent (e.g., a mechanical agent, such as a robot or vehicle), or in a data center, or in a cloud environment. In some cases, neural network 104 may be implemented in a distributed manner across two or more locations. For example, a first subnet of neural network 104 may be implemented on a user device or an agent, while a second subnet of neural network 104 may be implemented in a data center or a cloud environment.

[0096] Neural network 104 may have any suitable neural network architecture that enables the neural network to perform its described functions, e.g., processing network inputs to generate corresponding network outputs as part of performing a neural network task. For example, neural network 104 may include any suitable type of neural network layers (e.g., fully connected layers, convolutional layers, attention layers, etc.) in any suitable number (e.g., 5, 50, or 100 layers) and connected in any suitable configuration (e.g., as a directed graph of layers).

[0097] The neural network 104 includes one or more regularized attention layers 108. In particular, the neural network 104 can include any suitable number of regularized attention layers 108, e.g., 1, 5, 10, or 100 regularized attention layers, and the regularized attention layers can optionally be interleaved with other neural network layers. Optionally, the architecture of the neural network 104 can exclude normalization layers, or skip connections, or both. The regularized attention layers of the neural network 104 can achieve stable training and inference performance of the neural network 104, even without normalization layers and skip connections, and even when the neural network is a deep neural network with many layers (e.g., 30, 50, or 100 layers).

[0098] Each regularized attention layer is configured to process a layer input 106 of a regularized attention layer 108 to generate a corresponding layer output 110 of the regularized attention layer. The layer input 106 may include a set of input embeddings, and the layer output 110 may include a set of output embeddings. In some cases, the number of input embeddings may be equal to the number of output embeddings. The input embeddings and output embeddings may have any suitable dimensions, and optionally, the input embeddings may have a different dimension than the output embeddings.

[0099] The neural network 104 may provide any suitable layer input 106 to the regularized attention layer, e.g., a layer input comprising at least a portion of a layer output of another neural network layer, or a layer input comprising at least a portion of the network input 102, or both. Similarly, the neural network 104 may provide a layer output 110 of the regularized attention layer 108, e.g., for processing by another layer in the neural network 104, or as a network output 112 of the neural network 104.

[0100] Signal propagation through a neural network can be evaluated (at least in part) with reference to the evolution of the "kernel matrix" through the layers of the neural network. The kernel matrix associated with layer l of a neural network can be defined as:

[0101]

[0102] in, Denotes the activation sequence of length T at layer l, i.e., the sequence generated when the neural network processes the network input to generate the corresponding network output. Regularization problems can arise in neural networks, for example, in cases where the diagonal entries of the kernel matrix grow or shrink rapidly with depth, which can indicate an uncontrolled activation norm and can lead to saturation losses or other numerical problems. Another form of regularization problem that can arise is rank collapse, for example, where the rank of the kernel matrix converges to a small value, such as one, two, or three. Rank collapse can lead to zero gradients for certain neural network parameters, thereby hindering the training of the neural network. Mitigating numerical problems (such as rank collapse) is critical to achieving efficient training of neural networks. Reference Figure 5 A diagram depicting the evolution of a kernel matrix through the layers of a neural network, which will be described in more detail below.

[0103] Each regularized attention layer 108 is conditioned on (e.g., parameterized by) a corresponding set of shaping constants. The shaping constants have values ​​that are initialized prior to training of the neural network 104, and optionally, are not adjusted during training of the neural network. The values ​​of the shaping constants of the regularized attention layer are selected to regularize the set of output embeddings generated by the regularized attention layer, e.g., to prevent rank collapse by increasing the likelihood that the rank of the set of output embeddings exceeds a threshold.

[0104] More specifically, each regularized attention layer 108 is configured to apply an attention operation to a set of input embeddings in the layer input 106 of the regularized attention layer 108 as part of generating a layer output 110 of the regularized attention layer. Applying the attention operation involves processing the layer inputs to generate a set of attention scores, which, when represented in matrix form, can be referred to as an attention matrix. The kernel matrix of a regularized attention layer can be described as a function of the product of attention matrices generated by previous regularized attention layers. Therefore, the evolution of the kernel matrix (which, as described above, describes signal propagation through the neural network) depends on the properties of the product of attention matrices generated by the regularized attention layer. One method for selecting a shaping constant that has the effect of regularizing the output of an attention layer is to identify a shaping constant that increases the likelihood that the product of the attention matrices generated by the attention layer has desired numerical properties (e.g., full rank and bounded norm). Referring below to Figure 3 An example process for generating such shaping constants is described in more detail.

[0105] Next, the following reference Figure 2 An example process in which the regularized attention layer 108 may apply a regularized attention operation to a layer input to generate a layer output is described in more detail.

[0106] Figure 2is a flow chart of an example process 200 for processing a layer input to generate a layer output by a regularized attention layer. For convenience, process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system appropriately programmed according to the present specification, e.g., Figure 1 The neural network system 100 can perform process 200.

[0107] The system receives a layer input (202) to a regularized attention layer. The layer input includes a matrix is a set of columns or rows of input embeddings T, where T is the number of input embeddings and d is the dimension of each input embedding.

[0108] The system performs the operations of steps 204-208 for each of one or more attention heads of the regularized attention layer. The regularized attention layer can include any suitable number of attention heads, for example, 1, 5, or 10 attention heads. Each attention head of the regularized attention layer can have a corresponding set of attention head parameters that are specific to the attention head, and in particular, different from each other attention head.

[0109] For each attention head, the system processes the set of input embeddings according to the values ​​of the set of parameters of the attention head to generate: (i) a set of value embeddings (including a corresponding value embedding corresponding to each input embedding), and (ii) a set of intermediate attention scores (204).

[0110] The system can generate a matrix with rows or columns representing sets of value embeddings For example, where T is the number of value embeddings and d h is the dimension of each value embedding, d h = d / h, d is the dimension of the input embedding, and h is the number of attention heads, as shown below:

[0111] V=X·W V (2)

[0112] where X is a matrix representing the set of input embeddings, and is the matrix of (trainable) parameter values ​​for the regularized attention layer.

[0113] The system can generate a matrix of intermediate attention scores As shown below:

[0114]

[0115] Q(X)=X·W Q (4)

[0116] K(X)=X·W K (5)

[0117] in, is the query embedding matrix, is the key embedding matrix, and W Q , W K is the matrix of (trainable) parameter values ​​for the regularized attention layer.

[0118] In some cases, for example, for an autoregressive neural network that generates a sequence of output tokens one at a time, some intermediate attention scores (e.g., those involving attention between previously produced tokens) can be pre-computed and stored, so that the neural network does not regenerate these attention scores, but instead retrieves them from a cache.

[0119] For each attention head, the system transforms the intermediate attention scores generated by the attention head using a set of shaping constants to generate a set of transformed attention scores (206). The values ​​of the shaping constants are initialized prior to training of the neural network, and optionally, are not adjusted during training of the neural network. The values ​​of the shaping constants are selected to regularize the set of output embeddings generated by the regularized attention layer. Each attention head of the regularized attention layer can optionally use the same shaping constants, i.e., the shaping constants have the same values ​​as those used by each other attention head of the regularized attention layer.

[0120] In some implementations, the system generates a set of transformed attention scores as follows

[0121] A′=A+B (6)

[0122] Where A is the matrix of intermediate attention scores and B is the matrix of shaping constants. Figure 3 An example process for generating a matrix B of shaping constants is described in more detail. The matrix B of shaping constants can be derived from a "basis matrix" with diagonal equal to 1 (as shown in FIG. Figure 3 ), which makes the shaping constant provide an implicit mechanism to control the norm of the output embeddings at all sequence positions in the deep layer. In addition, the matrix B of the shaping constants leads to a "recency bias", in which the output embeddings are close to the positions with large cosine similarity.

[0123] In some implementations, the system generates a set of transformed attention scores as follows

[0124] A′=αI+βA (7)

[0125] Where α, β are trainable (scalar) parameters, I represents the identity matrix (i.e., a square matrix with 1s on the diagonal and 0s off the diagonal), and A is the matrix of intermediate attention scores. In this example, the shaping constants are the binary entries of the identity matrix I. The shaping constants represented in Equation (7) can regularize the output embedding, for example by preventing rank collapse, but may not achieve all the advantages of the shaping constants represented in Equation (6). For example, at initialization, when the intermediate attention scores are zero (or close to zero) and the shaping constants (in such a case, defined as the scaled identity matrix) control the attention operation, there may be useful information contained in certain sequence positions that is not used.

[0126] Optionally, the system can further transform the intermediate attention scores by one or more additional modifications, i.e., in addition to transforming the intermediate attention scores using the shaping constant. For example, the system can bias each attention score by a penalty based on the distance between the pair of input embeddings (i.e., from the layer inputs) corresponding to the attention score.

[0127] For each attention head, the system generates a set of output embeddings of the attention head using: (i) the set of value embeddings generated by the attention head, and (ii) the transformed attention scores generated by the attention head (208).

[0128] In some embodiments, the system generates a set of output embeddings of the attention heads as follows

[0129] O=A′V (8)

[0130] Where A′ represents the matrix of transformed attention scores and V represents the matrix of value embeddings.

[0131] In some embodiments, the system generates a set of output embeddings of the attention heads as follows

[0132]

[0133] Among them, the matrix Implements an embedding-specific rescaling of the set of output embeddings. softmax(·) denotes the softmax operator. Is to satisfy M i,j = 1{i≥j}, otherwise it is zero, and Γ is a large constant, for example, 10 30, and V denotes the matrix of value embeddings. The mask matrix (and the associated large constant Γ) enforces causal masking of the attention scores (i.e., it applies a causal masking operation that treats the input embeddings as being in a sequence (with an order corresponding to the order of consecutive columns or rows of X) and prevents a given output embedding corresponding to the corresponding input embedding in the sequence from being generated based on the transformed attention scores of the input embeddings of a sequence that is later than the input embedding corresponding to the given output embedding), and may be beneficial, for example, if the neural network is configured to generate an output sequence autoregressively. Reference Figure 3 An example procedure for generating a matrix D that parameterizes an embedding-specific rescaling operation is described.

[0134] The system provides a layer output (210) for the regularized attention layer based on the set of corresponding output embeddings generated by each attention head. For example, the system can generate a layer output as follows

[0135] O′=concat(O1,…,O h )W O (10)

[0136] Among them, concat(·) is the concatenation operator, is the corresponding output embedding generated by each attention head, and W O A matrix representing the (trainable) parameter values ​​of the regularized attention layer.

[0137] Figure 3 is a flow chart of an example process 300 for generating a set of shaping constants for regulating a regularized attention layer of a neural network. Figure 3 The shaping constant generated by the process described in can regularize the output embedding generated by the regularized attention layer. For convenience, process 300 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system appropriately programmed according to the present specification, such as, Figure 1 The neural network system 100 can perform process 300.

[0138] The system generates one or more "basis" matrices (302). Some examples of possible basis matrices are described below.

[0139] In some implementations, the system generates one or more basis matrices in which the values ​​of off-diagonal entries of the basis matrix decay exponentially based on the distance from the diagonal of the basis matrix. The entries on the diagonal of the basis matrix can have the same value, such as a value of 1. For example, the system can generate a first basis matrix and the second basis matrix in:

[0140] (Σ in ) i,j =exp(-γ in |ij|) (11)

[0141] (Σ out ) i,j =exp(-γ out |ij|) (12)

[0142] Among them, γ in and γ out is the positive exponential decay rate, γ in ≥γ out In the case where the neural network includes a sequence of L>1 regularized attention layers, the system can instantiate a decreasing sequence of exponential decay rates Among them, γ l ≥γ l+1 , and where for each layer l, the γ of layer l in Can be selected as γ l-1 , and the γ of layer l out Can be selected as γ l .

[0143] In some implementations, the system generates one or more basis matrices, wherein the diagonal entries of each basis matrix each have a first value (e.g., a value of 1), and the off-diagonal entries of each basis matrix each have a second value (e.g., a value strictly less than 1). For example, the system may generate a first basis matrix and the second basis matrix in:

[0144] Σ in =(1-ρ in )I T +ρ in 11 T (13)

[0145] Σ out =(1-ρ out )I T +ρ out 11 T (14)

[0146] Among them, ρ in and ρ out is a positive value, ρ out ≥ρ in In the case where the neural network consists of a sequence of L>1 regularized attention layers, the system can instantiate an increasing sequence of values Among them, ρ l+1 ≥ρ l, and where for each layer l, the ρ of layer l in can be chosen as l-1 , and the ρ of layer l out can be chosen as l .

[0147] The system generates a shaped matrix based on one or more basis matrices (304). For example, to generate the shaped matrix, the system can decompose (e.g., Cholesky decomposition or lower-upper (LU) decomposition) a first basis matrix into a set of first decomposition matrices, e.g., such that the first basis matrix is ​​equal to the product of the first decomposition matrices. The system can also decompose a second basis matrix into a set of second decomposition matrices, e.g., such that the second basis matrix is ​​equal to the product of the second decomposition matrices. The system can then generate the shaped matrix based on the product of: (i) a first decomposition matrix from the set of first decomposition matrices, and (ii) an inverse matrix of the second decomposition matrix from the set of second decomposition matrices.

[0148] More specifically, in some embodiments, the system may generate the shaping matrix S as follows:

[0149]

[0150] Among them, L out is the basis matrix Σ out The Cholesky decomposition of, for example, as defined in equation (12) or (14), and L in is the basis matrix Σ in Cholesky decomposition of , for example, as defined in equation (11) or (13).

[0151] As another example implementation, the system may generate the shaping matrix S as follows:

[0152]

[0153] Σ0=(1-r)I T +r11 T (19)

[0154] Where r is a positive constant, L out is the basis matrix Σ out The Cholesky decomposition of, for example, as defined in equation (12) or (14), and L in is the basis matrix Σ in Cholesky decomposition of , for example, as defined in equation (11) or (13).

[0155] The system decomposes the shaped matrix into the product of: (i) a diagonal matrix with positive diagonal entries, and (ii) a matrix with row sums equal to 1, referred to herein as a block matrix (306).

[0156] The system applies an element-wise logarithm operation to the block matrix (308).

[0157] The system outputs a set of shaping constants (310). For example, the entries of a matrix obtained by applying an element-wise logarithm operation to a block matrix may define the shaping constants. A diagonal matrix may be parameterized to embed a specific rescaling operation, for example, as described in reference Figure 2 As described in equation (9) in the description of .

[0158] Figure 4 is a flow chart of an example process 400 for training a neural network including one or more regularized attention layers. For convenience, process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system appropriately programmed according to the present specification, e.g., Figure 1 The neural network system 100 can perform process 400.

[0159] The system obtains training data for training a neural network (402). The training data includes a set of training examples. In some cases, the set of training examples may include, for example, labeled training examples, each training example including: (i) a training input to the neural network, and (ii) a target output that should be generated by the neural network by processing the training input. In some cases, the set of training examples may include, for example, reinforcement learning training examples, each example including: (i) a trajectory representing an agent's interaction with an environment over a sequence of one or more time steps, and (ii) data describing a reward generated during the trajectory.

[0160] The system initializes values ​​for a set of parameters of the neural network (404). In some cases, the system initializes values ​​for the set of parameters of the neural network such that processing the network inputs according to the initialized values ​​of the neural network parameters causes each regularized attention layer to generate an intermediate attention score having a value of zero (or a value within a tolerance range around zero). For example, the system may initialize the trainable weight matrix W Q The parameter values ​​of (as described with reference to Equation (4)) are initialized to zero (or within a tolerance range around zero, or sampled from a probability distribution with a mean of zero and a standard deviation within a tolerance range around zero). As another example, the system may initialize the trainable weight matrix W to KThe parameter values ​​of (as described with reference to equation (5)) are initialized to zero (or within a tolerance range around zero, or sampled from a probability distribution with mean zero and standard deviation within a tolerance range around zero).

[0161] Initializing the values ​​of the neural network parameters so that the intermediate attention scores are initially zero (or close to zero) causes the attention operation to be initially controlled by the value of the shaping constant. The value of the shaping constant is chosen to regularize the output of the regularized attention layer, and regularization can be particularly effective when the intermediate attention scores are zero (or close to zero). During training, the values ​​of the neural network parameters are iteratively adjusted so that the intermediate attention scores are no longer necessarily zero (or close to zero). However, initializing the neural network in an operating mode with strong regularization increases the likelihood that the benefits of the regularization will persist throughout training and carry through to the trained neural network.

[0162] The remaining parameters of the neural network, e.g., parameters not set for zeroing the intermediate attention scores, can be initialized using any appropriate initialization scheme, e.g., glorot initialization or random initialization.

[0163] The system trains the neural network on the set of training data (406). The system can train the neural network through a series of training iterations to iteratively adjust parameter values ​​of the neural network to optimize an objective function. The objective function can be, for example, a cross entropy objective function or a squared error objective function (e.g., if the neural network is trained to perform a classification or regression task), or a reinforcement learning objective function (e.g., a Q-learning objective function, if the neural network is trained to select actions for controlling an agent). At each training iteration, the system can determine the gradient of the objective function with respect to a set of neural network parameters of the neural network, for example using backpropagation, and use the gradient to adjust the values ​​of the set of neural network parameters, for example according to an update rule of an appropriate gradient descent optimization algorithm (e.g., RMSprop or Adam).

[0164] Figure 5 The evolution of the normalized kernel matrices of the attention layers in a neural network is shown. Specifically, row 502 shows the kernel matrices of layers 1, 5, 10, 20, 50, and 100 in a neural network with traditional attention layers that do not implement the regularized attention operation described in this specification. Rows 504 and 506 show the kernel matrices of layers 1, 5, 20, 20, 50, and 100 in a neural network with two possible implementations of the regularized attention operation described in this specification. Using traditional attention layers, as shown in row 502, results in rank collapse, where all entries of the normalized kernel converge to one. In contrast, using regularized attention layers as shown in rows 504 and 506, controlled signal propagation can be maintained even at greater depths.

[0165] Figure 6 Training losses are shown for a range of training step sizes for various possible neural network architectures, including: (i) "e-spa", "u-spa", and "value skipinit", which implement the regularized attention operations described in this specification, and (ii) other alternatives such as "skip+LN" (implementing skip connections and normalization layers), "skipless" (not implementing any skip connections), "skipless+LN" (implementing normalization layers but no skip connections), and "skipless+LN+DKS" (implementing normalization layers and deep kernel shaping but no skip connections). Using the techniques described in this specification to regularize the attention operations performed by attention layers can allow deep neural networks with many attention layers to achieve acceptable performance (i.e., training loss is higher than skip+LN, but not unacceptable), even if no skip connections or normalization layers are used.

[0166] Figure 7 The training performance of a neural network with a regularized attention layer is shown as a function of the trainable weight matrix W. Q (as described in reference equation (4)) and W K (as described with reference to equation (5)). Q and / or W K Initializing the matrix parameters to zero (or close to zero) can improve the performance of the neural network, as shown in Figure 4 From Figure 7 It can be understood that with W Q and / or W K As the initial parameter values ​​of deviate more significantly from zero (in particular, as a result of sampling from a probability distribution with a higher standard deviation σ, especially one closer to 1), the training loss increases.

[0167] This specification uses the term "configuration" in conjunction with system and computer program components. For a system of one or more computers to be configured to perform specific operations or actions, it means that the system has installed thereon software, firmware, hardware, or a combination thereof, which, in operation, causes the system to perform those operations or actions. One or more computer programs to be configured to perform specific operations or actions means that the one or more programs include instructions that, when executed by a data processing device, cause the device to perform those operations or actions.

[0168] Embodiments of the subject matter and functional operations described in this specification may be implemented in digital electronic circuits, tangibly embodied computer software or firmware, computer hardware, including the structures disclosed in this specification and their structural equivalents, or a combination of one or more of them. Embodiments of the subject matter described in this specification may be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non-volatile storage medium, for execution by a data processing apparatus or for controlling the operation of a data processing apparatus. A computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or serial access storage device, or a combination of one or more of them. Alternatively or additionally, program instructions may be encoded on an artificially generated propagation signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, which is generated to encode information for transmission to a suitable receiver device for execution by a data processing apparatus.

[0169] The term "data processing apparatus" refers to data processing hardware and includes all kinds of apparatus, devices and machines for processing data, including, for example, a programmable processor, a computer or a plurality of processors or computers. The apparatus may also be or further include special-purpose logic circuits, for example, an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). In addition to hardware, the apparatus may optionally include code that creates an execution environment for a computer program, for example, code constituting processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0170] A computer program may also be referred to or described as a program, software, software application, application, module, software module, script, or code, and may be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it may be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program may be stored in a portion of a file that holds other programs or data, such as one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple collaborative files, such as files storing one or more modules, subroutines, or portions of code. A computer program may be deployed to execute on one computer or on multiple computers located at one location or distributed at multiple locations and interconnected by a data communications network.

[0171] In this specification, the term "engine" is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Typically, an engine will be implemented as one or more software modules or components installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a specific engine; in other cases, multiple engines can be installed and run on the same one or more computers.

[0172] The processes and logic flows described in this specification can be performed by one or more programmable computers that execute one or more computer programs to perform functions by operating on input data and generating output. These processes and logic flows can also be performed by special purpose logic circuits (e.g., FPGAs or ASICs), or by a combination of special purpose logic circuits and one or more programmed computers.

[0173] Computers suitable for executing computer programs can be based on general or special microprocessors or both, or any other kind of central processing unit. Typically, the central processing unit will receive instructions and data from a read-only memory or a random access memory or both. The basic elements of a computer are a central processing unit for executing instructions and one or more storage devices for storing instructions and data. The central processing unit and the memory can be supplemented or incorporated therein by a dedicated logic circuit. Typically, the computer will also include or be operably coupled to one or more large-capacity storage devices for storing data, such as a disk, a magneto-optical disk, or an optical disk, to receive data from the large-capacity storage device or to transmit data to it, or both. However, the computer does not need to have such a device. In addition, the computer can be embedded in another device, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device, such as a universal serial bus (USB) flash drive, to name just a few examples.

[0174] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media, and storage devices, including, for example, semiconductor memory devices, such as EPROM, EEPROM, and flash memory devices; magnetic disks, such as internal hard disks or removable disks; magneto-optical disks; and CD-ROM and DVD-ROM disks.

[0175] To provide interaction with a user, embodiments of the subject matter described in this specification may be implemented on a computer having a display device for displaying information to the user, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, and a keyboard and pointing device, such as a mouse or trackball, that the user can use to provide input to the computer. Other types of devices may also be used to provide interaction with the user; for example, the feedback provided to the user may be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback; and input from the user may be received in any form, including sound, voice, or tactile input. In addition, the computer may interact with the user by sending documents to and receiving documents from a device used by the user; for example, by sending a web page to a web browser on a user's device in response to a request received from the web browser. In addition, the computer may interact with the user by sending a text message or other form of message to a personal device (e.g., a smartphone running a messaging application) and receiving a response message from the user.

[0176] The data processing apparatus for implementing the machine learning model may also include, for example, dedicated hardware accelerator units for processing the general and computationally intensive parts of machine learning training or production, i.e., inference workloads.

[0177] Machine learning models can be implemented and deployed using a machine learning framework (e.g., the TensorFlow framework or the Jax framework).

[0178] Embodiments of the subject matter described in this specification may be implemented in a computing system that includes a back-end component, e.g., as a data server, or includes a middleware component, e.g., an application server, or includes a front-end component, e.g., a client computer with a graphical user interface, a web browser, or an application through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back-end, middleware, or front-end components. The components of the system may be interconnected by any form or medium of digital data communication (e.g., a communication network). Examples of communication networks include local area networks (LANs) and wide area networks (WANs), e.g., the Internet.

[0179] A computing system may include a client and a server. The client and the server are typically remote from each other and typically interact via a communication network. The relationship of the client and the server is due to the computer programs running on the respective computers and have a client-server relationship with each other. In some embodiments, the server sends data, such as an HTML page, to a user device, for example, in order to display data to a user interacting with the device acting as a client and receive user input from the user. Data generated at the user device, such as the result of a user interaction, may be received at the server from the device.

[0180] Although this specification contains many specific implementation details, these should not be interpreted as limitations on the scope of any invention or the scope that may be claimed, but rather as descriptions of features that are unique to specific embodiments of specific inventions. Certain features described in this specification in the context of independent embodiments may also be implemented in combination in a single embodiment. Conversely, the various features described in the context of a single embodiment may also be implemented in multiple embodiments individually or in any suitable sub-combination. In addition, although features may be described above as working in certain combinations, and even initially claimed as such, one or more features from the claimed combination may be deleted from the combination in some cases, and the claimed combination may be directed to a sub-combination or a variant of a sub-combination.

[0181] Similarly, although operations are described in a particular order in the drawings and recited in a particular order in the claims, this should not be understood as requiring that the operations be performed in the particular order or order shown, or that all of the operations shown be performed, to achieve the desired results. In some cases, multitasking and parallel processing may be advantageous. In addition, the separation of various system modules and components in the above-described embodiments should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

[0182] Specific embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve the desired results. As an example, the processes depicted in the accompanying drawings do not necessarily require the particular order shown or sequential order to achieve the desired results. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. A method performed by one or more computers, the method comprising: Receive network input; processing the network input using a neural network comprising a plurality of neural network layers arranged as a directed graph to generate a network output for the network input, wherein the plurality of neural network layers include one or more regularized attention layers, and Among them, for each regularized attention layer, processing network input includes: Receiving a layer input to a regularized attention layer, wherein the layer input to the regularized attention layer comprises a set of input embeddings; and Apply a regularized attention operation to a set of input embeddings to produce a set of output embeddings, including: Processing the set of input embeddings according to the values ​​of the set of regularized attention layer parameters to generate: (i) a set of value embeddings, including a corresponding value embedding for each input embedding, and (ii) a set of intermediate attention scores; The intermediate attention scores are transformed using a set of shaping constants to generate a set of transformed attention scores, where: The values ​​of the shaping constants are initialized prior to training of the neural network and are not adjusted during training of the neural network; and The value of the shaping constant is chosen to regularize the set of output embeddings; and generating a set of output embeddings using (i) the set of value embeddings, and (ii) the set of transformed attention scores; and Provides layer outputs for the attention layer based on the set of output embeddings.

2. The method according to claim 1, wherein: The value of the shaping constant is chosen to maintain or increase the rank of the set of output embeddings.

3. The method according to claim 2, wherein: The value of the shaping constant is chosen to increase the likelihood that the rank of the set of output embeddings exceeds a threshold.

4. A method according to any preceding claim, wherein: The values ​​of the set of shaping constants are derived from the shaping matrix by the following operations, which include: determining a decomposition of the shaped matrix into the product of (i) a diagonal matrix and (ii) a block matrix, wherein the block matrix has row sums equal to 1; and Logarithms are applied to the block matrix, and the shaped matrix is ​​based on the result of applying the logarithms to the block matrix.

5. The method according to claim 4, wherein: The shaped matrix is ​​derived from at least one base matrix, wherein values ​​of off-diagonal entries of the base matrix decay exponentially based on distance from a diagonal of the base matrix.

6. The method according to claim 5, wherein: The entries on each diagonal of the basis matrix have the same value.

7. The method according to claim 6, wherein: Each diagonal entry of the basis matrix has the value 1.

8. The method according to claim 4, wherein: The shaped matrix is ​​derived from at least one base matrix, wherein diagonal entries of the base matrix each have a same first value and non-diagonal entries of the base matrix each have a same second value.

9. The method according to claim 8, wherein: The first value is 1 and the second value is strictly less than 1.

10. A method according to any preceding claim, wherein: Transforming the intermediate attention scores using the set of shaping constants to generate a set of transformed attention scores includes, for each intermediate attention score: By combining the intermediate attention score with the corresponding shaping constant, the corresponding transformed attention score is generated.

11. The method according to claim 10, wherein: For each intermediate attention score, generating the corresponding transformed attention score includes: Add the intermediate attention score to the corresponding shaping constant.

12. The method according to any one of claims 4 to 11, wherein: The set of input embeddings processed to generate intermediate attention scores includes: processing the collection of input embeddings to generate, for each input embedding: (i) a corresponding query embedding, and (ii) a corresponding key embedding; and Each intermediate attention score is generated based on a measure of similarity between the corresponding query embedding and the corresponding key embedding.

13. The method according to any one of claims 4 to 12, wherein: The set of output embeddings generated using (i) the set of value embeddings and (ii) the set of transformed attention scores includes: generating a set of final attention scores by applying a causal masking operation to the set of transformed attention scores followed by a nonlinear transformation; and Generate a set of output embeddings using (i) a set of value embeddings and (ii) a set of transformed attention scores.

14. The method according to claim 13, wherein: The nonlinear transformation is the soft-max transformation.

15. The method according to any one of claims 13-14, wherein: The set of output embeddings generated using (i) the set of value embeddings and (ii) the set of transformed attention scores includes: Each output embedding is generated based on a linear combination of a set of value embeddings, where coefficients of the linear combination are defined by corresponding transformed attention scores from the set of transformed attention scores.

16. The method according to claim 15, wherein: The set of generated output embeddings includes: Applies an embedding-specific rescaling to the set of output embeddings based on a diagonal matrix.

17. A method according to any preceding claim, wherein: Before training of the neural network, the values ​​of the regularized attention layer parameters are initialized so that the value of each of the intermediate attention scores is zero.

18. A method according to any preceding claim, wherein: Before training of the neural network, the values ​​of the regularized attention layer parameters are initialized to encourage the value of each of the intermediate attention scores to be close to zero.

19. A method according to any preceding claim, wherein: Prior to training of the neural network, the values ​​of the regularized attention layer parameters are chosen to be: (i) initialized to zero, or (ii) chosen from a tolerance range around zero, or (iii) sampled from a probability distribution with a standard deviation chosen from a tolerance range around zero.

20. A method according to any preceding claim, wherein: The neural network includes a plurality of regularized attention layers associated with the ranking, wherein each of the plurality of regularized attention layers is associated with a different set of shaping constants.

21. A method according to any preceding claim, wherein: The neural network is configured to autoregressively generate the output sequence.

22. A method according to any preceding claim, wherein: The neural network does not include skip connections associated with the regularized attention layers in the neural network.

23. A method according to any preceding claim, wherein: The neural network does not include the normalization layers associated with the regularized attention layers in the neural network.

24. A method according to any preceding claim, wherein: The set of regularized output embeddings includes one or more of the following: (i) modifying the set of output embeddings to maintain or reduce the norm of the output embeddings included in the set of output embeddings, (ii) modifying the set of output embeddings to preserve or increase the norm of the gradient of the objective function used to train the neural network, or (iii) modifying the set of output embeddings to maintain or increase the rank of the set of output embeddings.

25. A system comprising: one or more computers; as well as One or more storage devices, the one or more storage devices being communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform the operations of the corresponding method of any one of claims 1-23.

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

Images