Learning to memorize at inference

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for learning to memorize inputs in weights of a neural long-term memory at inference. In one aspect, a token from a sequence of tokens is received and projected using respective linear layers to obtain a key and a value for the token. The key is processed using a neural long-term memory to generate a memory output. A gradient of a loss with respect to weights of the neural long-term memory is computed. Weights of the neural long-term memory are updated using the gradient without updating parameters of a core neural network. Queries for one or more of the tokens are processed using the neural long-term memory to generate respective memory outputs and an input comprising the respective memory outputs is processed using the core neural network to generate one or more network outputs.

Description

[0001] Atorney Docket No. 56113-0862WO1

[0002] LEARNING TO MEMORIZE AT INFERENCE CROSS REFERENCE TO RELATED APPLICATIONS

[0003] This application claims priority to U.S. Application No.63 / 753,787, filed February 4, 2025. The disclosure of the foregoing application is hereby incorporated by reference in its entirety7.

[0004] BACKGROUND

[0005] This specification relates to neural networks.

[0006] Neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current value inputs of a respective set of parameters.

[0007] SUMMARY

[0008] This specification describes a neural long-term memory implemented as computer programs on one or more computers in one or more locations that efficiently leams to memorize inputs in weights of the neural long-term memory7at inference.

[0009] One innovative aspect of the subject matter described in this specification can be implemented in a method that includes processing, in sequence, each token in a sequence of tokens, the processing comprising: receiving a token from the sequence of tokens; projecting the token using a first linear layer to obtain a key for the token; projecting the token using a second linear layer to obtain a value for the token; processing the key for the token using a neural long-term memory to generate a memory output that corresponds to the key for the token; computing a gradient of a loss with respect to weights of the neural long-term memory7, wherein the loss is dependent on a difference between the memory output that corresponds to the key for the token and the value for the token; and updating weights of the neural long-term memory using the gradient without updating parameters of a core neural network; processing a respective query for each of one or more of the tokens using the neural long-term memory7to generate a respective memory7output for each query7; and processing an input comprising the respective memory7outputs using the core neural network to generate one or more network outputs.Atorney Docket No. 56113-0862WO1

[0010] Other implementations of these aspects includes corresponding computer systems, apparatus, and computer programs recorded on one or more computer storage devices, each configured to perform the actions of the methods. A system of one or more classical computers can be configured to perform particular operations or actions by virtue of having software, firmw are, hardware, or a combination thereof installed on the system that in operation causes or cause the system to perform the actions. One or more computer programs can be configured to perform particular operations or actions by virtue of including instructions that, when executed by data processing apparatus, cause the apparatus to perform the actions.

[0011] The foregoing and other implementations can each optionally include one or more of the following features, alone or in combination. In some implementations the method further comprises computing a surprise metric for the token, the surprise metric comprising a combination of a surprise metric for a previous token in the sequence of tokens and the computed gradient of the loss, wherein the w eights of the neural long-term memory are updated using the surprise metric for the token.

[0012] In some implementations updating the weights of the neural long-term memory using the gradient comprises weighting current weights of the neural long-term memory using a gating mechanism.

[0013] In some implementations updating w eights of the neural long-term memory using the gradient further comprises updating weights of the neural long-term memory without updating parameters of the first linear layer and the second linear layer.

[0014] In some implementations the neural long-term memory' comprises a single-layer perceptron.

[0015] In some implementations computing the gradient of the loss comprises multiplying a difference between i) a product of a weight matrix for the neural long-term memory and the token and ii) the token by a transpose of the token.

[0016] In some implementations processing the key using the neural long-term memory7to generate the memory7outputs that corresponds to the key comprises multiplying a weight matrix for the neural long-term memory by the key.

[0017] In some implementations the method further comprises computing gradients of the loss in parallel for each token in a chunk of tokens included in the sequence of tokens; and applying the computed gradients sequentially to update the w eights of the neural long-term memory.Atorney Docket No. 56113-0862WO1

[0018] In some implementations computing the gradients of the loss in parallel for each token in the chunk comprises performing a batched matrix multiplication operation over the tokens in the chunk.

[0019] In some implementations updating weights of the neural long-term memory comprises computing surprise metrics for the tokens in the chunk using a parallel scan operation and storing intermediate update terms for each chunk.

[0020] In some implementations the first linear layer and the second linear layer comprise learned parameters of the core neural network.

[0021] In some implementations the sequence of tokens comprises an attention output generated by the core neural network and the method further comprises: obtaining, by the core neural network, a network input; dividing, by the core neural network, the network input into a plurality of fixed-size segments; for each segment in the pl ural i ty of fixed-size segments: retrieving, by the core neural network, memory outputs for the segment from the neural long-term memory and retrieve persistent memory parameters from a persistent memory: applying, by the core neural network, an attention mechanism to a concatenation of the persistent memory parameters, the segment , and the memory outputs for the segment to obtain the sequence of tokens; and processing the sequence of tokens using the neural long-term memory' to update the weights of the neural long-term memory.

[0022] In some implementations processing a respective query for each of one or more of the tokens using the neural long-term memory^ to generate a respective memory' output for each query comprises processing each token in the sequence of tokens as a respective query using the neural long-term memory to generate a respective memory output for each token in the sequence of tokens; and the input comprises the respective memory outputs comprises a combination of the attention output and the memory outputs.

[0023] In some implementations retrieving memory outputs for the segment from the neural long-term memory' comprises: generating a query' using the segment, comprising multiplying the segment by a weight matrix of learned parameters; and processing the query using the neural long-term memory to retrieve the memory outputs for the segment without updating weights of the neural long-term memory'.

[0024] In some implementations processing the attention output using the updated neural long-term memory to retrieve the memory outputs for the attention output comprises processing the attention output using the long-term memory without updating weights of the neural long-term memory.Atorney Docket No. 56113-0862WO1

[0025] In some implementations the sequence of tokens comprises a concatenation of persistent memory parameters and a network input, and wherein the method further comprises: obtaining, by the core neural network, the network input; retrieving, by the core neural network, the persistent memory parameters from a persistent memory; applying, by the core neural network, an attention mechanism to the sequence of tokens to obtain an attention output; and processing the sequence of tokens using the neural longterm memory to recursively update weights of the neural long-term memory, wherein the input comprising the respective memory outputs comprises a combination of the attention output and the respective memory outputs.

[0026] In some implementations the attention mechanism comprises a sliding window attention with prefix, wherein the sliding window attention acts as a short-term memory.

[0027] In some implementations the combination of the attention output and final updated weights of the neural long-term memory comprises a non-linear gating of the attention output and the final updated weights of the neural long-term memory.

[0028] In some implementations the method further comprises normalizing, by the core neural network, the attention output and the final updated weights of the neural long-term memory using learned vector-valued weights and a non-linearity function.

[0029] In some implementations the sequence of tokens comprises a concatenation of persistent memory parameters and a network input, and wherein the method further comprises: obtaining, by the core neural network, the network input; retrieving, by the core neural network, the persistent memory parameters from a persistent memory; processing the sequence of tokens using the neural long-term memory to recursively update weights of the neural long-term memory; and applying, by the core neural network, an attention mechanism to a final updated state of the neural long-term memory to obtain an attention output, wherein the input comprising the respective memory outputs comprises the attention output.

[0030] In some implementations the attention mechanism comprises a sliding window attention.

[0031] In some implementations the neural long-term memory comprises a multi-layer perceptron with two or more layers.

[0032] In some implementations one or more of the layers apply an expansion factor. In some implementations one or more of the layers applies a Gaussian Error Linear Unit activation function.Atorney Docket No. 56113-0862WO1

[0033] In some implementations different layers in the neural long-term memory are assigned respective learning parameters for updating weights of the neural long-term memory.

[0034] In some implementations the learning parameters comprise a learning rate, wherein the learning rates for the two or more layers are scaled such that increasingly deeper layers in the neural long-term memory have slower learning rates.

[0035] In some implementations the learning parameters comprise input-dependent parameters determined using a linear projection of the input.

[0036] Some implementations of the subject matter described herein may realize, in certain instances, one or more of the following advantages.

[0037] Conventional sequence modeling techniques face several fundamental limitations that hinder their applicability to tasks requiring long-range dependency modeling, continual adaptation, and efficient utilization of computational resources. For example, pure attention-based architectures, such as Transformers, can provide accurate modeling of token-to-token dependencies but incur quadratic time and memory complexity with respect to sequence length. As a result, such architectures are constrained to fixed or relatively short context windows, making them impractical for applications involving very7long sequences, streaming data, or persistent contextual reasoning.

[0038] To mitigate these scalability limitations, various linear attention and recurrent models have been proposed. While such approaches reduce computational complexity, they typically rely on fixed-size vector-valued or matrix-valued state representations that compress historical information into a limited memory7footprint. This compression restricts the capacity of the model to faithfully represent long or complex histories and often leads to rapid memory saturation, degraded performance on long sequences, and poor length extrapolation. Moreover, many existing recurrent or linear models rely on simplified or purely additive memory update rules that lack effective mechanisms for selectively forgetting obsolete information or emphasizing salient events.

[0039] Additionally, some conventional architectures treat learning and inference as strictly separate phases, with model parameters fixed at inference time. As a result, these architectures are unable to adapt their internal representations to new or surprising inputs encountered during deployment, limiting robustness to distribution shift, long-context reasoning, and dynamic environments. Attempts to incorporate test-time adaptation or external memory mechanisms often introduce substantial computational overhead, disruptAtorney Docket No. 56113-0862WO1

[0040] inference throughput, or require specialized retrieval pipelines that operate outside the core model architecture.

[0041] Examples of the presently described techniques address these limitations by introducing a neural long-term memory module that is explicitly configured to leam, update, and forget information at test time. The described neural network architectures achieve high accuracy on long-context reasoning, recall-intensive tasks, time-series modeling, and sequence prediction, while maintaining efficient training and inference performance. In addition, trade-offs between scalability, expressiveness, and efficiency can be further reduced by incorporating one or more of inference-time memory' updates, controlled forgetting, non-linear memory' structures, or parallelizable computation.

[0042] For example, the neural long-term memory can be trained using a gradient-based update rule derived from an associative memory' loss, enabling the memory to store abstractions of historical input data directly in its parameters. This formulation allows the memory' to adapt dynamically to incoming data without modify ing parameters of the core neural network, thereby preserving task-level behavior while enabling continual memorization.

[0043] Further, in contrast to fixed-size linear memories, the presently described neural long-term memory' includes a deep, non-linear architecture, providing substantially greater expressive capacity for representing long-range dependencies. By incorporating momentum-based update terms, the memory’ can accumulate information about surprising or salient events over time, rather than responding solely to momentary inputs. In addition, an adaptive forgetting or decay mechanism can selectively attenuate previously stored information, enabling effective memory' management over extremely long sequences and preventing uncontrolled memory growth.

[0044] A further technical advantage of the presently described techniques is that the memory update process can be formulated to support efficient parallel execution. In some implementations, gradient computations for multiple tokens can be performed in parallel within chunks, while recurrence relations governing momentum and decay are evaluated using parallel associative scan or convolution-like operations. This decomposition preserves the semantic behavior of sequential memory' updates while significantly improving throughput on modem parallel hardware, such as GPUs or TPUs. As a result, the neural long-term memory' can be trained and updated with linear computational complexity in the sequence length, even for contexts spanning millions of tokens.Atorney Docket No. 56113-0862WO1

[0045] In addition, examples of the presently described techniques further provide flexible mechanisms for integrating the neural long-term memory with short-term attention and persistent task memory. By supporting multiple architectural variants — including memory as a context, memory as a gated branch, and memory as a layer — the system enables complementary processing of local dependencies, long-term history, and task-specific knowledge. This modular design allows each memory’ component to operate independently while remaining interoperable, improving robustness, interpretability, and extensibility.

[0046] The details of one or more embodiments of the subject matter described in this specification are set forth in the accompanying drawings and the description below. Other potential features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.

[0047] DESCRIPTION OF DRAWINGS FIG. 1 is a block diagram of an example neural network architecture.

[0048] FIG. 2 is a block diagram of an example memory-as-context neural network architecture.

[0049] FIG. 3 is a block diagram of an example gated-memory neural network architecture.

[0050] FIG. 4 is a block diagram of an example memory-as-layer neural network architecture.

[0051] FIG. 5 is a flow chart of an example process for updating and query ing a neural long-term memory during inference.

[0052] FIG. 6 illustrates an example process for training and updating a neural long-term memory using parallelized computation.

[0053] Like reference symbols in the various drawings indicate like elements.

[0054] DETAILED DESCRIPTION

[0055] This specification describes a neural long-term memory implemented as computer programs on one or more computers in one or more locations that efficiently learns to memorize inputs in weights of the neural long-term memory' at inference. The surprise of an input to the neural long-term memory is measured by the gradient of the current state of the neural long-term memory with respect to the input in associative memory loss. TheAtorney Docket No. 56113-0862WO1

[0056] surprise is included in an update rule which is used to update the weights of the neural long-term memory. By construction, the weight updates can be calculated using only matrix multiplications and summations, which can be performed in parallel. Further, a decaying mechanism that considers the proportion of memon size and the amount of data surprise can be implemented to improve memory7management.

[0057] The neural long-term memory modules and update mechanisms described herein can be understood within a broader theoretical framework for online sequence modeling. For example, the MIRAS (Memory, Attentional Bias, Retention, and Online Optimization) framework generalizes modem sequence models by viewing them as instances of associative memory systems in which memory architecture, attentional bias, retention gates, and online update algorithms are design choices that govern long-range information retention and adaptation. The MIRAS framework emphasizes real-time adaptation and test-time memorization, in which model parameters are actively updated as new data is processed rather than being fixed after offline training. Within this framework, the presently described neural long-term memory, surprise-driven gradient updates, and retention mechanisms provide specific choices for memory architecture and retention control that enable improved long-context retention, adaptability to surprising or novel inputs, and efficient sequence processing.

[0058] FIG. 1 is a block diagram of an example neural network architecture 100. The example neural network architecture 100 includes a neural long-term memory 102. a core neural network 104, and a persistent memory 106. The core neural network 104 is configured to receive an input 108 for a machine learning task. The input 108 can be, e.g., a sequence of tokens, an embedding sequence, or other representation of text, image, audio, video, multimodal, or control data. The core neural network 104 is configured to process the input 108 to generate a corresponding network output 110 for the machine learning task. During the processing of the input 108, the core neural network 104 generates intermediate token representations, e.g., embeddings, and uses the neural longterm memory 102 and the persistent memory7106 to condition the intermediate token representations and the resulting network output 110.

[0059] For example, the persistent memory 106 can provide one or more persistent parameter vectors that are incorporated into a context sequence processed by the core neural netw ork 104. The persistent parameter vectors can be concatenated to a start portion of an input token sequence or an internal context sequence processed by the core neural netw ork 104, such that the persistent parameter vectors are attended to orAtorney Docket No. 56113-0862WO1

[0060] otherwise processed jointly with tokens derived from the input 108. as described in more detail below. In this manner, the persistent memory 106 provides input-independent task information that influences processing performed by the core neural network 104.

[0061] The neural long-term memory 102 is configured to store and retrieve information that is dependent on tokens processed during inference. During inference, the core neural network 104 can provide to the neural long-term memory 102 a sequence of tokens or token representations generated while processing the input 108, and the neural long-term memory 102 can update its weights based on the sequence of tokens, as described in more detail below. Further, during inference, the core neural network 104 can generate one or more memory' queries, e.g., derived from a current token representation, and provide the memory queries to the neural long-term memory 102 to retrieve corresponding memory outputs. The neural long-term memory 102 provides the memory outputs to the core neural network 104 for use in computing subsequent token representations and / or the network output 110.

[0062] Accordingly, the core neural network 104 generates the network output 110 based on (i) representations derived from the input 108, (li) outputs derived from the persistent memory 106 that encode task-level or system-level information, and (iii) outputs retrieved from the neural long-term memory' 102 that encode abstractions of previously processed tokens at inference time. In some implementations, the neural long-term memory 102 adapts during inference based on the sequence of tokens provided by the core neural network 104, while parameters of the core neural network 104 and the persistent memory 106 remain fixed during inference, thereby enabling test-time memorization while maintaining stable task-level behavior.

[0063] In more detail, at inference the neural long-term memory 102 can receive a sequence of tokens from the core neural network. The neural long-term memory 102 processes the tokens in the sequence of tokens in sequence, e.g., one token at a time or in chunks of tokens. Each token is projected using a first linear layer to obtain a key for the token. Further, each token is projected using a second linear layer to obtain a value for the token. The first linear layer and the second linear layer are learned parameters of the core neural network, e.g., learned during training of the core neural network.

[0064] In this manner, the neural long-term memory 102 operates as a parameterized memory' function whose parameters encode an abstraction of previously processed tokens. The abstraction is stored directly in the weights of the neural long-term memory,Atorney Docket No. 56113-0862WO1

[0065] allowing historical information to be compressed into the parameter space of the memory rather than stored explicitly as raw token data.

[0066] The key for the token is processed using the neural long-term memory 102 to generate a memory output that corresponds to the key for the token. A gradient of a loss with respect to weights of the neural long-term memory 102 is computed. The loss is dependent on a difference between the memory output that corresponds to the key for the token and the value for the token.

[0067] In one implementation, given a token xt, the key and value are computed as

[0068]

[0069] >

[0070] and the loss is defined as an associative memory loss

[0071]

[0072] The gradient of this loss with respect to the weights of the neural long-term memory 102 serves as a data-dependent learning signal that reflects how much the current token deviates from information previously encoded by the memory .

[0073] In cases where the neural long-term memory 102 processes the tokens in the sequence of tokens as a chunk of tokens, the gradients for the tokens in the chunk can be computed in parallel for each token in the chunk, e.g., through a batched matrix multiplication operation over the tokens in the chunk, and surprise metrics for the tokens in the chunk can be computed using a parallel scan operation.

[0074] To propagate learning signals across time, the update process may maintain an accumulated update state Stthat integrates gradients across multiple tokens. For example, the accumulated update state may be computed as

[0075]

[0076] controls decay of prior accumulated updates and 6tcontrols a learning rate associated with the current token. This accumulated update state allows learning signals associated with earlier tokens to influence updates corresponding to later tokens in the sequence, even when individual gradients become small.Atorney Docket No. 56113-0862WO1

[0077] Weights of the neural long-term memory' 102 are updated using the gradient. The weights of the neural long-term memory 102 are updated without updates to the core neural network or the first linear layer and the second linear layer.

[0078] In one update formulation, the weights of the neural long-term memory 102 are updated according to

[0079] Mt= Mt-i + $t>

[0080] such that the neural long-term memory' 102 adapts during inference while parameters of the core neural network and the projection layers remain fixed. This separation enables test-time memorization without overfitting the core neural network.

[0081] In cases where neural long-term memory' 102 processes the tokens in the sequence of tokens as a chunk of tokens, the gradients for the chunk of tokens can be applied sequentially to update the weights of the neural long-term memory.

[0082] When chunk-wise processing is used, the update rule may be reformulated so that gradient contributions for all tokens in a chunk are computed in parallel using matrix multiplication operations and then applied sequentially using precomputed scaling factors. This formulation allows the update process to be expressed using matrix multiplication and summation operations, enabling efficient execution on hardware accelerators.

[0083] The update rule used to perform the updates to the weights includes a combination of the current weights of the neural long-term memory' 102 and the surprise metric. A forgetting mechanism can be implemented by weighting the current weights of the neural long-term memory 102 using a gating mechanism.

[0084] In some implementations, the forgetting mechanism can be incorporated directly into the update rule as

[0085]

[0086] where atG [0,1] is a gating parameter that controls how much previously stored information is retained. By adjusting at, the neural long-term memory 102 can preserve prior abstractions, partially decay older information, or clear the memory entirely when a context change is detected.Atorney Docket No. 56113-0862WO1

[0087] As well as training the neural long-term memory' 102 to leam to memorize at inference time, the core neural network queries the neural long-term memory 102 at inference to retrieve information from the neural long-term memory'. The neural longterm memory' 102 processes the query' to generate a respective memory' output and provides the memory output to the core neural network. The core neural network can use the memory’ output to generate a corresponding output, e.g.. a network output.

[0088] Retrieval from the neural long-term memory 102 can be performed using a forward inference operation without updating the weights of the memory. For example, a query' qtderived from an input token may be computed using a projection layer, and the memory' output may be obtained as

[0089] yt= Mt(qt

[0090] The retrieved memory output represents information stored in the neural long-term memory 102 that is relevant to the query and may be incorporated by the core neural network when generating the network output.

[0091] In some implementations the neural long-term memory 102 can be implemented as a multi-layer perceptron with one or more layers, e g., two layers. In some cases one or more of the layers can apply an expansion factor, e.g., an expansion factor of 4.

[0092] Further, one or more of the layers can apply a nonlinear activation function, such as a Gaussian Error Linear Unit. Different layers in the neural long-term memory' 102 can be assigned respective learning parameters for updating weights of the neural long-term memory' 102, e.g., layer-wise learning, decay, and momentum parameters. The learning rates for respective layers can be scaled such that increasingly deeper layers in the neural long-term memory have slower learning rates, e.g., a first layer can be scaled by 1 and a second layer can be scaled by 0.1. The learning parameters can be input-dependent parameters, e.g., determined using a linear projection of the corresponding input.

[0093] As described above, updating the neural long-term memory 102 during inference corresponds to optimizing an associative memory loss using gradient-based updates with momentum and optional weight decay. A naive implementation of such updates processes tokens sequentially and therefore incurs computational cost that scales linearly with the length of the token sequence. To improve computational efficiency and to enable execution on parallel hardware architectures, such as GPUs or TPUs, the updateAtorney Docket No. 56113-0862WO1

[0094] process for the neural long-term memory' 102 can be reformulated to support parallel computation.

[0095] In some implementations, a sequence of tokens processed during inference can be divided into one or more contiguous chunks, each chunk including a plurality' of tokens. For a given chunk, gradients of the associative memory loss with respect to the weights of the neural long-term memory’ 102 can be computed for each token in the chunk in parallel. For example, keys and values corresponding to tokens in a chunk can be aggregated into matrices, and loss gradients can be computed using batched matrix multiplication operations.

[0096] In some examples, the update of the neural long-term memory' 102 can include applying, for each token, a gradient term scaled by a learning-rate parameter and optionally a decay term applied to existing memory weights. When processing a chunk of tokens, the cumulative effect of these per-token updates can be expressed as a weighted sum of gradient contributions associated with the tokens in the chunk, where the weights depend on learning-rate parameters and decay parameters associated with the tokens or the chunk. By expressing the cumulative update in this form, the gradient contributions for all tokens in the chunk can be computed in parallel using matrix multiplication and summation operations.

[0097] To incorporate momentum or accumulated update effects across tokens, the update process can maintain an accumulated update state that follows a linear recurrence relation across the sequence of tokens. In these examples, the accumulated update state depends on a prior accumulated state, a gradient contribution associated with a current token, and one or more decay parameters. Because this recurrence is linear, accumulated update states for a chunk of tokens can be computed using a parallel associative scan operation, allowing the accumulated update state for each token position within the chunk to be computed concurrently.

[0098] In some implementations, parameters controlling the update process, such as learning-rate parameters, decay parameters, and momentum parameters, can be defined as functions of individual tokens, resulting in token-dependent update behavior. In other implementations, one or more of these parameters can be defined at the chunk level and shared across all tokens within a chunk. Defining update parameters at the chunk level can further reduce computational overhead and memory usage, for example by allowing certain scaling factors to be represented as scalars rather than token-specific values.Atorney Docket No. 56113-0862WO1

[0099] When chunk-level parameters are used, the update dynamics within a chunk can correspond to a linear time-invariant system, enabling the accumulated update state to be computed using convolution-like operations that are efficiently supported by hardware accelerators. Regardless of whether token-level or chunk-level parameters are used, the parallelized update process enables the neural long-term memory 102 to be updated using a sequence of matrix multiplication, summation, and scan operations, thereby reducing latency and improving throughput relative to strictly sequential update schemes.

[0100] Accordingly, the neural long-term memop 102 can be updated during inference using a parallelized computation pipeline that preserves the functional behavior of sequential gradient-based updates while enabling efficient execution on modem parallel computing hardware.

[0101] The core neural network 104 is responsible for the main flow of data processing in the neural network architecture. The core neural network 104 can be configured to perform a machine learning task to output data that includes, for example, text data, image data, video data, audio data, or multimodal data that includes data in two or more different modalities.

[0102] Some examples of machine learning tasks, including generative tasks, that the core neural network 104 when implemented using one of the architectures described herein or other known architectures can be configured to perform follow.

[0103] In some cases, the core neural network 104 is a neural network that is configured to perform an image processing task, i.e., receive a network input that includes an image and to process the image to generate a network output for the input image. For example, the task may be image classification and the output generated by the core neural network 104 for a given image may be scores for each of a set of object categories, with each score representing an estimated likelihood that the image contains an image of an object belonging to the category. As another example, the task can be image embedding generation and the output generated by the core neural network 104 can be a numeric embedding of the input image. As yet another example, the task can be object detection and the output generated by the core neural network 104 can identify locations in the input image at which particular types of objects are depicted. As yet another example, the task can be image segmentation and the output generated by the core neural network 104 can assign each pixel of the input image to a category from a set of categories. In some other cases, the core neural network 104 is a neural network that is configured to performAtorney Docket No. 56113-0862WO1

[0104] an image generation task, where the input is a conditioning input and the output is a sequence of intensity value inputs for the pixels of an image.

[0105] As one example, the task may be a neural machine translation task. For example, if the network input to the core neural network 104 is a sequence of text, e.g., a sequence of words, phrases, characters, or word pieces, in one language, the network output generated by the neural network may be a translation of the sequence of text into another language, i.e., a sequence of text in the other language that is a translation of the input sequence of text. The vocabulary for the input tokens may be words, wordpieces or characters of the first language, and the vocabulary for the output tokens may be words, wordpieces or characters of the other language. As a particular example, the task may be a multi-lingual machine translation task, where a single neural network is configured to translate between multiple different source language - target language pairs. In this example, the source language text may be augmented with an identifier that indicates the target language into which the core neural network 104 should translate the source language text.

[0106] Some implementations may be used for automatic code generation. For example the input tokens may represent words, wordpieces or characters in a first natural language and the output tokens may represent instructions in a computer programming or markup language, or instructions for controlling an application program to perform a task e.g. build a data item such as an image or web page.

[0107] As another example, the task may be an audio processing task. For example, if the network input to the core neural network 104 is a sequence representing a spoken utterance, the network output generated by the core neural network 104 may be a score for each of a set of pieces of text, each score representing an estimated likelihood that the piece of text is the correct transcript for the utterance. As another example, if the network input to the core neural network 104 is a sequence representing a spoken utterance, the network output generated by the core neural network 104 can indicate whether a particular word or phrase (“hotword'’) was spoken in the utterance. As another example, if the network input to the core neural network 104 is a sequence representing a spoken utterance, the network output generated by the core neural network 104 can be a classification of the spoken utterance into one of a plurality of categories, for example an identity of the natural language in which the utterance was spoken.

[0108] As another example, the task can be a natural language processing or understanding task, e.g., an entailment task, a paraphrase task, a textual similarity task, aAtorney Docket No. 56113-0862WO1

[0109] sentiment task, a sentence completion task, a grammaticality task, and so on, that operates on a sequence of text in some natural language.

[0110] As another example, the task can be a text to speech task, where the network input is text in a natural language or features of text in a natural language and the network output is a spectrogram, a waveform, or other data defining audio of the text being spoken in the natural language.

[0111] As another example, the task can be a health prediction task, where the network input is a sequence derived from electronic health record data for a patient and the network output is a prediction that is relevant to the future health of the patient, e.g., a predicted treatment that should be prescribed to the patient, the likelihood that an adverse health event will occur to the patient, or a predicted diagnosis for the patient. Such electronic health data may, for example, comprise one or more sequences of physiological data taken from a patient, with the output being a corresponding prediction that relates to those sequences of data. Examples of physiological data and a corresponding prediction include: blood glucose measurements, with the prediction being a predicted future blood glucose measurement or the prediction of a hyper- or hypo-glycemic event; a heart rate, with the prediction being the presence or absence of a heart condition, or a future cardiac event; blood pressure measurements, with the prediction being the risk of a future heart condition; or the like.

[0112] As another example, the task can be a text generation task, where the network input is a sequence of text, and the network output is another sequence of text, e.g., a completion of the input sequence of text, a response to a question posed in the input sequence, or a sequence of text that is about a topic specified by the first sequence of text. As another example, the network input to the text generation task can be an input other than text, e.g., an image, and the network output sequence can be text that describes the network input.

[0113] In some implementations the input sequence represents data to be compressed, e.g. image data, text data, audio data, or any other type of data; and the output sequence a compressed version of the data. The input and output tokens may each comprise any representation of the data to be compressed / compressed data e g. symbols or embeddings generat ed / decoded by a respective neural network.

[0114] As another example, the task can be an agent control task, where the network input is a sequence of observations or other data characterizing states of an environment and the network output defines an action to be performed by the agent in response to theAtorney Docket No. 56113-0862WO1

[0115] most recent data in the sequence. The agent can be, e.g., a real-world or simulated robot, a control system for an industrial facility, or a control system that controls a different kind of agent. The observations may comprise sensor data captured by sensors associated with (e.g. part of) the agent, for example visual data, LIDAR data, sonar data, agent configuration data (e g. joint angles), agent orientation data, or the like.

[0116] In some implementations, the environment is a real-world environment, the agent is a mechanical (or electro-mechanical) agent interacting with the real-world environment, e g., a robot or an autonomous or semi-autonomous land, air, or sea vehicle operating in or navigating through the environment, and the actions are actions taken by the mechanical agent in the real-world environment to perform the task. For example, the agent may be a robot interacting with the environment to accomplish a specific task, e.g., to locate or manipulate an object of interest in the environment or to move an object of interest to a specified location in the environment or to navigate to a specified destination in the environment.

[0117] In these implementations, the observations may include, e.g., one or more of: images, object position data, and sensor data to capture observations as the agent interacts with the environment, for example sensor data from an image, distance, or position sensor or from an actuator. For example in the case of a robot, the observations may include data characterizing the current state of the robot, e.g., one or more of: joint position, joint velocity, joint force, torque or acceleration, e.g.. gravity-compensated torque feedback, and global or relative pose of an item held by the robot. In the case of a robot or other mechanical agent or vehicle the observations may similarly include one or more of the position, linear or angular velocity, force, torque or acceleration, and global or relative pose of one or more parts of the agent. The observations may be defined in 1, 2 or 3 dimensions, and may be absolute and / or relative observations. The observations may also include, for example, sensed electronic signals such as motor current or a temperature signal; and / or image or video data for example captured by a camera or a LIDAR sensor, e.g., data from sensors of the agent or data from sensors that are located separately from the agent in the environment.

[0118] In these implementations, the actions may be control signals to control the robot or other mechanical agent, e.g., torques for the joints of the robot or higher-level control commands, or the autonomous or semi-autonomous land, air, sea vehicle, e.g., torques to the control surface or other control elements e.g. steering control elements of the vehicle, or higher-level control commands. The control signals can include for example, position.Atorney Docket No. 56113-0862WO1

[0119] velocity, or force / torque / accel eration data for one or more joints of a robot or parts of another mechanical agent. The control signals may also or instead include electronic control data such as motor control data, or more generally data for controlling one or more electronic devices within the 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 or air or sea vehicle the control signals may define actions to control navigation e.g. steering, and movement e.g., braking and / or acceleration of the vehicle.

[0120] In some implementations the environment is a simulation of the above-described real-world environment, and the agent is implemented as one or more computers interacting with the simulated environment. For example, a system implementing the core neural network 104 may be used to select actions in the simulated environment during training or evaluation of the system and, after training, or evaluation, or both, are complete, the action selection policy may be deployed for controlling a real-world agent in the particular real-world environment that was the subject of the simulation. This can avoid unnecessary wear and tear on and damage to the real-world environment or real-world agent and can allow the core neural network 104 to be trained and evaluated on situations that occur rarely or are difficult or unsafe to re-create in the real-world environment. For example the system may be partly trained using a simulation of a mechanical agent in a simulation of a particular real-world environment, and afterwards deployed to control the real mechanical agent in the particular real-world environment. Thus in such cases the observations of the simulated environment relate to the real-world environment, and the selected actions in the simulated environment relate to actions to be performed by the mechanical agent in the real-w orld environment.

[0121] In some implementations, as described above, the agent may not include a human being (e.g. it is a robot). Conversely, in some implementations the agent comprises a human user of a digital assistant such as a smart speaker, smart display, or other device. Then the information defining the task can be obtained from the digital assistant, and the digital assistant can be used to instruct the user based on the task.

[0122] For example, a system implementing the core neural network 104 may output to the human user, via the digital assistant, instructions for actions for the user to perform at each of a plurality of time steps. The instructions may for example be generated in the form of natural language (transmitted as sound and / or text on a screen) based on actions chosen by the system. The system chooses the actions such that they contribute to performing a task. A monitoring system (e.g. a video camera system) may be provided forAtorney Docket No. 56113-0862WO1

[0123] monitoring the action (if any) which the user actually performs at each time step, in case (e.g. due to human error) it is different from the action which the system instructed the user to perform. Using the monitoring system the system can determine whether the task has been completed. The system may identify actions which the user performs incorrectly with more than a certain probability7. If so, when the system instructs the user to perform such an identified action, the system may warn the user to be careful. Alternatively or additionally, the system may leam not to instruct the user to perform the identified actions, i.e. ones which the user is likely to perform incorrectly.

[0124] More generally, the digital assistant instructing the user may comprise receiving, at the digital assistant, a request from the user for assistance and determining, in response to the request, a series of tasks for the user to perform, e.g. steps or sub-tasks of an overall task. Then for one or more tasks of the series of tasks, e.g. for each task, e.g. until a final task of the series the digital assistant can be used to output to the user an indication of the task, e.g. step or sub-task, to be performed. This may be done using natural language, e.g. on a display and / or using a speech synthesis subsystem of the digital assistant. Visual, e.g. video, and / or audio observations of the user performing the task may be captured, e.g. using the digital assistant. A system as described may then be used to determine whether the user has successfully achieved the task e.g. step or sub-task, i.e. from the answer as previously described. If there are further tasks to be completed the digital assistant may then, in response, progress to the next task (if any) of the series of tasks, e.g. by¬ outputting an indication of the next task to be performed. Tn this way the user may be led step-by-step through a series of tasks to perform an overall task. During the training of the core neural network, training rewards may be generated e g. from video data representing examples of the overall task (if corpuses of such data are available) or from a simulation of the overall task.

[0125] In a further aspect there is provided a digital assistant device including a system as described above. The digital assistant can also include a user interface to enable a user to request assistance and to output information. In implementations this is a natural language user interface and may comprise a keyboard, voice input-output subsystem, and / or a display. The digital assistant can further include an assistance subsystem configured to determine, in response to the request, a series of tasks for the user to perform. In implementations this may comprise a generative (large) language model, in particular for dialog, e.g. a conversation agent such as Spanow (Glaese et al. arXiv:2209.14375) or Chinchilla (Hoffmann et al. arXiv:2203.15556). The digitalAtorney Docket No. 56113-0862WO1

[0126] assistant can have an observation capture subsystem to capture visual and / or audio observations of the user performing a task; and an interface for the above-described language model neural network (which may be implemented locally or remotely). The digital assistant can also have an assistance control subsystem configured to assist the user. The assistance control subsystem can be configured to perform the steps described above, for one or more tasks e.g. of a series of tasks, e.g. until a final task of the series. More particularly the assistance control subsystem and output to the user an indication of the task to be performed, capture, using the observation capture subsystem, visual or audio observations of the user performing the task, determine from the above-described answer whether the user has successfully achieved the task. In response the digital assistant can progress to a next task of the series of tasks and / or control the digital assistant, e.g. to stop capturing observations.

[0127] As another example, the task can be a genomics task, where the network input is a sequence representing a fragment of a DNA sequence or other molecule sequence and the network output is either an embedding of the fragment for use in a downstream task, e.g., by making use of an unsupervised learning technique on a data set of DNA sequence fragments, or an output for the downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting functional effects of non-coding variants, and so on.

[0128] In some cases, the machine learning task is a combination of multiple individual machine learning tasks, i.e., the core neural network 104 is configured to perform multiple different individual machine learning tasks, e.g., two or more of the machine learning tasks mentioned above. For example, the system can be configured to perform multiple individual natural language understanding tasks, with the network input including an identifier for the individual natural language understanding task to be performed on the netw ork input.

[0129] In some cases, the machine learning task is a multi-modal processing task that requires processing multi-modal data. In general, multi-modal data is a combination of two or more different types of data, e.g., two or more of audio data, image data, text data, or graph data. As one example the multi-modal data may comprise audio-visual data, comprising a combination of pixels of an image or of video and audio data representing values of a digitized audio waveform. As another example the multi-modal data may comprise a combination of i) text data representing text in a natural language and ii) pixels of an image or of video or audio data representing values of an audio waveform.Atorney Docket No. 56113-0862WO1

[0130] Optionally, but not necessarily, the different types of data may represent the same or overlapping objects using the different modalities (types), and when processing multimodal data the data may be mapped into a common embedding space.

[0131] As a particular example, the task is a multi-modal processing task that requires processing both text and image inputs, so that the core neural network 104 includes both a computer vision neural network and a text processing neural network. That is, the target output to be generated by the computer vision neural network for a given image depends on one or more outputs generated by the text processing neural network for one or more corresponding text inputs (and vice versa). Examples of such tasks include openvocabulary image classification, open-vocabulary object detection, image captioning, text-based image search, image-based retrieval, and so on.

[0132] More generally, the multi-modal processing task may correspond to any of the tasks previously described for any of the types of data making up the multi-modal combination. For example, an accuracy of the previously described tasks may be increased when the task is applied to multi-modal data combining the data for which the task has been previously described and another type of data. For example detection or classification of an object or event may be improved when data of multiple different types (modalities) is processed.

[0133] In particular, the core neural network 104 is capable of receiving network inputs and generating network outputs for multiple different machine learning tasks. Generally, two machine learning tasks are different if they have different desired outputs for the inputs received for the tasks. For example, two image classification tasks can be different if the object categories into which each task requires classifying input images are different. As another example, two robot learning tasks can be different if the two tasks require generating outputs defining actions to be performed by a robot to reach two different goals.

[0134] In practice, for any of these examples, the task to be performed by the core neural network 104 can be defined by (at least a part of) the network input, e.g., that is in the form of a prompt or a request, received by the core neural network. In other words, the core neural network 104 will be able to perform any of these tasks when an appropriate prompt or request is received.

[0135] In some situations, the core neural network 104 can be referred to as an autoregressive neural network when the neural network auto-regressively generates an output sequence of tokens as the network output. More specifically, the auto-regressivelyAtorney Docket No. 56113-0862WO1

[0136] generated output is created by generating each particular token in the output sequence conditioned on a current input sequence that includes an input sequence included in the network input and any tokens that precede the particular token in the output sequence, i.e., the tokens that have already been generated for any previous positions in the output sequence that precede the particular position of the particular token.

[0137] As a particular example, the core neural network 104 can have any of a variety of Transformer-based neural network architectures, e.g., encoder-only Transformer architectures, encoder-decoder Transformer architectures, decoder-only Transformer architectures, diffusion Transformer architectures, other attention-based architectures, and so on.

[0138] Examples of such Transformer-based neural network architectures include those described in Colin Raffel, et al., Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 1910.10683, 2019; Daniel Adiwardana, et al., Towards a human-like open-domain chatbot. CoRR, abs / 2001.09977, 2020;

[0139] Aakanksha Chowdhery. et al., PaLM: Scaling Language Modeling with Pathways, arXiv preprint arXiv: 2204.02311; Rohan Anil, et al. Palm 2 technical report. arXiv preprint arXiv:2305.10403, 2023; and Gemini Team, et al., Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805 (2023).

[0140] For example, the core neural network 104 may be a (generative) language mode neural network. Examples of generative language model neural networks include Sparrow (Glaese et al. arXiv:2209.14375), Chinchilla (Hoffmann et al. arXiv:2203.15556), and PaLM 2 (Anil, et al. arXiv:2305: 10403).

[0141] As another example, the core neural network 104 may be a multi-modal model neural network, e.g., a vision language model (VLM) neural network. Examples of multimodal neural networks include Flamingo (Alayrac et al. arXiv: 2204.14198), PaLI (Chen et al. arXiv: 2209.06794), and PaLI-X (Chen et al. arXiv: 2305.18565).

[0142] As another example, the core neural network 104 may be a foundation model neural network. A foundation model neural network is a large-scale machine learning model trained on a broad data set that can be adapted and fine-tuned for a wide variety of applications and downstream tasks. Examples of foundation model neural networks include Imagen (Saharia et al. arXiv: 2205.11487) and Parti (Yu et al. arXiv:

[0143] 2206.10789).Atorney Docket No. 56113-0862WO1

[0144] The persistent memory 106 includes a set of learnable parameters that are independent of input data and are configured to store task-level or system-level information. The persistent memory 106 differs from the neural long-term memory 102 in that the parameters of the persistent memory 106 are not updated as a function of individual input tokens during inference, but instead encode relatively stable abstractions associated with a task, domain, or operating context.

[0145] In some implementations, the persistent memory 106 can include one or more learnable parameter vectors

[0146]

[0147] >

[0148] where Np> 1. The parameter vectors can be incorporated into processing through concatenation with input tokens provided to the core neural network. For example, given an input sequence x and a context window of length N, a modified input sequence can be formed as

[0149]

[0150] >

[0151] where II represents concatenation. The parameters of the persistent memory' 106 thereby occupy predefined positions within the sequence and are processed jointly with input tokens by the core neural network.

[0152] From a memory-system perspective, the persistent memory 106 provides inputindependent storage of task knowledge that complements the adaptive, input-dependent storage provided by the neural long-term memory 102. While the neural long-term memory 102 captures abstractions derived from recent input history, the persistent memory 106 encodes information describing how a task is performed, constraints of the task, or structural regularities that should remain stable across different input sequences.

[0153] From an architectural perspective, the persistent memory 106 operates analogously to data-independent components of a feedforward network. In some neural network architectures, such as Transformer-based models, fully connected layers can be interpreted as implementing attention-like operations with parameters that are independent of the input data. The parameters of the persistent memory' 106 may therefore function as fixed key -value representations that influence attention orAtorney Docket No. 56113-0862WO1

[0154] aggregation operations performed by the core neural network, without being conditioned on the current input tokens.

[0155] Incorporating persistent memory parameters at the beginning of the input sequence can mitigate biases introduced by causal or masked attention mechanisms. In particular, attention mechanisms often exhibit a preference for early positions in a sequence. By introducing learnable, input-independent parameters at those positions, the architecture redistributes attention weights more effectively, improving stability and performance over long contexts.

[0156] During inference, the persistent memory 106 is accessed by the core neural network 104 in the same manner as other sequence elements. Outputs generated by processing the persistent memory parameters may influence representations of subsequent tokens and may be combined with outputs of the neural long-term memory 102 when generating network outputs. In some implementations, the persistent memory parameters can be trained during an offline training phase and remain fixed during inference, while in other implementations the persistent memory parameters can be updated at a slower timescale than the neural long-term memory.

[0157] In some implementations, one or more components of the neural network architecture 100 apply residual connections, gating operations, and normalization operations when processing the input 108 or intermediate token representations. For example, in some implementations the core neural network 104 computes query, key, and value representations using anon-linear activation function, such as a SiLU activation function, and normalizes one or more of the query' and key representations using an ( 2-norm. In some implementations, one or more convolution operations are incorporated into the computation of query, key, and value representations. For example, a onedimensional depthwise-separable convolution layer can be applied after each of the query, key, and value projection operations to incorporate local contextual information in a computationally efficient manner. In some implementations, prior to generating an output projection, one or more normalization operations and gating operations are applied, including gating using a learned linear transformation and application of a nonlinearity, such that the resulting network output 110 is generated from a gated and normalized combination of intermediate representations.

[0158] Generally, deep learning architectures can incorporate memory functionality using different architectural components. For example, attention-based mechanisms in Transformer models implement associative memory' over a limited context windowAtorney Docket No. 56113-0862WO1

[0159] through key and value projection matrices. Due to their reliance on a fixed attention span and their focus on modeling dependencies within a current context window, such mechanisms can be viewed as providing a form of short-term memory.

[0160] In contrast, the neural long-term memory 102 is configured to continuously leam from input data and to store abstractions of previously processed information directly in its weights. As a result, the neural long-term memory 102 can operate as a long-term memory that complements short-term, attention-based mechanisms of the core neural network 104. The neural long-term memory 102 can be incorporated into the core neural network 104 in various architectural configurations, each providing different trade-offs between computational efficiency, memory capacity, and effectiveness when processing long input sequences.

[0161] FIG. 2 is a block diagram of an example memory-as-context neural network architecture 200. The example memory-as-context neural network architecture 200 includes a neural long-term memory 102, core neural network 104, and persistent memory 106 as described above with reference to FIG. 1. In this example memory-as-context neural network architecture 200, the neural long-term memory 102 is treated as a context to current information.

[0162] During operation, the core neural network 104 obtains a network input that is represented as a sequence x e IR'Vxdin, where N represents a sequence length and dinrepresents an input dimensionality. The input 202 is divided into multiple fixed-size segments Sty), i = 1, ... , N / C, where each segment has a fixed size C. The segments can be contiguous, non-overlapping segments of the input sequence, or may be defined using a sliding window, overlapping segmentation, or another deterministic or learned segmentation strategy, provided that each segment has a bounded size C. The operations described below are then performed iteratively for each segment in sequence, such that the neural long-term memory state Mtproduced after processing a segment Sty) is used when processing a subsequent segment S(t + 1).

[0163] For each segment, e.g., segment 202, the core neural network 104 retrieves memory outputs 204 for the segment from the neural long-term memory 102 and retrieves persistent memory parameters 206 from the persistent memory 106. In particular, for a current segment Sty), the neural long-term memory 102 stores a memory state

[0164]

[0165] that encodes abstractions of previously processed segments.

[0166] The core neural network 104 derives a query for the current segment, for exampleAtorney Docket No. 56113-0862WO1

[0167] qt= S(t)WQ,

[0168] where WQis a learned projection associated with the core neural network 104. Using this query, the core neural network 104 retrieves historical information ht(the memory outputs 204) that correspond to to the current segment from the neural long-term memory 102 according to

[0169]

[0170] denotes a retrieval operation of the neural long-term memory 102. In this context, memory retrieval refers to obtaining information from the neural long-term memory 102 using a forward-pass operation without updating the memory weights. During inference, a query is derived from an input, for example by projecting the input using a learned projection matrix WQto obtain qt= xtWQ. The neural long-term memory' 102 processes the query using its current parameters to generate a memory' output yt= M ’ (<7t)- where the memory output represents information stored in the neural long-term memory 102 that is relevant to the query.

[0171] The persistent memory 106 provides a set of input-independent parameters (the persistent memory parameters 206)

[0172] P = [Pl> P2> — >PNP]>

[0173] which encode task-level or system-level knowledge and are retrieved independently of the cunent segment.

[0174] The core neural network 104 applies an attention mechanism 208 to a concatenation 210 of the persistent memory' parameters 206, the segment 202, and the memory' outputs 204 for the segment 202 to obtain an attention output 212. In particular, the core neural network constructs an augmented context sequence

[0175]

[0176] Atorney Docket No. 56113-0862WO1

[0177] where II represents concatenation. The augmented context sequence is processed by an attention mechanism to generate the attention output 212 given by

[0178] yt= Attn(S(t ).

[0179] The attention mechanism jointly attends over persistent memory parameters, retrieved long-term memory outputs, and current segment information to determine which information is relevant for processing and storage.

[0180] Weights of the neural long-term memory' 102 are updated 214 using the attention output. In particular, the neural long-term memory updates its internal parameters based on the attention output via a forward-pass-based update

[0181] Mt = Mt^(yt),

[0182] thereby enabling the neural long-term memory' 102 to store abstractions derived from the current segment. This update occurs without modifying parameters of the core neural network 104.

[0183] The neural long-term memory 102 processes the attention output to retrieve memory outputs for the attention output, for example

[0184]

[0185] The core neural network 104 outputs an output 216 that is a combination of the attention output 212 and a memory' output 218 retrieved from the neural long-term memory 102 using the attention output as a query, e.g., for further processing by subsequent layers in the core neural network 104. In one example, the output for the segment is computed as

[0186]

[0187] where ® denotes a combination operation, such as element-wise multiplication, gating, or another learned fusion operation.Atorney Docket No. 56113-0862WO1

[0188] During inference, parameters associated with the persistent memory 106 remain fixed to preserve task-level knowledge. Parameters of the core neural network 104 operate as in-context learners via attention over the augmented context sequence, while parameters of the neural long-term memory 102 continue to adapt based on the sequence of segments. In this manner, the MAC architecture enables test-time memorization of long-range contextual information while maintaining stable task behavior.

[0189] FIG. 3 is a block diagram of an example gated-memory neural network architecture 300. The example gated-memory neural network architecture 300 includes a neural long-term memory 102, a core neural network 104, and a persistent memory' 106, as described above with reference to FIG. 1. In this architecture, the neural long-term memory 102 is implemented as a gated memory that interacts with the core neural network 104 through a gating mechanism rather than being incorporated as an explicit contextual sequence.

[0190] During operation, the core neural network 104 obtains a network input 302 represented as a sequence x e Dk'Vxd'" where N represents a sequence length and dinrepresents an input dimensionality. The core neural network 104 retrieves persistent memory parameters 304 from the persistent memory 106. The persistent memory parameters include a set of learned, input-independent vectors P —

[0191] [Pi> P2> ■■■ ’ PNP]’ which encode task-level or system-level information.

[0192] The core neural network 104 constructs an augmented input sequence 306 by' concatenating the persistent memory' parameters with the network input, for example

[0193]

[0194] The core neural network 104 applies an attention mechanism 308 to the concatenation of the persistent memory parameters 304 and the network input 302 to obtain an attention output 310. In some implementations, the attention mechanism is a sliding- window attention with prefix, where the persistent memory parameters occupy prefix positions and the sliding window attention acts as a short-term memory' over a bounded context window of the input sequence. For example, the attention output 310 can be given by

[0195] y — SW — Attn*(x),Atorney Docket No. 56113-0862WO1

[0196] where SW — Attn* (•) represents sliding-window atention with prefix.

[0197] In parallel with the atention operation 308, the neural long-term memory 102 is used to process, in sequence, tokens in the concatenation 306 of the persistent memory parameters 304 and the network input 302. As the tokens of x are processed, the neural long-term memory 102 recursively updates its internal weights, thereby accumulating abstractions of the input sequence over time. After processing all tokens of the sequence x, the neural long-term memory 102 produces a final memory state. A memory output 312 derived from the final memory state is used by the core neural network 104 and is denoted herein as M(%).

[0198] The core neural network 104 outputs a combination of the attention output 310 and a memory output 312 derived from the final updated weights of the neural long-term memory' 102. In particular, the atention output y and the neural long-term memory' output M(x) are combined to generate an output 316 using anon-linear gating mechanism 314, for example

[0199] o = y 0 f(x),

[0200] where 0 represents a non-linear gating or fusion operation. In some implementations, the atention output 310 and the neural long-term memory' output 312 are first normalized using learned vector-valued weights and then combined using a non-linearity function (- )•

[0201] Accordingly, FIG. 3 illustrates an architecture in which the atention mechanism provides a short-term memory over recent input tokens, while the neural long-term memory' 102 acts as a fading memory' that continuously updates its parameters based on sequential processing of the input. The gated combination of these two pathways allows the core neural network 104 to balance precise short-term contextual reasoning with longer-term memorization captured in the weights of the neural long-term memory 102.

[0202] During inference, parameters associated with the persistent memory' 106 remain fixed to preserve task-level knowledge. Parameters of the core neural network 104 perform in-context learning via the sliding-window atention mechanism, while parameters of the neural long-term memory 102 continue to adapt based on the inputAtorney Docket No. 56113-0862WO1

[0203] sequence, enabling test-time memorization of long-range information while maintaining stable task behavior.

[0204] FIG. 4 is a block diagram of an example memory-as-layer neural network architecture 400. The example memory-as-layer neural network architecture 400 includes a neural long-term memory 102, a core neural network 104, and a persistent memory 106, as described above with reference to FIG. 1. In this architecture, the neural long-term memory 102 is implemented as a layer within the core neural network 104 and is positioned to process input data prior to application of an attention mechanism.

[0205] The core neural network 104 obtains a network input 402 represented as a sequence x e IRJVxdi’1, where N represents a sequence length and dinrepresents an input dimensionality. The core neural network 104 retrieves persistent memory parameters 404 from the persistent memory 106 and concatenates the persistent memory parameters 404 with the network input 402.

[0206] In particular, the core neural network constructs an augmented input sequence 406 as

[0207] x = [P1p2... pNp] || x,

[0208] where P = [pltp2, ■■■ , pNp] represents a set of learned, input-independent persistent memory parameters encoding task-level or system-level information, and II represents concatenation.

[0209] The neural long-term memory 102 processes, in sequence, tokens in the concatenation 406 of the persistent memory7parameters and the network input to recursively update weights of the neural long-term memory 102. As the tokens of x are processed sequentially, the neural long-term memory 102 updates its internal parameters, thereby7compressing abstractions of both past and current context into its weights. After processing all tokens of the augmented sequence x, the neural long-term memory7102 produces a final memory state and a corresponding memory output 408 given by

[0210]

[0211] where M(-) represents the forward-pass operation of the neural long-term memory 102 after completion of recursive updates.Atorney Docket No. 56113-0862WO1

[0212] The core neural network 104 applies an attention mechanism 410 to the output 408 of the neural long-term memory 102 to obtain an attention output 412. In some implementations, the attention mechanism is a sliding-window attention that operates over the memory output y, thereby providing a bounded short-term context over representations already compressed by the neural long-term memory7. For example, the attention output 412 can be given by

[0213] o = SW — Attn(y),

[0214] where SW — Attn ') represents a sliding-window attention operation.

[0215] Accordingly, in the memory-as-layer architecture 400, the neural long-term memory 102 acts as a sequence modeling layer that precedes the attention mechanism 410 and is responsible for compressing historical and current context into a learned representation before attention is applied. This differs from the memory-as-context architecture of FIG. 2 and the gated-memory architecture of FIG. 3. in that attention operates on the output of the neural long-term memory 102 rather than jointly over raw input tokens and memory outputs.

[0216] In some implementations, the attention mechanism 410 may be omitted entirely, such that the neural long-term memory7102 alone functions as a sequence model without explicit short-term memory. In such variants, the neural long-term memory 102 remains responsible for learning abstractions over long sequences via recursive weight updates, even in the absence of an attention-based short-term memory7.

[0217] During inference, parameters associated with the persistent memory 106 remain fixed to preserve task-level knowledge. Parameters of the neural long-term memory 102 continue to adapt as tokens are processed sequentially, enabling test-time memorization of long-range dependencies, while parameters of the attention mechanism 410, when present, provide localized contextual refinement over the memory-compressed representation. In this manner, the memory-as-layer architecture 400 enables hierarchical processing in which long-term memory precedes and conditions attention-based reasoning.

[0218] While FIG. 4 illustrates an example memory-as-layer architecture 400 in which the neural long-term memory7102 is applied prior to a single attention mechanism, this configuration is not limiting. In general, the neural long-term memory 102 can beAtorney Docket No. 56113-0862WO1

[0219] integrated at arbitrary depths within the core neural network, and the relative ordering of memory layers and attention layers may be varied. For example, in some implementations, the neural long-term memory 102 can be applied after one or more attention layers, between successive attention layers, or multiple times at different depths of the core neural network 104.

[0220] Further, the core neural network 104 may include multiple attention layers, and the neural long-term memory 102 may be interleaved with these layers to enable hierarchical accumulation of long-term information across representational levels. The attention mechanisms used in such architectures are not limited to sliding-window attention, and may include global attention, sparse attention, block-sparse attention, routing-based attention, or hybrid combinations thereof.

[0221] In these architectures, the neural long-term memory 102 continues to operate as a learnable sequence-processing component whose parameters are updated based on sequential token processing, while the attention mechanisms provide complementary short-term or structured context aggregation. By allowing flexible placement and repetition of neural long-term memory layers and attention layers, the architecture supports a broad class of memory-attention compositions that trade off expressiveness, computational cost, and context length without altering the fundamental learning and update principles described herein.

[0222] FIG. 5 is a flow chart of an example process 500 for updating and querying a neural long-term memory during inference. For convenience, the process 500 will be described as being performed by a system of one or more computers located in one or more locations. For example, the neural network architecture 100 of FIG. 1, appropriately programmed, can perform example process 500.

[0223] The system processes, in sequence, each token in a sequence of tokens (step 502). As part of this sequential processing, the system receives a token from the sequence of tokens (step 504). The received token is then projected using a first linear layer to obtain a key for the token (step 506). The token is further projected using a second linear layer to obtain a value for the token (step 508). In some implementations, the first linear layer and the second linear layer include learned parameters of a core neural network.

[0224] The system then processes the key for the token using a neural long-term memory to generate a memory' output that corresponds to the key for the token (step 510). In some implementations, the neural long-term memory includes a single-layer perceptron. In some implementations, processing the key using the neural long-term memory includesAtorney Docket No. 56113-0862WO1

[0225] multiplying a weight matrix for the neural long-term memory by the key. In other implementations, the neural long-term memory is a multi-layer perceptron with two or more layers. In such implementations, one or more layers may apply an expansion factor, and one or more layers may apply a Gaussian Error Linear Unit activation function. Different layers in the neural long-term memory may be assigned respective learning parameters for updating weights of the neural long-term memory, including learning rates that are scaled such that increasingly deeper layers have slower learning rates. In some implementations, the learning parameters are input-dependent parameters determined using a linear projection of the input.

[0226] After generating the memory output corresponding to the key, the system computes a gradient of a loss with respect to weights of the neural long-term memory (step 512). The loss is dependent on a difference between the memory output that corresponds to the key for the token and the value for the token. In some implementations, computing the gradient of the loss includes multiplying a difference between (i) a product of a weight matrix for the neural long-term memory and the token and (ii) the token by a transpose of the token.

[0227] In some implementations, the system further computes a surprise metric for the token. The surprise metric includes a combination of a surprise metric for a previous token in the sequence of tokens and the computed gradient of the loss. The system then updates weights of the neural long-term memory using the gradient, without updating parameters of a core neural network (step 514). In some implementations, updating the weights of the neural long-term memory' includes weighting current weights of the neural long-term memory using a gating mechanism. In some implementations, updating the weights is performed without updating parameters of the first linear layer and the second linear layer. In implementations that compute a surprise metric, the weights of the neural long-term memory are updated using the surprise metric for the token.

[0228] In some implementations, the sequence of tokens is processed in chunks (e.g., implementations where the neural long-term memory is implemented as memory as a context, as described above with reference to FIG. 2).

[0229] In some implementations corresponding to a memory-as-context architecture, as described above with reference to FIG. 2, the sequence of tokens processed by the neural long-term memory is derived from attention outputs generated over segmented portions of a network input.Atorney Docket No. 56113-0862WO1

[0230] In such implementations, the system computes gradients of the loss in parallel for each token in a chunk of tokens included in the sequence of tokens. In some implementations, computing the gradients of the loss in parallel comprises performing a batched matrix multiplication operation over the tokens in the chunk. The system then applies the computed gradients sequentially to update the weights of the neural long-term memory. In some implementations, updating the weights comprises computing surprise metrics for the tokens in the chunk using a parallel scan operation and storing intermediate update terms for each chunk.

[0231] After or during updating of the neural long-term memory. the system processes a respective query for each of one or more of the tokens using the neural long-term memoiy to generate a respective memory output for each query (step 516). In some implementations, each token in the sequence of tokens is processed as a respective query using the neural long-term memory to generate a respective memory output for each token in the sequence of tokens. In some implementations, processing the respective query is performed without updating weights of the neural long-term memory'.

[0232] The system then processes an input comprising the respective memory outputs using the core neural network to generate one or more network outputs (step 518). In some implementations, the input comprises a combination of an attention output and the memory outputs.

[0233] In some implementations, the sequence of tokens processed by the neural longterm memory' comprises an attention output generated by the core neural network.

[0234] In some implementations corresponding to a memory-as-context architecture, as described above with reference to FIG. 2, the attention output is generated by attending over a concatenation of persistent memory parameters, a current input segment, and memory outputs retrieved from the neural long-term memory.

[0235] In these implementations, the system obtains, by the core neural network, a network input and divides, by the core neural network, the network input into a plurality of fixed-size segments. For each segment, the system retrieves, by the core neural network, memoiy outputs for the segment from the neural long-term memory and retrieves persistent memory parameters from a persistent memory7. The system applies, by the core neural network, an attention mechanism to a concatenation of the persistent memory7parameters, the segment, and the memory7outputs for the segment to obtain the sequence of tokens. The system then processes the sequence of tokens using the neural long-term memory to update the weights of the neural long-term memory7.Atorney Docket No. 56113-0862WO1

[0236] In some implementations, retrieving memory' outputs for the segment from the neural long-term memory comprises generating a query using the segment by multiplying the segment by a weight matrix of learned parameters, and processing the query using the neural long-term memory to retrieve the memory outputs for the segment without updating weights of the neural long-term memory. In some implementations, processing the attention output using the updated neural long-term memory to retrieve memory¬ outputs for the attention output comprises processing the attention output using the neural long-term memory without updating weights of the neural long-term memory-.

[0237] In some implementations, the sequence of tokens comprises a concatenation of persistent memory parameters and a network input.

[0238] In some implementations corresponding to a gated-memory architecture, as described above with reference to FIG. 3, the neural long-term memory processes the concatenated sequence to recursively update its weights while the core neural network applies an attention mechanism to the same sequence to generate an attention output in parallel.

[0239] In these implementations, the system obtains, by the core neural network, the network input and retrieves, by the core neural network, the persistent memory parameters from a persistent memory-. The system applies, by the core neural network, an attention mechanism to the sequence of tokens to obtain an attention output, and processes the sequence of tokens using the neural long-term memory to recursively update weights of the neural long-term memory. The input comprising the respective memory outputs then comprises a combination of the attention output and the respective memory- outputs. In some implementations, the attention mechanism comprises a sliding window attention with prefix, wherein the sliding window attention acts as a short-term memory. In some implementations, the combination of the attention output and final updated weights of the neural long-term memory comprises a non-linear gating of the attention output and the final updated weights of the neural long-term memory-. In some implementations, the system normalizes the attention output and the final updated weights of the neural long-term memory using learned vector-valued weights and a non-linearity function.

[0240] In some implementations, the sequence of tokens comprises a concatenation of persistent memory parameters and a network input, and the system processes the sequence of tokens using the neural long-term memory to recursively update weights of the neural long-term memory.Atorney Docket No. 56113-0862WO1

[0241] In some implementations corresponding to a memory-as-layer architecture, as described above with reference to FIG. 4, the neural long-term memory operates as a sequence-processing layer that produces a final updated memory state prior to application of an attention mechanism by the core neural network.

[0242] The system then applies, by the core neural network, an attention mechanism to a final updated state of the neural long-term memory to obtain an attention output. In some implementations, the attention mechanism comprises a sliding window attention.

[0243] During training of the neural network architecture, the system leams parameters that enable the memory to store and retrieve abstractions of previously processed information.

[0244] In general, training a neural long-term memory involves repeatedly processing sequences of tokens, computing a loss that reflects how well the memory reproduces or predicts target representations associated with the tokens, and updating weights of the neural longterm memory based on gradients of the loss. Unlike conventional neural network layers whose parameters are typically updated only during offline training, the neural long-term memory described herein is configured to support gradient-based updates that can also be performed during inference, enabling continual learning and memorization over long input sequences.

[0245] FIG. 6 illustrates an example process 600 for training and updating a neural longterm memory using parallelized computation. As shown in FIG. 6. the system processes a sequence of input tokens and updates weights of the neural long-term memory based on the tokens, while parameters of a core neural network remain fixed. The illustrated process decomposes memory training into a set of operations that can be executed efficiently using parallel hardware, such as GPUs or TPUs, while preserving the functional behavior of sequential gradient-based updates.

[0246] In the example of FIG. 6, the sequence of input tokens is divided into one or more contiguous chunks, each chunk including multiple tokens. For each chunk, the system computes representations of the tokens, including keys and values derived from the tokens using learned projection layers. Based on these representations, the system computes an associative memory loss for each token that measures a discrepancy between a memory' output produced by the neural long-term memory and a target representation associated with the token. Operations performed independently within each chunk correspond to a linear, within-chunk computation stage 602, in which gradient contributions for tokens of the chunk are computed in parallel.Atorney Docket No. 56113-0862WO1

[0247] Gradients of the associative memory loss with respect to weights of the neural long-term memory are computed for tokens within a chunk in parallel, for example using batched matrix multiplication operations. These parallel gradient computations produce per-token update terms that reflect how each token contributes to updating the neural long-term memory. The resulting per-token update terms are then provided to a crosschunk update stage 604, which integrates update information across successive portions of the sequence.

[0248] To preserve temporal dependencies across the sequence, the system accumulates the per-token update terms using a recurrence relation that propagates update information across tokens. This accumulation is illustrated in FIG. 6 as a momentum computation stage 606, which aggregates update contributions over time. In some implementations, the momentum computation includes a parallel associative summation path 606a that computes accumulated update states for tokens within a chunk using a scan or summation operation. In other implementations, or in addition, the momentum computation includes a global kernel path 606b that implements a convolution-like or linear time-invariant update over a chunk or across chunks.

[0249] The accumulated update states are then applied sequentially7across chunks to update the weights of the neural long-term memory. This sequential application ensures that the neural long-term memory evolves consistently with the order of tokens in the original sequence, while the parallel computation within each chunk reduces overall latency and improves throughput. The application of accumulated updates is further conditioned by a weight-decay control stage 608. In one implementation, illustrated as 608a, the update is performed without weight decay, such that previously stored memory content is preserved while new update terms are added. In another implementation, illustrated as 608b, the update includes a decay term that attenuates previously stored memory content prior to incorporating new update terms, thereby7implementing controlled forgetting.

[0250] In the example illustrated in FIG. 6, the weight-decay control stage 608 operates on update terms expressed in matrix form. In these expressions, Worepresents an initial or prior weight matrix of the neural long-term memory at the start of processing a chunk, and X represents a matrix formed from token representations within the chunk, for example by stacking token vectors as columns or rows. The term (V1 / 0X — X)XTAtorney Docket No. 56113-0862WO1

[0251] represents an aggregated gradient contribution derived from discrepancies between memory outputs and corresponding target representations for tokens in the chunk.

[0252] In implementations that apply weight decay, as illustrated as stage 608b, the aggregated gradient contribution is further scaled by one or more decay-related factors. In particular, Qbrepresents a chunk-level learning-rate or update-scaling matrix, and Bbrepresents a chunk-level decay or retention matrix that controls how previously stored memory content is attenuated when incorporating new updates. By contrast, in implementations without weight decay, illustrated as stage 608a, the aggregated gradient contribution is applied without scaling by decay-related factors, thereby preserving previously stored memory content while incorporating new update terms.

[0253] These formulations allow gradient-based updates, with or without decay, to be expressed using matrix multiplication operations, enabling efficient implementation on parallel hardware while maintaining the temporal semantics of sequential memory updates.

[0254] The update of the neural long-term memory is performed independently of updates to the core neural network. During this process, parameters of the core neural network and any persistent memory remain fixed, while only the weights of the neural long-term memory are modified. This separation allows the neural long-term memory to adapt dynamically to observed input data without destabilizing task-level behavior encoded in the core neural network.

[0255] Accordingly, FIG. 6 illustrates a training and update mechanism in which gradient-based learning of a neural long-term memory is reformulated to support parallel computation across tokens and chunks, including linear within-chunk computation 602, non-linear cross-chunk integration 604, momentum accumulation via parallel summation and global kernel paths 606a, 606b, and optional weight-decay control 608a, 608b. This formulation enables efficient execution on modem parallel hardware while maintaining the semantics of sequential memory updates. In some implementations, this training process is performed during an offline training phase, while in other implementations the same process is performed during inference or test-time execution to enable continual memorization of long-range contextual information.

[0256] This specification uses the term “configured” in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause theAtorney Docket No. 56113-0862WO1

[0257] system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions.

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

[0259] The term “data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0260] A computer program, which may also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it can be deployed in any form, including as a stand alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds otherAtorney Docket No. 56113-0862WO1

[0261] programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g.. files that store one or more modules, sub programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communication network.

[0262] In this specification, the term “database” is used broadly to refer to any collection of data: the data does not need to be structured in any particular way, or structured at all, and it can be stored on storage devices in one or more locations. Thus, for example, the index database can include multiple collections of data, each of which may be organized and accessed differently.

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

[0264] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0265] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read only memory or a random access memory or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., aAtorney Docket No. 56113-0862WO1

[0266] mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

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

[0268] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g.. visual feedback, auditory’ feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web brow ser on a user's device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.

[0269] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardw are accelerator units for processing common and compute-intensive parts of machine learning training or production, i.e., inference, workloads.

[0270] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework or a JAX framework.

[0271] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front end component, e.g., a client computer having a graphical user interface, a web brow ser, or an app through which a user can interact with an implementation of theAtorney Docket No. 56113-0862WO1

[0272] subj ect matter described in this specification, or any combination of one or more such back end, middleware, or front end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.

[0273] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at the server from the device.

[0274] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.

[0275] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally beAttorney Docket No. 56113-0862WO1

[0276] integrated together in a single software product or packaged into multiple software products.

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

[0278] What is claimed is:

Claims

Attorney Docket No. 56113-0862WO1CLAIMS1. A computer implemented method comprising:processing, in sequence, each token in a sequence of tokens, the processing comprising:receiving a token from the sequence of tokens;projecting the token using a first linear layer to obtain a key for the token; projecting the token using a second linear layer to obtain a value for the token;processing the key for the token using a neural long-term memory to generate a memory output that corresponds to the key for the token;computing a gradient of a loss with respect to weights of the neural longterm memory, wherein the loss is dependent on a difference between the memory output that corresponds to the key for the token and the value for the token; andupdating weights of the neural long-term memory using the gradient without updating parameters of a core neural network;processing a respective query7for each of one or more of the tokens using the neural long-term memory to generate a respective memory output for each query; and processing an input comprising the respective memory outputs using the core neural network to generate one or more network outputs.

2. The method of claim 1, further comprising computing a surprise metric for the token, the surprise metric comprising a combination of a surprise metric for a previous token in the sequence of tokens and the computed gradient of the loss, wherein the weights of the neural long-term memory7are updated using the surprise metric for the token.

3. The method of claim 1 or claim 2, wherein updating the weights of the neural long-term memory^ using the gradient comprises weighting current weights of the neural long-term memory using a gating mechanism.

4. The method of any one of claims 1 to 3, wherein updating weights of the neural long-term memory using the gradient further comprises updating weights of the neuralAttorney Docket No. 56113-0862WO1long-term memory without updating parameters of the first linear layer and the second linear layer.

5. The method of any one of claims 1 to 4, wherein the neural long-term memory comprises a single-layer perceptron.

6. The method of claim 5, wherein computing the gradient of the loss comprises multiplying a difference between i) a product of a weight matrix for the neural long-term memory and the token and ii) the token by a transpose of the token.

7. The method of claim 5, wherein processing the key using the neural long-term memory to generate the memory outputs that corresponds to the key comprises multiplying a weight matrix for the neural long-term memory by the key.

8. The method of any one of claims 1 to 7, further comprising:computing gradients of the loss in parallel for each token in a chunk of tokens included in the sequence of tokens; andapplying the computed gradients sequentially to update the weights of the neural long-term memory.

9. The method of claim 8, wherein computing the gradients of the loss in parallel for each token in the chunk comprises performing a batched matrix multiplication operation over the tokens in the chunk.

10. The method of claim 8, wherein updating weights of the neural long-term memory comprises computing surprise metrics for the tokens in the chunk using a parallel scan operation and storing intermediate update terms for each chunk.

11. The method of any one of claims 1 to 10, wherein the first linear layer and the second linear layer comprise learned parameters of the core neural network.

12. The method of any one of claims 1 to 11, wherein the sequence of tokens comprises an attention output generated by the core neural network and the method further comprises:Attorney Docket No. 56113-0862WO1obtaining, by the core neural network, a network input;dividing, by the core neural network, the network input into a plurality of fixed-size segments;for each segment in the plurality of fixed-size segments:retrieving, by the core neural network, memory' outputs for the segment from the neural long-term memory and retrieve persistent memory parameters from a persistent memory:applying, by the core neural network, an attention mechanism to a concatenation of the persistent memory' parameters, the segment , and the memory outputs for the segment to obtain the sequence of tokens; andprocessing the sequence of tokens using the neural long-term memory to update the weights of the neural long-term memory.

13. The method of claim 12, wherein:processing a respective query’ for each of one or more of the tokens using the neural long-term memory to generate a respective memory output for each query comprises processing each token in the sequence of tokens as a respective query using the neural long-term memory' to generate a respective memory output for each token in the sequence of tokens; andthe input comprising the respective memory outputs comprises a combination of the attention output and the memory outputs.

14. The method of claim 12, wherein retrieving memory outputs for the segment from the neural long-term memory’ comprises:generating a query using the segment, comprising multiplying the segment by a yveight matrix of learned parameters; andprocessing the query using the neural long-term memory' to retrieve the memory' outputs for the segment yvithout updating weights of the neural long-term memory.

15. The method of any one of claims 12-14, yvherein processing the attention output using the updated neural long-term memory to retrieve the memory outputs for the attention output comprises processing the attention output using the long-term memory without updating weights of the neural long-term memory'.Atorney Docket No. 56113-0862WO116. The method of any one of claims 1 to 11, wherein the sequence of tokens comprises a concatenation of persistent memory parameters and a network input, and wherein the method further comprises:obtaining, by the core neural network, the network input;retrieving, by the core neural network, the persistent memory7parameters from a persistent memory;applying, by the core neural network, an attention mechanism to the sequence of tokens to obtain an attention output; andprocessing the sequence of tokens using the neural long-term memory to recursively update weights of the neural long-term memory,wherein the input comprising the respective memory outputs comprises a combination of the attention output and the respective memory outputs.

17. The method of claim 16, wherein the attention mechanism comprises a sliding window attention with prefix, wherein the sliding window attention acts as a short-term memory.

18. The method of claim 16 or claim 17, wherein the combination of the attention output and final updated weights of the neural long-term memory comprises a non-linear gating of the attention output and the final updated weights of the neural long-term memory.

19. The method of any one of claims 16 to 18, further comprising normalizing, by the core neural network, the attention output and the final updated weights of the neural longterm memory using learned vector-valued weights and a non-linearity7function.

20. The method of any one of claims 1 to 11, wherein the sequence of tokens comprises a concatenation of persistent memory parameters and a network input, and wherein the method further comprises:obtaining, by the core neural network, the network input;retrieving, by the core neural netw ork, the persistent memory parameters from a persistent memory;processing the sequence of tokens using the neural long-term memory to recursively update weights of the neural long-term memory; andAttorney Docket No. 56113-0862WO1applying, by the core neural network, an attention mechanism to a final updated state of the neural long-term memory to obtain an attention output,wherein the input comprising the respective memory outputs comprises the attention output.

21. The method of claim 20, wherein the attention mechanism comprises a sliding window attention.

22. The method of claim 1 , wherein the neural long-term memory' comprises a multilayer perceptron with two or more layers.

23. The method of claim 22, wherein one or more of the layers apply an expansion factor.

24. The method of claim 22, wherein one or more of the layers applies a Gaussian Error Linear Unit activation function.

25. The method of claim 22, wherein different layers in the neural long-term memory are assigned respective learning parameters for updating weights of the neural long-term memory7.

26. The method of claim 25, wherein the learning parameters comprise a learning rate, wherein the learning rates for the two or more layers are scaled such that increasingly deeper layers in the neural long-term memory have slower learning rates.

27. The method of claim 25, wherein the learning parameters comprise inputdependent parameters determined using a linear projection of the input.

28. A system comprising one or more computers and one or more storage devices storing instructions that when executed by7the one or more computers cause the one or more computers to perform operations according to the method of any one of the preceding claims.Atorney Docket No. 56113-0862WO129. One or more computer storage media storing instructions that when executed by one or more computers cause the one more computers to perform operations according to the method of any one of the preceding claims.

Images