Latent thought chain for machine learning models

Abstract

A computer-implemented method for generating a response to a query. The method comprises receiving one or more query tokens, the one or more query tokens indicative of the query, providing the one or more query tokens as input to a machine learning model, outputting, from a first head of the machine learning model, a first embedding based upon the one or more query tokens, generating an intermediate input embedding based upon the one or more query tokens and the first embedding, outputting, from a second head of the machine learning model, output data based upon the intermediate input embedding, and generating the response to the query based upon the output data.

Description

LATENT THOUGHT CHAIN FOR MACHINE LEARNING MODELSBACKGROUND

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

[0002] Machine learning models receive an input and generate an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models and generate the output based on the received input and on values of the parameters of the model.

[0003] Some machine learning models are deep models that employ multiple layers of models to generate an output for a received input. For example, a deep neural network is a deep machine learning model that includes an output layer and one or more hidden layers that each apply a non-linear transformation to a received input to generate an output.SUMMARY

[0004] This specification describes systems and methods implemented as computer programs on one or more computers in one or more locations for generating a response to a query using a machine learning model.

[0005] According to a first aspect there is provided a computer-implemented method for generating a response to a query'. The method comprises receiving one or more query tokens, the one or more query tokens indicative of the query, providing the one or more query tokens as input to a machine learning model, outputting, from a first head of the machine learning model, a first embedding based upon the one or more query tokens, generating an intermediate input embedding based upon the one or more query tokens and the first embedding, outputting, from a second head of the machine learning model, output data based upon the intermediate input embedding, and generating the response to the query based upon the output data. One or more of the query tokens represent image data, video data, or audio data.

[0006] According to a second aspect there is provided a computer-implemented method of generating a neural network output. The method comprises (i) obtaining an input sequence of tokens selected from a predetermined vocabulary of tokens, the input sequence of tokens representing a sy stem input, (ii) at one or more first processing iterations generating a first portion of an output sequence of tokens by processing the input sequence of tokens and a current output sequence of tokens, using a neural network, to generate a neural network output for selecting a next output token to extend the current output sequence of tokens, processing the neural network output using a token generation neural network to generate a soft token that is not part of the predetermined vocabulary of tokens, the soft token comprising one or morecontinuous values, and extending the current output sequence of tokens by appending the soft token to the current output sequence of tokens, (iii) at one or more second processing iterations generating a second portion of the output sequence of tokens by processing the input sequence of tokens and the current output sequence of tokens including the soft token, using the neural network, to generate the neural network output for selecting a next output token from the predetermined vocabulary of tokens to further extend the current output sequence and selecting the next output token from the predetermined vocabulary of tokens using the neural network output, and (iv) at a final iteration obtaining a system output based upon at least the second portion of the output sequence of tokens. One or more of the input sequence of tokens represents image data, video data, or audio data.

[0007] There is also provided a computing system comprising one or more processors and one or more non-transitory computer-readable media storing computer-readable instructions configured to cause one or more processors to perform a method according to the first and / or second aspect.

[0008] There is also provided one or more non-transitory computer-readable media storing computer-readable instructions configured to cause one or more computing devices to perform a method according to the first and / or second aspect.

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

[0010] Implementations described herein enable machine learning models, during both training and inference, to reason in a latent space (i.e. a vector space comprising one or more continuous values) distinct from a discrete token space (i.e. the predetermined vocabulary of tokens). In this way, implementations significantly decrease computational complexity and inference time by eliminating the need to generate discrete tokens explicitly indicating rationale (e.g. in words or natural language), while enabling machine learning models to leam abstract and task-agnostic representations in a self-supervised manner - i.e. without ground truth data explicitly indicating such representations. These abstract and task-agnostic representations are advantageous for a wide variety' of commonly performed tasks (e.g. many real-world tasks require abstract reasoning capabilities) and may be suitable for training the machine learning model for multi-task learning. Such abstract and task-agnostic representations provide a compressed and efficient representation of rationale for the response. Implementations described herein significantly outperform both standard fine-tuning methods and state-of-the- art baselines across six major benchmarks, covering tasks that include arithmetic reasoning,question-answering, and complex spatial reasoning tasks (e.g. tasks occurring in the real-world requiring physical interaction between objects).

[0011] As described above, the described methods and systems can be used to enhance the training and inference of machine learning models. As described herein, query data (i.e. query tokens) and output data may take various forms, including images, video, audio and text. It will be appreciated that such methods and systems enable a wide range of analysis and data processing tasks to be completed.

[0012] The techniques described herein are based upon at least the consideration that conventional machine learning models (e.g. autoregressive models), which are often required to generate a significant number of tokens (e.g. using chain-of-thought via natural language) to generate responses to a query, are computationally inefficient and ineffective for many real- world and / or real-time applications. For example, rationale expressed in natural language requires a relatively large number of sequential forward passes through the machine learning model, in addition to being ineffective at expressing abstract concepts (e.g. relative spatial positions of objects). By generating a compressed rationale in a latent space, the techniques described herein advantageously provide significantly reduced latency and / or memory bandwidth consumption, making the techniques suitable for deployment on resource- constrained hardware or in applications requiring rapid responses, such as the control of a device operating in the real-world.

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

[0014] FIG.1 depicts an example system for generating a response to a query in accordance with the techniques described herein.

[0015] FIG.2 depicts an example of a machine learning model implementing Chain-of- Thought reasoning for generating a response to a query.

[0016] FIG.3 depicts an example of a machine learning model implementing Latent Thought Chain reasoning for generating a response to a query in accordance with the techniques discussed herein.

[0017] FIG.4A depicts an example predetermined vocabulary of tokens.

[0018] FIG.4B depicts a first set of discrete tokens for generating a response to a query.

[0019] FIG.4C depicts a second set of tokens including a plurality of discrete tokens and a soft token for generating a response to a query.

[0020] FIG.4D depicts a space of discrete values defined by a predetermined vocabulary of tokens.

[0021] FIG.4E depicts a latent embedding space of continuous values.

[0022] FIG.5A depicts a Gaussian distribution for sampling a first embedding.

[0023] FIG.5B depicts a Gaussian head of a machine learning model configured to sample a first embedding from a Gaussian distribution.

[0024] FIG.6A depicts an example training algorithm for training a machine learning model in accordance with the techniques described herein.

[0025] FIG.6B depicts a system for training a machine learning model in accordance with the techniques described herein.

[0026] FIG.7 A depicts a first set of experimental results for a machine learning model trained using a plurality7of different methods.

[0027] FIG.7B depicts a second set of experimental results for a machine learning model trained using a plurality of different methods.

[0028] FIG.8 depicts a third set of experimental results for a machine learning model trained with and without Latent Thought Chain reasoning.

[0029] FIG.9 depicts a fourth set of experimental results for a machine learning model trained with Latent Thought Chain reasoning.

[0030] FIG.10 depicts a fifth set of experimental results for a machine learning model trained with Latent Thought Chain reasoning.

[0031] FIG. 11 depicts a flow diagram of a method for generating a response to a query7.

[0032] FIG. 12 depicts a flow diagram of a method for generating a neural network output.

[0033] Like reference numbers and designations in the various drawings indicate like elements.

[0034] Machine learning models may be trained to generate a response to a query . The query may comprise one or more tokens, referred to herein as “query tokens" (e.g. a word or a portion of a word). The response may also comprise one or more tokens. Such machine learning models are typically autoregressive, i.e. they are trained to predict a next output token (e.g. a word or a portion of a word) based upon preceding tokens, i.e. the query7tokens and a current output sequence of tokens. Some of the query7tokens and the output tokens may be selected from a predetermined vocabulary7of tokens (e.g. a list of all possible words or portions of words each corresponding to a numeric value) and appended to the current output sequence of tokensfor generation of further next output tokens. The current output sequence of tokens may be provided, at a final iteration, as the response to the query.

[0035] Some machine learning models can generate rationale as part of their response. The rationale could include text, images, video, audio, etc., to be used by the machine learning model as a logical and / or statistical basis for generating the response. For example, the querymay be ‘"Please let me know whether the following movie review is positive or negative and explain your reasoning: [MOVIE REVIEW]” and the response (i.e. the cunent output sequence of tokens) may be '‘The movie review was negative because the reviewer described the narrative of the movie as subpar and unoriginal in the final paragraph”. Some particular machine learning models may be trained using chain-of-thought (“CoT”) or tree-of-thought (“ToT”) reasoning to enhance their performance. Machine learning models trained to use CoT or ToT may be trained specifically to generate tokens that indicate rationale for the response such that proceeding next tokens take into account the rationale (i.e. some preceding tokens), thus improving the prediction for the next token. For example, the query (or "‘prompt”) may indicate “What is the answer to the question 33 x 2? Explain your reasoning step by step” and the response may indicate “The answer to your question may be decomposed into two separate steps: 3 x 2 = 6; and 30 x 2 = 60. The answer is therefore 6 + 60 = 66”. Machine learning models may leverage CoT and ToT training methods to improve their reasoning capabilities by training the model (e.g. using supervised fine-tuning) to break a problem down into a longer sequence of simpler steps.

[0036] However, there are many substantial drawbacks when training a model to explicitly generate rationale as part of their response. Generating training datasets (i.e. for supervised fine-tuning) that include optimal rationale for the response requires substantial data curation and can be both economically and computationally expensive to generate. Generating responses (e.g. an output sequence of tokens) that include a large number of tokens indicative of the rationale for the response using autoregressive decoding can be both computationally and time inefficient. That is, large-scale language models require significant processing power to generate each next token for the output sequence. As the number of output tokens increases so does the computational load and energy consumption. In addition, the overall time to generate the final response to the query increases with every iteration, which can lead to a poor user experience and / or act as a bottleneck for other systems that may require the response to the query for further processing.

[0037] For complex tasks, articulating (i.e. outputting one or more next tokens) the rationale for generating the response may not always be possible, e.g. where spatial reasoningis required. For example, the query could be ‘‘Imagine a cube sitting on a flat surface. If you paint all six sides of the cube and then cut the cube into 27 smaller, equal-sized cubes (like a 3x3x3 grid), how many of the smaller cubes will have exactly two sides painted?”. In another example, the query could include an image depicting a scene with elements in the foreground and elements in the background. In this example, the response could be for determining which elements are in the background or foreground, thus requiring spatial reasoning capabilities inferred based upon low-level features from the image. Generating rationale for the response in these cases may be impossible by only sampling from the predetermined vocabulary of tokens. As an additional drawback, machine learning models trained to explicitly generate rationale as part of their response may lack the fundamental ability to plan the response prior to generation. That is, machine learning models that generate a next output token at each iteration, e.g. autoregressively, may only be concerned with the next predicted token at any particular iteration, rather than considering subsequent tokens as well (i.e. the response looking forward) or a general structure for the response.

[0038] FIG. l depicts an example system 1 for generating a response 150 to a query 100. As will become readily apparent, generating the response 150 to the query 100 may use a machine learning model 1 10. In some examples, the machine learning model 110 is a neural network such as a Transformer-based neural network comprising one or more attention layers 116. A method for generating the response 150 to the query- 100 is described by way of example below, e.g. using the example system 1.

[0039] The method comprises receiving one or more query tokens 102. In such implementations, the one or more query- tokens 102 are indicative of the query- 100. In some examples, one or more of the query tokens 102 are selected from a predetermined vocabulary of tokens. Specific detail regarding the predetermined vocabulary of tokens is provided below with reference to FIG.4. In some examples, the predetermined vocabulary of tokens is finite. The query tokens 102, as described herein, may- represent text, e.g., words, wordpieces, characters, etc., in a natural or computer language. That is, the query- may be received in the form of text, e.g.. as a series of encoded characters, e.g. UTF-8 encoded characters, and the text may be tokenized according to the predetermined vocabulary of tokens, e.g. using a tokenizer layer 112 of the machine learning model 110, into one or more tokens 102 representative of portions of the text. However, in other examples, the tokens 102 may represent image, video, audio (i.e. a whole or portion of an image, video, and / or audio) or other suitable modality . In a specific example, the query 100 may be “What is the answer to 33 x 2?” and the one or more query tokens 102 may be representative of “What”, “is”, “the”, “answer”, “to”, “33”, “x”, “2”,and “?”. A query token 102 may, in some examples, be a numeric value corresponding to one of those elements of the query 100, the numeric value being suitable for processing using the machine learning model 110. That is, the predetermined vocabulary of tokens in this example may be {'‘What”: 6304, “is”: 1873, “the”: 1531, “answer”: 9431, “to”: 2892, “33”: 4613, “x”: 5852, “2”: 4554, “?”: 5008, ... } where the query tokens 102 are 6304, 1873, 1531, 9431, 2892, 4613, 5852, 4554, 5008. The predetermined vocabulary of tokens may define every possible token 102 for representing the query 100. and may in some examples include tokens for indicating portions of the query 100 that do not appear in the predetermined vocabulary’, e g. an unknow n token corresponding to [UNK], tokens for indicating the start of the query' 100, e.g. a start-of-string token representing [SOS], and tokens for indicating the end of the query 100, e.g. an end-of-string token representing [EOS],

[0040] The method further comprises providing the one or more query tokens 102 as input to the machine learning model 110. In some examples, the one or more query tokens 102 are provided as input to an embedding layer 114 of the machine learning model 110. For example, where the query tokens 102 are 6304, 1873. 1531, 9431, 2892, 4613. 5852, 4554, 5008, these numeric values may be provided as input to an embedding layer 114 of the machine learning model 110 for further processing, e g. to generate one or more query' embeddings, as described below . In another example, each query' token 102 may be a one-hot vector encoding e.g. the particular word, word piece, image, etc. which the query token represents. That is, an example query token 102 [0, 0, 0, 1. 0, 0, . . . ] may represent a word associated with a fourth index of the predetermined vocabulary of tokens.

[0041] The method further comprises outputting, from a first head 118 of the machine learning model 110, a first embedding 130 based upon the one or more query tokens 102. In some examples, the first embedding 130 may comprise one or more continuous numeric values. For example, the first embedding 130 may be an n-dimensional vector such as [-0.8384, 0.7654, 0.4273, -0.3692, 0.2321, ... ]. Outputting from the first head 118 may occur in a first forward pass of the machine learning model 110, i.e. a first iteration of providing input to the machine learning model. The first head 118 may be a latent output layer of the machine learning model 110. i.e. an output layer configured to generate a latent or intermediate output which is not a final prediction for generating the response 150 to the query' 100. In some examples, the first head 118 is connected to a final attention layer 116 of the machine learning model 110, as described below.

[0042] As described herein, a “head” of the machine learning model 110 is any part of the machine learning model 110 that is configured to generate an output associated with a task-specific purpose. For example, as will become apparent, the first embedding 130 output by the first head 118 has a task-specific purpose associated with the internal functioning of the machine learning model 110 described herein. A head of the machine learning model 1 10 may comprise a final layer of the machine learning model 110. A head of the machine learning model 110 is generally connected to preceding layers and is at least configured to generate an output for the task-specific purpose.

[0043] The method further comprises generating an intermediate input embedding 104 based upon the one or more query tokens 102 and the first embedding 130. In some examples, the intermediate input embedding 104 may be a latent embedding generated by processing the one or more query tokens 102 and the first embedding 130. For example, the intermediate input embedding 104 may be generated based upon the query tokens 102: 6304. 1873, 1531, 9431, 2892, 4613, 5852, 4554, 5008 and the first embedding 130: [-0.8384, 0.7654, 0.4273, -0.3692, 0.2321, ... ]. In this example, the intermediate input embedding 104 may also comprise one or more continuous numeric values. That is, the intermediate input embedding 104 that is generated based upon the query tokens 102 and the first embedding 130 may also be an n- dimensional vector such as [[-0.1003. 0.2410, 0.2310, -0.0612, 0.1812. -0.0606, 0.1807 ... J. ... ].

[0044] The method further comprises outputting, from a second head 120 of the machine learning model, output data 140 based upon the intermediate input embedding 104. In some examples, the machine learning model 110 may process the intermediate input embedding 104 to output the output data 140 from the second head 120 of the machine learning model 1 10. As will become readily apparent, the output data 140 may be any suitable t pe of data. That is, the machine learning model 110 may be trained to infer any suitable prediction based upon the query 100 for generating the response 150. In some examples, the output data 140 indicates at least a portion of the response 150 (e.g. where the output data indicates a word and the response 150 indicates a plurality of words including the word indicated by the output data 140). In a specific example, the machine learning model 110 may process [-0.1003, 0.2410, 0.2310, -0.0612, 0.1812. -0.0606, 0.1807 ... ] to generate output data 140 indicating [0.76, 0. 21, 0.02, 0.006, ... ]. In this specific example, the output data 140 may be a probability distribution indicating one or more probability values each corresponding to a different token in the predetermined vocabulary of tokens, i.e. 0.76 may correspond to a token representing a w ord such as “The”. An output sequence of tokens may therefore be generated autoregressively by performing multiple iterations each for predicting the next output token in the sequence. That is, in this specific example, the output data 140 may be used to select a token representing theword “The'’ as the next output token for the output sequence of tokens. In other examples, the output data 140 may be a token (e.g. 6769) from the predetermined vocabulary of tokens, rather than a probability distribution indicating a next token in the sequence. In other examples, the output data 140 may be classification data for classifying the query 100 or signal data for signaling the response 150 and / or controlling a device, as described in further detail below. Outputting from the second head 120 may occur in a second forward pass of the machine learning model 110. i.e. a second iteration after the first iteration of providing input to the machine learning model 110. The second head 120 may be a final output layer of the machine learning model 110, i.e. an output layer configured to generate a final output for generating the response 150 to the query 100. In some examples, the second head 120 is connected to a final attention layer 116 of the machine learning model 110, as described below.

[0045] The method further comprises generating the response 150 to the query 100 based upon the output data 140. As described above, the output data 140 is outputted from the second head 120 of the machine learning model 110. In the specific example above, the output data 140 is [0.76, 0. 21, 0.02, 0.006, ... ] indicative of a probability distribution for the predetermined vocabulary of tokens. The output data 140 may thus be used to select a next output token for the output sequence of tokens. The next output token may in this specific example be representative of “The'’. The multiple iterations, as previously described, may generate one or more output tokens representing “The”, “answer”, “to”, “33”, “x”. “2”, “is”, and “66” as the output sequence of tokens. The output sequence of tokens may be provided as the response 150 to the query 100. As such, the response to the query 100 “What is the answer to 33 x 2?” is generated based upon the output data 140. By generating and leveraging the first embedding 130, the method of the first aspect demonstrates one or more of the advantages, as described above and further described below with reference to FIG.7 - FIG. 10, when generating the response 150 to the query 100.

[0046] In some implementations, the first embedding 130 is indicative of rationale for the response. As previously described, machine learning models trained to leverage CoT (e.g. as described with reference to FIG.2 below) or ToT may be specifically trained to generate tokens explicitly indicating rationale for their response. Implementations described herein remove the need to explicitly recite this rationale by leveraging the first embedding 130 instead (and the second embeddings described below with reference to FIG.3) to confer the advantages described above. That is, a machine learning model trained to use CoT reasoning may receive a prompt such as “What is the answer to the question 33 x 2? Explain your reasoning step by step” and in response generate lengthy rationale as part of the response, e.g. “To solve themultiplication problem 33 x 2. let's break it down step by step: Step 1: 33 x 2 may be broken down into 3 x 2 and 30 x 2. Step 2: 3 x 2 = 6. Step 3: 30 x 2 = 60. Step 4: 6 + 60 = 66. The answer is therefore 66.”. For more complicated problems (e.g. spatial reasoning problems or complex games), number of tokens indicating the rationale may be hundreds or thousands of tokens. In contrast, the first embedding 130 may be indicative of such rationale in a latent embedding space of the machine learning model 110. which may be conceptualized as an internal and native "‘language” of the machine learning model 1 10, thus significantly reducing the number of inference cycles required to generate the response 150. This is advantageous in a number of different specific contexts and purposes (e.g. image analysis, audio analysis, video analysis, robotic or control methods, and contexts where spatial reasoning is vital for accurate predictions / inferences). A machine learning model implementing CoT will now be described with reference to FIG.2.

[0047] FIG.2 depicts an example of a machine learning model 210 implementing CoT reasoning for generating a response 250 to a query 100.

[0048] At a first iteration 260a. the machine learning model 210 may process the query 100 to generate one or more query tokens 202. For example, in the example above, the one or more query tokens may include a token representing “33”, “x”, “2”, and “=”. The one or more query’ tokens may be generated by processing the query' 100 using a tokenizer layer 212 of the machine learning model 210. At the first iteration 260a, the machine learning model 210 may process the one or more query tokens 202 using one or more token embedding layers 214 to generate an intermediate input embedding 204. The intermediate input embedding 204 may represent and / or be based upon the one or more query tokens 202. The intermediate input embedding may comprise one or more query embeddings generated by the token embedding lay er(s) 214 for each respective token of the one or more tokens 202. At the first iteration 260a, the machine learning model 210 may process the intermediate input embedding 204 using one or more attention layers 216 to generate a latent output embedding 206. At the first iteration 260a, the machine learning model 210 may process the latent output embedding 206 using an output head 218 to generate a next output token in a current output sequence of tokens. At the first iteration 260a, the next output token may be appended to the current output sequence of tokens. It will be appreciated that, at the first iteration 260a, the current output sequence of tokens does not comprise any tokens until the first next output token is appended. At subsequent iterations 260b - 260N, the current output sequence of tokens may be processed by the successive token embedding layer(s) 214, attention layer(s) 216, and output head 218 to generate a new next output token. With reference to the example depicted in FIG.2, at the firstiteration 260a, the next output token may be a token representing “6”. In the same example depicted in FIG.2, at the second iteration 260b, the next output token may be a token representing£'+”, and so on for subsequent iterations 260c - 260e. The process described above in this paragraph may be repeated for N iterations (e.g. 260a - 260e). At one or more of the iterations 260a - 260d, tokens explicitly indicating rationale 200 for the response 250 may be generated. For example, in the example depicted in FIG.2, the machine learning model 210 may output a current output sequence of tokens representing ”6+60=’' as the rationale for the response. Of course, in practice, the rationale 200 may be different and may include thousands of tokens to form the basis of the rationale 200 for the response 250. In other words, the machine learning model 110 may be trained to generate an output which decomposes or extracts information from the query 100 using CoT reasoning to provide, at a final iteration, basis to infer the response 250. That is, at a final iteration 260e, the process described above may be repeated by processing the current output sequence of tokens (i.e. including the tokens explicitly indicating the rationale 200 for the response 250) to generate a final output token 240 for the response 250 which is at least based upon the tokens explicitly indicating the rationale 200.

[0049] FIG.3 depicts an example of the machine learning model 110 implementing Latent Thought Chain (LATCH) reasoning for generating the response 150 to the query 100 in accordance with the techniques discussed herein.

[0050] At a first iteration 360a. the machine learning model 110 may process the query 100 using the tokenizer layer 1 12 to generate the one or more query tokens 102. At the first iteration 360a, the machine learning model 110 may process the one or more query tokens 102 using one or more token embedding layers 114 to generate the intermediate input embedding 104. The token embedding layer(s) 114 may be configured to generate a respective embedding for each of the one or more query tokens 102. The intermediate input embedding 104 may be based upon the one or more query tokens 102 in that the intermediate input embedding 104 represents the one or more query tokens 102 at the first iteration 360a. At the first iteration 360a. the machine learning model 110 may process the intermediate input embedding 104 using one or more attention layers 116 to generate a latent output embedding 106. At the first iteration 360a, the machine learning model 110 may process the latent output embedding using the first head 118 to generate the first embedding 130. The first embedding 130 with reference to FIG.3 is denoted zl. As will become apparent, at subsequent iterations, respective second embeddings 330a - 330c may be generated. That is. at a second iteration 360b. the first embedding 130 may in examples be concatenated with (or appended / prepended to) theintermediate input embedding 104 generated at the first iteration 360a to update the intermediate input embedding 104. At the second iteration 360b. the updated intermediate input embedding 104 may be processed by the machine learning model 110 using the one or more attention layers 116 to generate the latent output embedding 106. The latent output embedding 106 (now at the second iteration 360b being based upon the updated intermediate input embedding 104 comprising the query token embeddings and first embedding 130) may be processed by the machine learning model 110 using the first head 118 to generate a second embedding 330a for the second iteration 360b. This process may be repeated for multiple respective second embeddings (e.g. 330b and 330c) to further indicate rationale for the response 150. At a final iteration 360e, the updated intermediate input embedding 104 may be processed by the second head 120 to generate the output data 140. The response 150 may, accordingly, be based upon the output data 140 generated at the final iteration 360e.

[0051] The iterative process of generating and appending soft tokens described with reference to FIG.3 allows the machine learning model 110 to build a chain of reasoning within a continuous vector space. The architectural choice for the machine learning model 110 (i.e. incorporation of a first head 118) provides the specific advantage of representing complex, potentially non-linguistic reasoning (such as spatial relationships or abstract plans) in a compact and computationally efficient format. Each subsequent soft token 130, 330a- 330c is conditioned on the previous ones, enabling the model to refine its internal representation of the problem-solving trajectory, over multiple orders, before committing to a final output, thereby improving the quality and accuracy of the response 150 for complex reasoning tasks (e.g. image classification; depth estimation) while reducing computational complexity.

[0052] CoT may be combined with LATCH such that inference using the second head 120 may be repeated multiple times to, for example, either (i) additionally generate output tokens explicitly indicating rationale for the response or (ii) generate further tokens indicating the response per se, for example where the answer cannot be represented by a single token such as the token representing “66” as the answer to “33x2=” in the example above.

[0053] FIG.4A depicts an example predetermined vocabulary of tokens 400. FIG.4B depicts a first set of discrete tokens 420a for generating a response to a query. FIG.4C depicts a second set of tokens 420b including a plurality of discrete tokens 102 and a soft token 130 for generating a response to a query. FIG.4D depicts a space of discrete values 430 defined by a predetermined vocabulary of tokens 400. FIG.4E depicts a latent embedding space 440 of continuous values.

[0054] The predetermined vocabulary' of tokens 400 may include a plurality of discrete values (e.g. 0, 1, 2, 3. ... . 274) each corresponding to a respective word or word piece (e.g. “x”, “ / ”, “+”, ... , “[EOS]”). The predetermined vocabulary7of tokens 400 may thus define a space (e.g. numeric space) of discrete values (e.g. numeric values) for representing e.g. yvords, word pieces, image chunks, etc. An example of such a space of discrete values 430 is depicted in FIG.4D in which space 430 only includes values that are discrete (i.e. does not include any continuous values). The machine learning model 210 of FIG.2. for example, when generating its response to the query 100, may only generate tokens as output, or its rationale for the response 250, that are selected or part of the predetermined vocabulary of tokens 400 defining the discrete space of values. In contrast, the machine learning model 110 which is described with reference to FIG. l and FIG.3 may also generate / output tokens that are not part of the predetermined vocabulary of tokens. That is, the machine learning model 110 may be configured to generate soft tokens that, rather than being sampled from a discrete space 430, comprise one or more continuous values that are sampled from a latent embedding space 440 of continuous values as depicted in FIG.4E. For example, the first embedding 130, zl, may be sampled from the latent embedding space 440. In other examples, the respective second embeddings 330a - 330c may also be sampled from the latent embedding space 440. An input to the machine learning model 110, 210 may include the query token(s) 102, 202 and tokens indicative of rationale for the response 150. 250. For example, for a machine learning model 210 implementing CoT, an input 420a to the machine learning model 210 may be a combination of the query token(s) 202 (e.g. discrete values 73, 0, 42, 4 corresponding to “33”, “x”, “2”, and “=” in the example above) and tokens explicitly indicating the rationale 200 (e.g. discrete values 42, 2, 100 corresponding to “6”, “+”, “60” in the example above). In another example, for a machine learning model 110 implementing LATCH, an input 420b to the machine learning model 1 10 may be a combination of the query token(s) 102 (e.g. discrete values 73, 0, 42, 4 corresponding to “33”, “x”, “2”, and “=” in the example above) and one or more soft tokens, e.g. the first embedding 130, indicating the rationale 300 for the response 150.

[0055] That is, in some implementations, the first embedding 130 is a soft token not part of a predetermined vocabulary of tokens 400 for selecting at least one of the one or more query tokens. In such implementations, the predetermined vocabulary7of tokens 400 defines a space of discrete values 430. In such implementations, the soft token comprises one or more continuous values sampled from a latent embedding space 440 of continuous values. In some examples, the latent embedding space 440 of continuous values is distinct from thepredetermined vocabulary of tokens 400, i.e. the latent embedding space may not be exactly the same as the predetermined vocabulary of tokens, which may be a space of discrete values 430 rather than continuous values 440. In some examples, the latent embedding space 440 may include a projection (i.e. vector projection) of each token in the predetermined vocabulary of tokens 400. To illustrate, the predetermined vocabulary of tokens may comprise one or more discrete values each representing a different element (e.g. numeric values 0, 1, 2, 3, ... 274 each representing one of a word, wordpiece, image, portion of an image, etc.). In contrast, soft tokens may not represent or be selected based upon any particular element or portion of the query’ 100. As described above, in a specific example, the predetermined vocabulary of tokens 400 may be {“What”: 6304, “is”: 1873, “the”: 1531, “answer”: 9431, “to”: 2892, “33”: 4613, “x”: 5852, “2”: 4554, “?”: 5008. . .. }. In this example, there may not be any word or portion of the query 100 corresponding to the first embedding 130 as [-0.8384, 0.7654, 0.4273, -0.3692, 0.2321, . . . ]. The latent embedding space 440 may be a space of continuous values (e.g. one or more values between -1 and 1) which the machine learning model 110 is trained to sample from (i.e. when outputting the first embedding from the first head) during inference. That is, in a similar manner to “soft prompts”, which are embedding vectors that are learnt during a training phase and configured to adapt operation of a machine learning model 110 once inserted into an input sequence, “soft tokens” are also learnt during training and may be configured to modify operation of the machine learning model 110 (i.e. affect the output data 140 and thus the response). In general, the soft token may comprise an embedding, i.e. an ordered collection of numerical values, in particular a vector with continuous-valued elements, that represents a token but is not restricted to discrete values that represent tokens in the predetermined vocabulary of tokens 400 (i.e. a soft token does not represent any particular item in the predetermined vocabulary of tokens, such as a word or sub-word). In implementations, the soft token can be an embedding in a learned embedding space 440 of the machine learning model 110.

[0056] The use of soft tokens 130, 330a - 330c comprising continuous values sampled from the latent embedding space 440allows reasoning steps for which no adequate natural language descriptor exists to be represented. For example, the precise spatial relationship between objects in an image / video or a planned trajectory for a robotic device or its endeffector, which could not otherwise be represented by discrete tokens, can be represented in the latent embedding space 440. The same is true for other datatypes, such as the spatial relationship of different sounds represented in audio data. Furthermore, from a computational standpoint, a single soft token (e.g. the first embedding 130) can encode more information in asingle reasoning step (e.g. a forward pass with the first head 118) than a discrete word-token, thereby providing a more compressed and efficient representation of the rationale. This compression directly contributes to reduced memory usage and faster processing during inference.

[0057] In some implementations, the method further comprises generating, for each of the one or more query tokens 102, one or more query embeddings each indicative of the respective query token 102. In such implementations, outputting, from the first head 118 of the machine learning model 110, the first embedding 130, comprises processing the one or more query embeddings using the machine learning model 110. That is, the machine learning model 110 may process the one or more query tokens 102 and in response generate the one or more query embeddings. For example, the machine learning model 110 may process 6304. 1873, 1531, 9431, 2892, 4613, 5852, 4554, 5008 and in response generate an embedding for each of those tokens 102. In some examples, each query embedding may be generated using one or more token embedding layers 114, as described below. In some examples, each query7embedding may be a projection of the respective token into an n-dimensional embedding space. In some examples, each query embedding may comprise one or more continuous numenc values. In some examples, each query embedding may be an n-dimensional vector such as [-0.5123, 0.9431, 0.3195, -0.2784, 0.4186, ... ] corresponding to a token 102 representing "What" and [-0.6749, 0.8576, 0.5213, -0.4561, 0. 1932, . . . ] corresponding to a token 102 representing "is", etc. To output the first embedding 130. the machine learning model 110 may process, using one or more of its trainable parameters, the one or more query embeddings each indicative of a query7token 102 to output the first embedding 130 from the first head 118. In some examples, the one or more query7embeddings is embedded to a same latent embedding space 440 of continuous values as the soft token, i.e. the latent embedding space 440 of continuous values previously described. As will become readily apparent, processing the one or more query embeddings (and the first 130 and respective second embeddings 330a - 330c, as below in subsequent iterations) using the machine learning model 110 (e.g. using one or more attention layers 116) may generate a latent output embedding 106 suitable for generating the output (i.e. the first embedding 130 in a first iteration) from the first head 118. The latent output embedding 106 may be indicative of the one or more query7embeddings (and the first 130 and respective second embeddings 330a - 330c, as below in subsequent iterations).

[0058] In some implementations, generating the intermediate input embedding 104 comprises concatenating the one or more query embeddings with the first embedding 130. For example, if the query embeddings were [-0.5123, 0.9431, 0.3195, -0.2784, 0.4186, ... ] and[-0.6749, 0.8576, 0.5213, -0.4561, 0.1932, ... ], and the first embedding 130 output from the first head 118 based upon those query tokens 102 was [-0.8384. 0.7654, 0.4273, -0.3692, 0.2321, ... ], the intermediate input embedding 104 may be [[-0.5123, 0.9431, 0.3195, -0.2784, 0.4186, ... ], [-0.6749, 0.8576, 0.5213, -0.4561, 0.1932, ... ], [-0.8384, 0.7654, 0.4273, -0.3692, 0.2321, ... ]]. Generating the intermediate input embedding 104 may be achieved in any suitable way such that the intermediate input embedding 104 indicates at least the one or more query tokens 102 and the first embedding 130. For example, mean pooling (i. e. an average of the set of one or more query embeddings and first embedding 130) could be used in some implementations. In some examples, concatenating the one or more query embeddings with the first embedding 130 may comprise appending the first embedding 130 to the one or more query embeddings, e.g. at each iteration 360a - 360d, as below.

[0059] Concatenating the query embedding(s) with the first embedding 130 allows the proceeding iteration of the reasoning process of the machine learning model 110 to be conditioned on the entirety of the preceding reasoning chain. By explicitly appending the generated embedding(s) (i.e. the first embedding 130 or second embedding(s) 330a - 330c) to the sequence of query embedding(s). the subsequent forward pass through the attention layer(s) 116 can directly attend to this compressed representation of the rationale. This ensures that the state of the machine learning model 110 is updated based upon the output of the previous reasoning step, enabling a structured, multi-step problem-solving process within the latent space 440, which has the advantage of improving accuracy on tasks that cannot be solved in a single inference step.

[0060] In some implementations, the method further comprises, for a number, k. of iterations outputting, from the first head 118 of the machine learning model 110, a respective second embedding 330a - 330c based upon the intermediate input embedding 104 and updating the intermediate input embedding 104 by concatenating the intermediate input embedding 104 with the respective second embedding 330a - 330c. In such implementations, outputting the output data 140 is based upon the updated intermediate input embedding 104. That is, the process of outputting an embedding from the first head 118 may be repeated one or more times (e.g. 360a - 360d) to generate a sequence comprising the one or more query embeddings and one or more additional embeddings (e.g. EM, EM, EM, EM, zl, z2, ... , zk depicted in FIG.3), the additional embeddings being output from the first head 118 of the machine learning model 110, referred to in this part as respective second embeddings 330a - 330c. However, unlike with the first embedding 130, each additional embedding may be generated based upon every preceding embedding output from the first head 118 (e.g. zk 330c may be generated based uponthe first embedding 130 and the second embedding(s) 330a, 330b). In other words, at each subsequent iteration, the next soft token may take into account both the query 100 and each preceding soft token (e.g. 130, 330a, 330b for zk). Accordingly, the intermediate input embedding 104 may be extended by including (e.g. concatenating, appending, etc.) an additional second embedding 330a - 330c for the sequence. For example, as above, the intermediate input embedding 104 generated based upon the one or more query embeddings and the first embedding 130 was, in the specific example. [[-0.5123, 0.9431, 0.3195. -0.2784. 0.4186, ... ], [-0.6749, 0.8576, 0.5213, -0.4561, 0.1932, ... ], [-0.8384, 0.7654, 0.4273, -0.3692, 0.2321, ... ]]. At a next iteration, the first head 118 of the machine learning model 110 may output a second embedding 330a - 330c, such as [-0.1256, 0.7369, 0.4903, -0.3624, 0.8921. ... ] based upon the current intermediate input embedding 104 [[-0.5123, 0.9431, 0.3195, -0.2784, 0.4186, ... ], [-0.6749, 0.8576, 0.5213, -0.4561, 0.1932, ... ], [-0.8384, 0.7654, 0.4273, -0.3692, 0.2321, ... ]], e.g. by providing the intermediate input embedding 104 as input to the machine learning model 110 in a first forward pass. Subsequently, the current intermediate input embedding 104 may be updated (i.e. modified, for example by concatenation) to generate [[-0.5123, 0.9431. 0.3195. -0.2784. 0.4186. ... J, [-0.6749, 0.8576. 0.5213, -0.4561, 0.1932, ... ], [-0.8384, 0.7654, 0.4273, -0.3692, 0.2321, ... ], [-0.1256, 0.7369, 0.4903, -0.3624, 0.8921, ... ]] as an updated intermediate input embedding 104. This updated intermediate input embedding 104 may be provided as input to the machine learning model 110 in a second forward pass to output the output data 140 from the second head 120, as described above. Thus, each subsequent output embedding may take into account each preceding output embedding. Accordingly, this enables the machine learning model 110 to reason recursively, i.e. generate abstract representations which are dependent upon multiple different orders of output from the first head 118, and therefore represent the query 100 at different levels of abstraction without having to generate explicit rationale which can be computationally expensive. In experiments, the number of iterations used to generate subsequent embeddings 330a - 330c was ten iterations or less, which is far fewer than the number of iterations required using CoT or ToT reasoning to achieve adequate performance for particular tasks, which can in some cases reach hundreds or thousands of iterations. Accordingly, inference time is improved and computational complexity is reduced in a number of different contexts.

[0061] In some implementations, the machine learning model 110 comprises a tokenizer layer 112 for receiving the query 100 and generating the one or more query tokens 102, one or more token embedding layers 114 for generating the one or more query embeddings basedupon the one or more query tokens 102, and one or more attention layers 116 configured to process the one or more query embeddings for outputting the first embedding 130 from the first head 118. In such implementations, the one or more attention layers 116 are further configured to process the intermediate input embedding 104 for outputting the output data 140 from the second head 120. In such implementations, the first head 118 and the second head 120 are connected to the one or more attention layers 116. In some examples, the first head 118 and the second head 120 may be included in separate portions of the machine learning model 110 (e.g. separate portions of a neural network), but which are connected to a same tokenizer 112, embedding layers 114, and / or attention layers 116 (e.g. a same Transformer trunk). In other examples, the first 118 and second 120 heads may be connected to entirely separate neural networks (e.g. ones not connected to the same tokenizer layer, token embedding layer(s), and attention layer(s)). Portions of the machine learning model 110 connected to the first head 118 of the machine learning model 110 may comprise far fewer parameters (e.g. 100 times fewer) than those portions of the machine learning model 110 connected to the second head 120. Thus, inferring the output from the first head 118 may be much more computationally efficient than e.g. the second forward pass for generating output from the second head 120.

[0062] FIG.5A depicts a Gaussian distribution 500 for sampling the first embedding 130. FIG.5B depicts a Gaussian head of the machine learning model 110 configured to sample the first embedding 130 from the Gaussian distribution 500.

[0063] In some implementations, the first head 118 is a Gaussian head configured to sample the first embedding 130 from the Gaussian distribution 500. In such implementations, the Gaussian distribution 500 is parameterized based upon a latent output embedding 106 generated by the machine learning model 110 in response to providing the one or more query tokens 102 as input to the machine learning model 110. In general, the Gaussian distribution 500 may be parameterized using one or more Gaussian parameters 510a, 510b. For example, the Gaussian distribution 500 may be parameterized based upon mean and variance parameters, as discussed below. The sampling from the Gaussian distribution may occur using any suitable sampling mechanism 520 (e.g. Box-Muller Transform, Reparameterization Trick, and Inverse Transform Sampling). In some examples, the latent output embedding 106 may be an output of a final one of the attention layers 11 prior to being processed by the first head 118, and e.g. may be indicative of the one or more query embeddings. The latent output embedding 106 may be a vector comprising a plurality of continuous values. In some examples, the Gaussian head 118 (e.g. an output layer using the Gaussian distribution 500) may utilize a multivariate Gaussian distribution to generate the first embedding. For example, the Gaussian distributionmay be parameterized by a mean vector and covariance matrix each comprising a plurality' of values. In such examples, the mean vector and covariance matrix may be derived from the latent output embedding 106 in any usual way (e.g. an average value for each dimension of the latent output embedding 106 for the mean vector). By sampling from the Gaussian distribution 500 in this way, the first head 118 extracts relevant information for generating the output data 140 from both the prompt (i.e. one or more query tokens 102), model weights (i.e. embedded information leamt during training), and in some examples, preceding output embeddings (i.e. first 130 or second 330a - 330c embeddings) from the first head 118.

[0064] While two linear projection vectors are described above for inferring the Gaussian parameters 510 from the latent output embedding 106 to serve as the mean (p) and a component of the variance (v), several alternative inference methods are also envisaged. For example, instead of a simple linear projection, a more complex non-linear mapping could be used, where a multi-layer perceptron (MLP) with activation functions processes the latent output embedding 106 to generate the Gaussian distribution's parameters 510. This provides an even more expressive relationship between the machine learning model's 110 internal state and the resulting embedding 130, 330a- 330c. In general, the Gaussian head 118 may use one or more linear and / or non-linear function(s) to process the latent output embedding 106 to generate the Gaussian parameters 510.

[0065] Using the Gaussian distribution 500 to sample the soft token (i.e. the first embedding 130 or the second embedding(s) 330a - 330c), rather than directly outputting a deterministic embedding (e g. by sampling a latent vector from one of the internal layers of the machine learning model 110), introduces controlled stochasticity into the reasoning process. This is advantageous, both during training and inference, because it encourages the machine learning model 110 to explore a wider region of the latent embedding space 440 and prevents it from collapsing and relying on a single, suboptimal reasoning trajectory. Use of the Gaussian distribution 500 also allows for the generation of multiple diverse reasoning trajectories (i.e. tentative embeddings - as described in more detail below) which can be evaluated, for example using importance sampling, to identify a more robust and accurate final response 150. This improves the reliability of the machine learning model 110 when faced with novel or ambiguous queries 100.

[0066] In some implementations, parameterizing the Gaussian distribution 500 based upon the latent output embedding 106 comprises generating a first linear projection vector of the latent output embedding 106, generating a second linear projection vector of the latent output embedding 106, and parameterizing the Gaussian distribution 500 according to: J\f (p, ev. Insuch implementations, p is the first linear projection vector and v is the second linear projection vector. For example, a linear transformation may be applied to the latent output embedding 106 such that the latent output embedding 106 is mapped to two d-dimensional parameter vectors, i.e. the first and second linear projection vectors, or p and v. The Gaussian distribution 500 may be parameterized based upon these two d-dimensional parameter vectors, as above. In some examples, evmay comprise one or more continuous values of v raised as a power of e. As previously mentioned, the Gaussian distribution 500 parameterized in this way may be a multivariate Gaussian distribution. The first embedding 130 may then be sampled from the Gaussian distribution 500 in any suitable way, as above. In general, the first embedding 130 may be the first output from the first head 118 for generating the response 150 to the query 100 (i.e. zt) and each subsequent second embedding 330a- 330c may be z2, z3, z4, ... and so forth (e.g. up to k).

[0067] In some implementations, sampling the first embedding 130 comprises reparameterizing z based upon:In such implementations, z represents output from the first head 118 and e represents a random value sampled from a Gaussian distribution 500. That is, the methods described herein may employ a reparameterization trick (Kingma et al., 2013). The random value may be a value sampled from a standard normal distribution (i.e. N (0, 1)).

[0068] In some implementations, the output data 140 is data indicating an output token (e.g. data indicating '‘66”, as depicted in FIG.3). In such implementations, the one or more query tokens 102 and the output token 140 each represent text data, image data, video data, audio data, signal data, sensor data, point-cloud data, or categorical data. That is, in some examples, the machine learning model 110 described herein may be a machine learning model configured to determine, e.g. autoregressively, a next output token in a current sequence of output tokens where the query 100 and the response 150 include any number of different modality inputs or outputs. For example, the machine learning model 110 may be trained to be integrated as part of smart home assistant system. In this example, the query tokens 102 may represent image, video, and audio tokens, but the output token may be an audio token only. In another example, the machine learning model 110 may be trained to be integrated as part of a chatbot assistant system. In this example, the query tokens 102 may represent text tokens only, but the response 150 may include one or more output tokens representing both images and text.

[0069] It will be understood that the query tokens 102 and the output token may, independently, therefore represent text, an image, a video, audio, a signal, a sensor input / output, a point-cloud, categorical information, or a combination thereof in the case of multiple query 102 and / or output tokens. Likewise, the tokens may represent a portion of the foregoing datatypes (e.g. a wordpiece, a chunk of an image, etc ).

[0070] In some implementations, the response 150 to the query 100 is for performing a task, the task comprising a classification task, a control task for a robotic device, a machine translation task, a summarization task, a question-answering task, a data generation task, a signaling task, a spatial reasoning task, a programming code generation task, a transcription task, a captioning task, or a dialogue task.

[0071] The task, as referred to herein, may include any suitable task that may be performed based upon the output data 140. In some examples, the one or more query tokens 102 may each represent text data, image data, video data, audio data, or sensor data. That is, the task may be performed by processing text data, image data, video data, audio data, or sensor data using the machine learning model 110.

[0072] The task may include a classification task such as classifying the query 100 (i.e. text classification, image classification, audio classification, video classification, sensor classification, or time-series classification). For example, the query tokens 102 could represent time-series data and the task could be to classify that time-series data into one or more categories by outputting categorical data (i.e. the output data).

[0073] The task may include a control task for a robotic device. For example, the output data 140 may include signal data that causes the robotic device to be controlled (i.e. by providing the signal data to the robotic device) in response to processing query tokens representing image or video data. In another example, in a control task for a robotic device, the query tokens 102 may represent sensor data generated by a camera viewing a scene (e.g. an object on a table). The generated soft tokens (i.e. the first embedding 130 or second embedding(s) 330a - 330c) may represent a sequence of abstract motor control primitives, such as a planned multi-stage trajectory’ for a robotic end effector to pick up the object, including intermediate waypoints and gripper orientations. This non-linguistic rationale cannot be efficiently or adequately expressed using discrete tokens from a natural language vocabulary 400. The final output data 140 in this example may then be the specific low-level control signals, generated based on this latent plan, for actuating the motors of the robot. This provides for a more fluid and computationally efficient control mechanism compared to systems that must first generate and then parse a natural language description of the intended action(s).

[0074] The task may include a machine translation task. For example, the query tokens 102 could represent text in a first language and the response could represent a natural language translation of that text in a second language. The task may include a summarization task. For example, the response 150 could represent text indicating a summary of some text represented by the query7tokens 102.

[0075] The task may include a question-answering task. For example, the query tokens 102 could represent text indicating a question ‘‘What is this a picture of?” in addition to representing an image depicting an elephant. In this example, the response 150 may represent text indicating a response 150 to the question which indicates “The picture is of an elephant”. Additionally or alternatively the output data 140 could be categorical data indicating a classification of an elephant (i.e. a classification inferred based upon the low-level features of the image depicting the elephant).

[0076] The task may include a generation task, i.e. atask for generating data. For example, the query tokens 102 may represent an image depicting an elephant and the response 150 may represent a video (e.g. one or more image tokens each representing a different frame of the video) of the elephant, i.e. where the machine learning model 110 has been trained to generate videos based upon images.

[0077] The task may include a signaling task. For example, the query7tokens 102 may represent audio data indicating a user (e.g. a person in the real-world) speaking the words “Please schedule a meeting for next Monday” and the output data 140 could be a signal to a client device of the user which causes a meeting to be scheduled in a digital calendar of the client device.

[0078] The task may include a spatial reasoning task. For example, the query tokens 102 could represent an image depicting a 3D model of a cube with different colours on each side and the output data 140 may be configured to predict an orientation of the 3D model in 3D space. In another example, the query tokens 102 could represent an initial state of a game of Sudoku and the output data 140 (e.g. output token) at each iteration could indicate a particular move (e.g. an indication of a placement of a value between 0 and 9) in the game of Sudoku.

[0079] The task may include a programming code generation task. For example, the query tokens 102 may represent a query such as “Please could you generate some code in Python for solving a quadratic equation” and the response 150 (i.e. one or more output tokens) may represent the Python code for solving the quadratic equation.

[0080] The task may include a transcription task. In examples, the transcription task may be to transcribe the query 100. For the transcription task, the query 100 (e.g. as indicated bythe query token(s) 102) may represent text data, an image, audio data, a video, signal data, sensor data, point-cloud data, and / or categorical data that includes aspects of the data to be transcribed (e.g. audio or video represented by the data to be transcribed).

[0081] The task may include a captioning task. In examples, a captioning task may be to caption the query 100. For the captioning task, the query 100 (i.e. as indicated by the query token(s) 102) may represent text data, an image, audio data, a video, signal data, sensor data, point-cloud data, and / or categorical data that includes aspects to be captioned (e.g. to provide a text description of an image, audio, video, etc.).

[0082] The task may include a dialogue task. For example, the query tokens 102 may represent text including dialogue from multiple different persons and a question querying for information about each person and their relationship with each other person. In this example, the response 150 may represent text including information about each person and their relationship with every' other person.

[0083] In some implementations, the task is a task for a real-world environment. In such implementations, the output data 140 is configured to control one or more mechanical and / or electronic devices for performing the task in the real-world environment. That is, the machine learning model 1 10 may be trained to generate a response 150 for performing tasks in the real- world environment based upon real-world inputs (e.g. image or video data). For example, the task could include tasks in the field of autonomous self-driving, industrial robotics, healthcare, smart homes, agriculture, environmental monitoring etc. In some examples, the one or more query' tokens 102 may each represent text data, image data, video data, audio data, or sensor data. That is, the task may be performed by processing text data, image data, video data, audio data, or sensor data using the machine learning model 110. To perform tasks in the real-world environment, the machine learning model 110 via the output data 140 may be configured to control (e.g. by outputting a signal as previously described) one or more mechanical and / or electronic devices for performing the task in the real-world environment. In other words, the mechanical and / or electronic devices may be controlled for performing the task in the real- world environment in response to the machine learning model 110 outputting data inferred based upon real-world inputs such as image, video, audio, etc. The mechanical and / or electronic devices are devices configured to perform specific tasks utilizing electrical and / or mechanical systems. Electronic devices may operate by e.g. using electrical circuits including components such as transistors, diodes, and microprocessors (e.g. to process instructions for controlling the mechanical and / or electronic device) to perform the task. Mechanical devices may operate by e.g. using physical movement and mechanical forces to perform the task.Mechanical devices may, in some examples, include gears, levers, motors, end effectors, etc. An example electro-mechanical device may be a robotic device, where the robotic device includes a circuit for processing instructions derived from the output data 140 and an end effector for performing the task, such as operating a lever arm in the real-world.

[0084] In some implementations, generating the response 150 to the query 100 for performing the task requires abstract or non-linguistic reasoning. That is, the task as previously described may require reasoning, e.g. one or more tokens in the case of CoT or ToT models previously discussed, which cannot be expressed or represented completely in natural language or is not included in the predetermined vocabulary of tokens 400. For example, abstract or non- linguistic reasoning may include an inference about a relative spatial position of objects depicted in an image (i.e. the query 100). In this example, no discrete tokens may exist within the predetermined vocabulary of tokens 400 to precisely represent the particular spatial relationship of the objects. By leveraging latent representations of said reasoning, rather than discrete representations (e.g. words or wordpieces), tasks that require more complex and abstract capabilities may be enhanced. The subject matter described herein has been shown to significantly outperform other state-of-the-art methods on such tasks.

[0085] In some implementations, the method further comprises updating one or more parameters of the machine learning model 110 by computing a loss value based upon the output data 140. That is, the machine learning model 110 as described herein may be trained using any suitable method. In one example, a training dataset is obtained comprising a plurality of training examples, where each example comprises a pair of inputs (e.g. an input sequence of query’ tokens 102) and outputs (e.g. an output token, or the next token in the output sequence). Where the machine learning model 110 is trained for a task different to next token prediction, the pairs of inputs and outputs in the training dataset may be different (e.g. an input may be the query tokens 102 and an output may be classification data classifying the query 100 into one or more categories). During training, the input for a given training example may be provided to the machine learning model 110 to perform the method described above to output the output data 140. The loss value may be computed based upon the output data 140 and the output of the training example. In other words, the output data 140 may be compared with ground truth data (i.e. the output of the training example) to compute the loss (e.g. difference between the two). The loss value may be computed using any suitable loss function, however, particular suitable loss functions and their corresponding loss values are described in detail below. Accordingly, the loss value may be used in any suitable way to update the trainable parameters of the machine learning model 110. For example, any suitable optimization algorithm may beused to compute gradients (e g. a value which indicates an adjustment for a respective trainable parameter) for the one or more parameters using backpropagation. The parameters may be adjusted or updated accordingly. This process may occur, during training, over multiple iterations for multiple different training examples and may occur in batches (e.g. batches of 64 training examples for each update).

[0086] In some implementations, the loss value comprises a marginal likelihood loss value defined by:In such implementations, x represents the query 100, z represents output from the first head 118 of the machine learning model, y represents output from the second head 120 of the machine learning model, and 6 represents the one or more parameters of the machine learning model 110. That is, the machine learning model 110 may be trained using a marginal likelihood loss function for computing the loss value. In some examples, the marginal likelihood loss function and loss value are a maximum marginal likelihood loss function and maximum marginal likelihood loss value respectively. That is, the machine learning model 110 may be trained to optimize the parameters of the model II 0 to maximize the marginal likelihood.

[0087] FIG.6A depicts an example training algorithm for training the machine learning model 110 in accordance with the techniques described herein. The example training algorithm using one or more training examples each comprising one or more query tokens 102 and corresponding output data 140 (i.e. Data pair (x, y)). The example training algorithm includes an outer loop for trajectories t=l to T and an inner loop for iterations i=l to Niatents. In the outer loop, the inner loop is repeated i=l to N times for each trajectory. In the inner loop, at each iteration a first embedding 130 (or indeed second embedding 330a- 330c for i>l) is generated using the trainable parameters of the machine learning model 110 based upon the query 100 (i.e. represented by the query token(s) 102)) and each preceding “latent thought7’ - i.e. the preceding first / second embeddings 130, 330a - 330c, or z i. At the end of each trajectory, t, a loss value is computed using the negative log function which means that a relatively high probability results in a relatively low loss value. The loss value is computed on the basis that, given the query tokens 102 and the output from the first head 118 (e.g. the first embedding 130), how likely is the ground truth output from the second head 120. At the end of the iterations, an importance weight is assigned to each of the loss values computed for the T trajectories. An importance weight is assigned by transforming the negative of each of the loss values to one or more respective probabilities using a softmax function - since the exampletraining algorithm uses a negative loss, trajectories with a low loss value are assigned a higher weight than traj ectories with a high loss value (i. e. less desirable traj ectories). A final loss value is computed as a weighted sum of the trajectory losses Lt . The machine learning model 110 trained using the example training algorithm does not update its trainable parameters with backpropagation to affect a change in the importance weights that are generated at each training iteration, as indicated by “s lopgrad" in the example training algorithm. For example, the importance weights may not be used to influence activations of the machine learning model 110 during training.

[0088] FIG.6B depicts a system 6 for training the machine learning model 110 in accordance with the techniques described herein. One or more training examples 600 are provided (e.g. obtained from a training dataset). Each training example 600 may comprise an example query 602 and ground truth data 604. The example uery 602 may be processed by the machine learning model 110 to generate output data 140 from the second head 120 (i.e. by generating respective embeddings from the first head 118 at one or more iterations, and inferring the output data 140 based therefrom). A loss value 606 may be computed based upon the output data 140. For example, the loss value 606 may be computed based upon the output data 140 and the ground truth data 604 of the particular training example 600 corresponding to the processed example query 602. The loss value 606 may be used by a training algorithm to compute one or more gradients for the machine learning model 110. An optimization algorithm (e.g. Adam) may backpropagate the gradients to update one or more parameters of the machine learning model 1 10. The result of this training process is a trained machine learning model 110. The training may occur, in specific examples, using the training algorithm depicted in FIG.6A and described above.

[0089] In some implementations, the first head 118 is trained to generate the first embedding 130 (and optionally the second embedding(s) 330a - 330c) using self-supervised learning (e.g. learning from data generated from the machine learning model 110 itself, such as the first and second embeddings 130, 330a - 330c). In some implementations, the second head 120 is trained to generate the output data 140 using supervised learning (e.g. learning from labelled data). In some implementations, the first head 118 is trained to generate the first embedding 130 (and optionally the second embedding(s) 330a - 330c) without using one or more training examples 600 comprising ground truth data 604 indicating an example output for the first head 118. In some implementations, the second head 120 is trained to generate the output data 140 using one or more training examples 600 comprising ground truth data 604 indicating example output for the second head 120. In some implementations, the machinelearning model 110 is trained without using one or more training examples 600 comprising example tokens indicating example rationale for an example of the response 150. For example, the machine learning model 1 10 may be trained with training example(s) 600 comprising an example query 602 representing “33 x 2 =” and ground truth data 604 representing “66”. In this example, the machine learning model 110 may not be trained with training example(s) 600 comprising example tokens indicating “6”, “+”. “60”, “=”, i.e. analogous to the rationale 200 depicted in FIG.2, as may otherwise be used in CoT to generate the response 250 to the query 100.

[0090] By training the machine learning model 110 in this way, i.e. by learning latent thoughts via self-supervision, the machine learning model 110 may effectively and efficiently leam abstract and task-agnostic representations which improve model accuracy in a number of different contexts (e.g. spatial reasoning). Moreover, collecting and annotating large amounts of training data is a practical burden and one that also incurs a high computational cost, e.g. to collate and process the data via human-machine interactions. The incorporation of selfsupervision into the training process, specifically for training the machine learning model 110 to generate rationale, removes the need to collect and annotate large amounts of training data indicating example rationale for the response. Additionally, processing training inputs with a large amount of tokens indicating rationale for the response can also be very' computationally expensive. Training the machine learning model 110 to incorporate self-supervision and generate its own rationale reduces computational expense in this regard by reducing the amount of tokens processed during training (and inference, as discussed above).

[0091] In some implementations, outputting the first embedding 130 based upon the one or more query tokens 102 further comprises, for a number of iterations, outputting, from the first head 118 of the machine learning model 110, a tentative embedding based upon the one or more query tokens 102 and determining the first embedding 130 based upon each tentative embedding using one or more corresponding importance weights. That is, the method described above may leverage importance sampling. In some examples, each tentative embedding is independently sampled from the first head 118 of the machine learning model 110, i.e. independent from every' other tentative embedding. As previously described, the method described above generates the first embedding 130 (and subsequent output embeddings from the first head 118), e.g. as indicative of rationale 300 for the response 150, rather than discrete tokens explicitly defined in the predetermined vocabulary of tokens 400 and recited in the sequence of output tokens. This removes explicit dependence on the output from the secondhead 120 (i.e. y), which is a significant problem for other methods e.g. CoT and ToT methods. That is, CoT and ToT methods commonly suffer from a “cheating” effect whereby output tokens may be generated as rationale 200 for the response 250, but those output tokens merely emulate proper logic or provide incorrect reasoning that only serves to optimize for the correct, or ground truth, next token in the current output sequence of tokens. To further enhance training, and account for the removal of explicit dependence as previously described, multiple outputs from the first head 118 may be generated and weighted using importance weighting to focus the loss function on the “better performing” embeddings 130, 330a - 330c - i.e. those embeddings which contribute most towards e.g. correctly predicting the next output token in the sequence of output tokens. The resulting machine learning model 110. once trained, may accurately generate abstract and task-agnostic representations of the query 100 (and other “latent thoughts”) for generating the response 150 to the query 100.

[0092] In some implementations, the loss value further comprises an importance weighted autoencoder loss value defined by:In such implementations. wtrepresents an importance weight for each of the number of iterations, softmaxtrepresents a softmax activation function for the number of iterations, and represents one or more of the tentative output embeddings for the number of iterations. That is, the machine learning model 110 may be trained using an importance weighted autoencoder loss function such that the loss value is computed as previously described above. The loss function includes a softmax activation function configured to output a distribution of values, w, for weighting each output from the first head 118. By sampling multiple outputs from the first head 118 and weighting accordingly, during training (e.g. for each training example), the machine learning model 110 may be updated based upon the parameters of the machine learning model 110 that achieve the best loss value. In some examples, the randomness of the activation function may be controlled by selecting a temperature parameter, 7'. for the softmax activation function, e.g. a value between zero and one, thus affecting the output distribution of values, i.e. a low temperature parameter such as zero may cause the output distribution of values (e.g. the result of a softmax activation function) to become less flat or exhibit less variation. In some examples, a stop gradient function, stopgrad, is applied to the importance weight. wt, for preventing gradients being computed for wtduring training (i.e. stopgrad[wt]). That is, the importance weight in some examples may not be a trainableparameter of the machine learning model 110; the importance weight may not influence activations of the machine learning model 110 generated during training.

[0093] The use of an importance weighted autoencoder loss directly addresses challenges associated with training the first head 118 to generate embeddings 130, 330a - 330c that are causally effective for the task solved by the second head 120. By weighting the loss contributions of different sampled reasoning trajectories (i.e. the first and second embeddings 130, 330a - 330c. denoted “zt” for a given trajectory), the training process is focused on rewarding the “latent thoughts” that lead to a higher probability of the “correct” output. This results in a more efficient training process, leading to faster convergence and a machine learning model 110 that produces accurate results for a specific task, such as correctly identifying object positions in an image / video (e.g. object positions for a robotic grasping operation targeted at the object).

[0094] In some implementations, generating the response 150 based upon the output data 140 comprises a number, 2N, of forward passes using the machine learning model 110, where N is a positive integer indicating a number of output tokens for the response 150. That is, in an example where the output data 140 is a next output token in a sequence of output tokens, the machine learning model 110 may perform a first forward pass to infer one or more embeddings 130, 330a - 330c as output from the first head 118 of the machine learning model 110, and may perform a second forw ard pass to infer the output data 140 as output from the second head 120 of the machine learning model 110. For each next output token in the response 150, a further iteration of the first and second forward passes may be required. For example, an output sequence of tokens comprising “The”, “answer”, “is”, and “66” may comprise processing input (e.g. query tokens 102 and optionally output from the first head 118) with the machine learning model 110 to generate an output a total of 8 times. Thus, only one forward pass is required to output e.g. rationale 300 for the response 150. In contrast, the computational complexify of state-of-the-art methods is much higher due to the significantly higher number of forward passes required to generate tokens (i.e. from the predetermined vocabulary' of tokens 400) indicating rationale 200 for the response 250.

[0095] By reducing the number of forward passes with the machine learning model 110 a tangible improvement in the internal operation of the computing system upon which it is implemented is provided. Specifically, the overall processor cycles, memory -to-processor data transfer (i.e. memory bandwidth), and energy consumption required to generate the response 150 are reduced. These advantages are particularly significant for deploying the machine learning model 110 on mobile or embedded devices with limited computational power andbatery life, or in large-scale data center applications where improvements in computational efficiency translate directly to significantly lower operational costs and / or higher throughput.

[0096] According to the techniques described herein, there is also provided a method of generating a neural network output. For example, the machine learning model 110 previously described may be a neural network.

[0097] The method comprises obtaining an input sequence of tokens (e.g. the query tokens 102) selected from a predetermined vocabulary of tokens 400. In such implementations, the input sequence of tokens represents a system input (e.g. the query 100). In some examples, the input sequence of tokens may represent a prompt or other input for the neural network.

[0098] The method further comprises, at one or more first processing iterations (e.g. iterations 360a - 360d), generating a first portion (e.g. the first embedding 130 and the respective second embedding(s) 330a - 330c) of an output sequence of tokens by processing the input sequence of tokens (e.g. the query tokens 102) and a current output sequence of tokens (e.g. the tokens generated as output from the first head 118 in FIG.3), using a neural network, to generate a neural network output for selecting a next output token to extend the current output sequence of tokens. In some examples, the neural network output are logits generated by the neural network for further processing.

[0099] The method further comprises, at the one or more first processing iterations, processing the neural network output using a token generation neural network (e g. a neural network including the first head 118) to generate a soft token (e.g. the first embedding 130) that is not part of the predetermined vocabulary of tokens 400, the soft token comprising one or more continuous values. In some examples, the soft token comprises (e.g. is defined by) one or more continuous values sampled from a latent embedding space 440.

[0100] The method further comprises, at the one or more first processing iterations, extending the current output sequence of tokens by appending the soft token to the cunent output sequence of tokens. For example, with reference to FIG.3, a soft token, z2 330a, may be appended to a current output sequence of tokens: {EM, EM, EM, EM, zl} to extend the current output sequence of tokens to: {EM, EM, EM, EM, zl, z2}.

[0101] The method further comprises, at one or more second processing iterations (e.g. the final iteration 360e depicted in FIG.3), generating a second portion (e.g. a token representing “66”, as depicted in the response 150 of FIG.3) of the output sequence of tokens by processing the input sequence of tokens and the current output sequence of tokens including the soft token (e.g. the second embedding 330c. zk. with reference to FIG.3), using the neuralnetwork, to generate the neural network output for selecting a next output token from the predetermined vocabulary of tokens 400 to further extend the current output sequence.

[0102] The method further comprises, at one or more second processing iterations, selecting the next output token from the predetermined vocabulary of tokens 400 using the neural network output.

[0103] The method further comprises, at a final iteration, obtaining a system output (e.g. the output data 140 or a full response 150) based upon at least the second portion of the output sequence of tokens. For example, the system output may require decoding (e.g. translating a numeric token to its corresponding text according to the predetermined vocabulary of tokens 400).

[0104] In some implementations, the system input defines a reasoning task. In such implementations, the system output comprises a result of the reasoning task.

[0105] In some implementations, the neural network output comprises a score distribution that assigns a respective score to each token in the predetermined vocabulary of tokens 400.

[0106] In some implementations, the neural network output comprises an intermediate output of the neural network from an intermediate layer of the neural network. In such implementations, the method further comprises processing the neural network output using one or more subsequent layers of the neural network, after the intermediate layer, to generate a score distribution that assigns a respective score to each token in the vocabulary- of tokens. For example, the score distribution may be a probability distribution generated using a softmax activation function.

[0107] In some implementations, the neural network has a Transformer-based architecture. In some implementations, the neural network comprises one or more attention layers.

[0108] In some implementations, the method further comprises obtaining the neural network with a pre-trained set of trainable parameters, obtaining a training dataset comprising one or more training data items, each training data item comprising an example system input and a corresponding example system output, training the neural network and token generation neural network by processing the example system input of a training data item using the method described with reference to the second aspect to generate the system output, and backpropagating gradients for the one or more trainable parameters of the neural network using a training objective function dependent on a difference between the example system output and the corresponding system output through the neural network and token generation neural network to update the trainable parameters of both the neural network and token generationneural network. In some examples, the set of trainable parameters are weights of the neural network.

[0109] FIG.7A depicts a first set 7a of experimental results for machine learning models trained using a plurality of different training methods. The results 700 for the machine learning model 110 trained according to the techniques previously described indicate better performance than the other training methods across at least six different industry standard benchmarks (i.e. CommonsenceQA. HellaSwag, PhysicalQA, SQuAD, GSM8K, and LAMBADA) - a higher number for each of these benchmarks indicates better model performance. These industry standard benchmarks evaluate a machine learning model’s capability in a number domains including reasoning ability (e.g. spatial reasoning), question-answering, mathematical problem-solving, and understanding (e.g. image / video or natural language understanding). For example, PhysicallQA is a benchmark for evaluating a machine learning model’s capability for understanding physical interaction of objects in the real-world (spatial reasoning). The machine learning model 110 trained using LATCH exhibits a score of 77. 10 for PhysicallQA outperforming models trained rather using few-shot direct prompting and several variations of model fine-tuning.

[0110] FIG.7B depicts a second set 7b of experimental results for a machine learning model trained using a plurality of different methods. The results 702a for the machine learning model 110 trained according to the techniques previously described also indicate better performance than the other methods of fine-tuning (e.g. standard; scratch token; PAUSE) tested. That is, the machine learning model 1 10 exhibited, in the results 702a, higher test accuracy during training across every Low-Rank Adaption (LoRA) rank tested. LoRA rank is a training condition that determines how many parameters of a given machine learning model are available for updates / fine-tuning. In general, a higher LoRA rank indicates that more parameters of the respective machine learning model are available for updates / fine-tuning.

[0111] FIG.8 depicts a third set 8 of experimental results for a machine learning model trained with and without Latent Thought Chain reasoning. That is, the results 800 for the machine learning model 110 trained using LATCH exhibited a test accuracy greater than 60 for each LoRA rank 1. 2. 8, 16, and 64. In contrast, a machine learning model trained via other means (i.e. without using LATCH; without using the first head 1 18) exhibited a test accuracy less than 60 for each of the same LoRA ranks, and significantly less for LoRA ranks 1, 2, and 8 specifically. The machine learning model 110 implementing the first head 118 is therefore particularly well suited in situations where the hardware used for training the model is compute or memory constrained. For example, a lower LoRA rank means that training the machinelearning model requires fewer computational resources (e.g. video random access memory (VRAM)). LATCH thus allows machine learning models to be effectively trained on hardware that may otherwise not be possible (e.g. due to the hardware not comprising enough VRAM to load model weights into memory). LATCH is therefore at least based upon considerations relating to the internal functioning and configuration of the computer hardware used for training. The same improvement in test accuracy for lower LoRA ranks is exhibited with reference to the results 702a depicted in FIG.7B.

[0112] FIG 9 depicts a fourth set 9 of experimental results for the machine learning model 110 trained with Latent Thought Chain reasoning. The fourth set 9 of experimental results are for the machine learning model 110 trained to generate several different numbers of first / second embeddings 130, 330a - 330c in each respective test. The different numbers of soft tokens generated in each respective test were 1, 2, 5, 10, 20, 50, and 100.

[0113] FIG.10 depicts a fifth set 10 of experimental results for the machine learning model 110 trained with Latent Thought Chain reasoning. The fifth set 10 of experimental results are for the machine learning model 110 trained to generate output 130, 330a - 330c from the first head 118 having different dimensions. For example, the first embedding 130 may be a vector [-0.123, 0.691, 0.032] having three dimensions. The different number of dimensions for each soft token generated for the results were 64, 128, 256, 512, 1024, and 2048.

[0114] FIG. 11 depicts a flow diagram of a method for generating a response to a query.

[0115] At step 1100, receiving one or more query tokens, the one or more query tokens indicative of the query.

[0116] At step 1102, the method comprises providing the one or more query tokens as input to a machine learning model.

[0117] At step 1104, the method comprises outputting, from a first head of the machine learning model, a first embedding based upon the one or more query tokens.

[0118] At step 1106, the method comprises generating an intermediate input embedding based upon the one or more query tokens and the first embedding.

[0119] At step 1108, the method comprises outputting, from a second head of the machine learning model, output data based upon the intermediate input embedding.

[0120] At step 1110, the method comprises generating the response to the query based upon the output data.

[0121] One or more of the query tokens represent image data, video data, or audio data.

[0122] FIG. 12 depicts a flow diagram of a method for generating a neural network output.

[0123] At step 1200, the method comprises obtaining an input sequence of tokens selected from a predetermined vocabulary’ of tokens, the input sequence of tokens representing a system input.

[0124] At step 1202, the method comprises, at one or more first processing iterations generating a first portion of an output sequence of tokens by performing steps 1204, 1206, and 1208 described below.

[0125] At step 1204, the method comprises processing the input sequence of tokens and a current output sequence of tokens, using a neural network, to generate a neural network output for selecting a next output token to extend the current output sequence of tokens.

[0126] At step 1206, the method comprises processing the neural network output using a token generation neural network to generate a soft token that is not part of the predetermined vocabulary of tokens, the soft token comprising one or more continuous values.

[0127] At step 1208, the method comprises extending the current output sequence of tokens by appending the soft token to the current output sequence of tokens.

[0128] At step 1210, the method comprises, at one or more second processing iterations generating a second portion of the output sequence of tokens by performing steps 1212 and 1214 described below.

[0129] At step 1212, the method comprises processing the input sequence of tokens and the current output sequence of tokens including the soft token, using the neural network, to generate the neural network output for selecting a next output token from the predetermined vocabulary of tokens to further extend the current output sequence.

[0130] At step 1214, the method comprises selecting the next output token from the predetermined vocabulary of tokens using the neural netw ork output.

[0131] At step 1216, the method comprises, at a final iteration obtaining a system output based upon at least the second portion of the output sequence of tokens.

[0132] One or more of the input sequence of tokens represents image data, video data, or audio data.

[0133] The machine learning models as described herein may be neural networks. For example, the machine learning models may comprise a neural network having one or more (self-)attention layers, such as a Transformer neural network. The neural networks may be any of a variety7of Transformer-based neural netw ork architectures for example. Examples of such architectures include those described in J. Hoffmann, S. Borgeaud. A. Mensch, E. Buchatskaya, T. Cai. E. Rutherford. D. d. L. Casas, L. A. Hendricks. J. Welbl, A. Clark, et al. Training compute-optimal large language models, arXiv preprint arXiv:2203. 15556, 2022; J. W. Rae, S.Borgeaud, T. Cai, K. Millican. J. Hoffmann, H. F. Song, J. Aslanides, S. Henderson, R. Ring, S. Young, E. Rutherford. T. Hennigan. J. Menick, A. Cassirer, R. Powell, G. van den Driessche, L. A. Hendricks, M. Rauh, P. Huang, A. Glaese, J. Welbl, S. Dathathri, S. Huang, J. Uesato, J. Mellor, I. Higgins, A. Creswell, N. McAleese, A.Wu, E. Eisen, S. M. Jayakumar, E. Buchatskaya, D. Budden, E. Sutherland, K. Simonyan, M. Paganini, L. Sifire, L. Martens, X. L. Li, A. Kuncoro. A. Nematzadeh, E. Gribovskaya, D. Donato, A. Lazaridou, A. Mensch, J. Lespiau, M. Tsimpoukelli. N. Grigorev, D. Fritz. T. Sottiaux, M. Pajarskas, T. Pohlen. Z. Gong. D. Toyama, C. de Masson d’Autume, Y. Li, T. Terzi, V. Mikulik, I. Babuschkin, A. Clark, D. de Las Casas, A. Guy, C. Jones, J. Bradbury, M. Johnson, B. A. Hechtman, L. Weidinger, I. Gabriel, W. S. Isaac, E. Lockhart, S. Osindero, L. Rimell, C. Dyer, O. Vinyals, K. Ayoub, J. Stanway. L. Bennett, D. Hassabis. K. Kavukcuoglu, and G. Irving. Scaling language models: Methods, analysis & insights from training gopher. CoRR, abs / 2112. 11446, 2021; Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to- text transformer. arXiv preprint arXiv: 1910. 10683, 2019; Daniel Adiwardana, Minh-Thang Luong. David R. So, Jamie Hall, Noah Fiedel, Romal Thoppilan, Zi Yang. Apoorv Kulshreshtha, Gaurav Nemade, Yifeng Lu, and Quoc V. Le. Towards a human-like opendomain chatbot. CoRR, abs / 2001.09977, 2020; and Tom B Brown. Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan. Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.

[0134] Generally, however, the Transformer-based neural network includes a sequence of attention blocks, and, during the processing of a given input sequence, each attention block in the sequence receives a respective input hidden state for each input token in the given input sequence. The attention block then updates each of the hidden states at least in part by applying self-attention to generate a respective output hidden state for each of the input tokens. The input hidden states for the first attention block are embeddings of the input tokens in the input sequence and the input hidden states for each subsequent attention block are the output hidden states generated by the preceding attention block. It will be readily appreciated that such neural networks having a Transformer-based architecture may be used to generate the embeddings as described herein, for example, by sampling the input hidden states for a given block.

[0135] As described herein, the inputs and outputs may comprise tokens. In some implementations, the tokens can represent text, e.g., words, wordpieces or characters, in a natural or computer language. For example, text may be received, e.g., as a series of encodedcharacters, e.g. UTF-8 encoded characters; such “characters'’ can include Chinese and other similar characters, as well as logograms, syllabograms and the like. A text encoder, i.e. a tokenizer, can process a sequence of text to represent the text as a series of text tokens from a vocabulary of text tokens, e.g. that each represent words, wordpieces or characters in a natural or computer language. The computer language may be any formal language used to communicate with a computer, e.g. a markup language, or a command or configuration language, or a data exchange language such as JSON. or a programming language. The tokenizer can, e g., implement BPE (Byte Pair Encoding) or Wordpiece tokenization. Optionally the text can be obtained from audio data representing speech; the output tokens may be converted into audio data that represent speech corresponding to the text.

[0136] Also or instead the tokens may represent an image. For example, a set (sequence) of input or output tokens can represent an image. Each image token may comprise a block encoding of values of the pixels in a different region of an image that maps a set of values of the pixels to a respective image token. The block encoder may comprise a neural network, e.g. having one or more (self-)attention layers, such as a Transformer neural network as previously described.

[0137] Also or instead the tokens may represent an audio waveform. For example, a set (sequence) of input or output tokens can represent audio data representing a waveform e.g. instantaneous audio amplitude values or time-frequency audio data. Each image token may comprise a block encoding of the audio waveform in a different time segment of the audio that maps a set of values representing the audio waveform to a respective image token.

[0138] In some implementations, the machine learning models described herein are pretrained. e.g., trained on a particular modeling task prior to further training or inference. For example, the machine learning models described herein may be language models, vision models, multi-modal models, or any other suitable type of machine learning model that has been trained prior to inference and is suitable for processing the database data items described herein.

[0139] To illustrate, a system may pre-train a language model on a language modeling task, e.g.. atask that requires predicting, given a current sequence of text tokens, the next token that follows the current sequence in the training data. As a particular example, the language model can be pre-trained on a maximum-likelihood objective on a large dataset of text, e.g., text that is publically available from the Internet or another text corpus. It will be readily appreciated that the machine learning models described herein may further be fine-tuned to a particular task.

[0140] A description of self-atention, as may be employed by some of the machine learning models described herein, now follows.

[0141] A self-atention block, as referred to above, is a neural network layer that includes an atention mechanism that operates over the self-atention block input (or an input derived from the layer input) to generate the self-atention block output. A self-atention mechanism may be causally masked so that any given position in an input sequence does not attend over (e.g. use data from) any positions after the given position in the input sequence. There are many different possible atention mechanisms. Some examples of self-atention layers including attention mechanisms, are described in Vaswani et al. “Atention is all you need”, 31st Conference on Neural Information Processing Systems (NIPS 2017), Long Beach, CA, USA; Colin Raffel. Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 1910.10683, 2019; Daniel Adiwardana, Minh- Thang Luong, David R. So, Jamie Hall, Noah Fiedel, Romal Thoppilan, Zi Yang, Apoorv Kulshreshtha, Gaurav Nemade, Yifeng Lu, and Quoc V. Le. Towards a human-like opendomain chatbot. CoRR. abs / 2001.09977. 2020; and Tom B Brown, Benjamin Mann. Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.

[0142] Generally, an atention mechanism maps a query and a set of key-value pairs to an output, where the query, keys, and values are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function, e.g. a dot product or scaled dot product, of the query' with the corresponding key.

[0143] Generally, a self-atention mechanism is configured to relate different positions in the same sequence to determine a transformed version of the sequence as an output. For example, the atention layer input may comprise a vector for each element of the input sequence. These vectors provide an input to the self-atention mechanism and are used by the self-atention mechanism to determine a new representation of the same sequence for the attention layer output, which similarly comprises a vector for each element of the input sequence. An output of the self-atention mechanism may be used as the atention layer output, or it may be processed by one or more of feed-forward layers, skip connections, or normalization operations to provide the atention layer output.

[0144] In some implementations the atention mechanism is configured to apply each of a query' transformation e g. defined by a matrix WQ, a key transformation e.g. defined by a matrixWK, and a value transformation e.g. defined by a matrix IV17. to the attention layer input which is the input data X to the attention layer, to derive a query matrix Q = XWQthat includes a respective query for each vector in the input sequence, key matrix K = XWKthat includes a respective key for each vector in the input sequence, and value matrix V = XWVthat includes a respective value for each vector in the input sequence, which are used determine an attended sequence for the output. For example, the attention mechanism may be a dot product attention mechanism applied by applying each query vector to each key vector to determine respective weights for each value vector, then combining the value vectors using the respective weights to determine the self-attention layer output for each element of the input sequence. The selfattention layer output may be scaled by a scaling factor e.g. by the square root of the dimensions of the queries and keys, to implement scaled dot product attention. Thus, for example, andimension of the key (and value) vector. In another implementation the attention mechanism be comprise an “additive attention7' mechanism that computes the compatibility function using a feed-forward network with a hidden layer. The output of the attention mechanism may be further processed by one or more fully-connected, feed forward neural network layers.

[0145] The attention mechanism may implement multi-head attention, that is, it may apply multiple different attention mechanisms in parallel. The outputs of these may then be combined, e.g. concatenated, with a learned linear transformation applied to reduce to the original dimensionality if necessary.

[0146] The robotic device, as described above, may take various forms, including but not limited to a telepresence robot, a robotic arm, a bi-arm robotic device, a humanoid robot or other bipedal robot, a quadraped robot such as a “robot dog”, a wheeled robot, an aquatic robot, and so forth. The robotic device may include control logic. Control logic may take various forms, such as a real time controller, one or more processors, one or more field-programmable gate arrays (“FPGA”), one or more application-specific integrated circuits (“ASIC”), and so forth. In some implementations, the logic may be operably coupled with memory’. Memory may take various forms, such as random-access memory (“RAM”), dynamic RAM (“DRAM”), read-only memory (“ROM”), Magnetoresistive RAM (“MRAM”), resistive RAM (“RRAM”), NAND flash memory', and so forth. In some implementations, the control logic may be operably coupled with one or more joints, one or more end effectors, and / or one or more sensors. A joint of a robot may broadly refer to actuators, motors (e.g., servo motors), shafts, gear trains, pumps (e.g., air or liquid), pistons, drives, propellers, flaps, rotors, or othercomponents that may create and / or undergo propulsion, rotation, and / or motion. An end effector may broadly refer to a variety of tools that may be operated by the robotic device in order to accomplish various tasks. For example, an end effector may take the form of a claw with two opposing “fingers” or “digits.” Such a claw is one type of “gripper” known as an “impactive” gripper. The gripper may have more than tw o digits, for example, three, four or five digits. Other types of grippers may include but are not limited to “ingressive” (e.g., physically penetrating an object using pins, needles, etc.), “astrictive” (e.g.. using suction or vacuum to pick up an object), or “contigutive” (e.g., using surface tension, freezing or adhesive to pick up object). More generally, other types of end effectors may include but are not limited to drills, brushes, force-torque sensors, cutting tools, deburring tools, welding torches, containers, trays, and so forth. In some implementations, end effector may be removable, and various types of modular end effectors may be installed onto robot. Some robots, such as some telepresence robots, may not be equipped with end effectors.

[0147] In this specification, the term "configured" is used in relation to computing systems and environments, as well as computer program components. A computing system or environment is considered "configured" to perform specific operations or actions when it possesses the necessary software, firmw are, hardware, or a combination thereof, enabling it to carry' out those operations or actions during operation. For example, configuring a system might involve installing a software library with specific algorithms, updating firmware with new instructions for handling data, or adding a hardware component for enhanced processing capabilities. Similarly, one or more computer programs are "configured" to perform particular operations or actions when they contain instructions that, upon execution by a computing device or hardware, cause the device to perform those intended operations or actions.

[0148] The embodiments and functional operations described in this specification can be implemented in various forms, including digital electronic circuitry, software, firmware, computer hardware (encompassing the disclosed structures and their structural equivalents), or any combination thereof. The subject matter can be realized as one or more computer programs, essentially modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by or to control the operation of a computing device or hardware. The storage medium can be a storage device such as a hard drive or solid-state drive (SSD), a storage medium, a random or serial access memory device, or a combination of these. Additionally or alternatively, the program instructions can be encoded on a transmitted signal, such as a machine-generated electrical, optical, or electromagnetic signal, designed to carry information for transmission to a receiving device or system for execution by a computingdevice or hardware. Furthermore, implementations may leverage emerging technologies like quantum computing or neuromorphic computing for specific applications, and may be deployed in distributed or cloud-based environments where components reside on different machines or within a cloud infrastructure.

[0149] The term "computing device or hardware" refers to the physical components involved in data processing and encompasses all types of devices and machines used for this purpose. Examples include processors or processing units, computers, multiple processors or computers working together, graphics processing units (GPUs), tensor processing units (TPUs), and specialized processing hardware such as field-programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs). In addition to hardware, a computing device or hardware may also include code that creates an execution environment for computer programs. This code can take the form of processor firmware, a protocol stack, a database management system, an operating system, or a combination of these elements. Embodiments may particularly benefit from utilizing the parallel processing capabilities of GPUs, in a General-Purpose computing on Graphics Processing Units (GPGPU) context, where code specifically designed for GPU execution, often called kernels or shaders, is employed. Similarly, TPUs excel at running optimized tensor operations crucial for many machine learning algorithms. By leveraging these accelerators and their specialized programming models, the system can achieve significant speedups and efficiency gains for tasks involving artificial intelligence and machine learning, particularly in areas such as computer vision, natural language processing, and robotics.

[0150] 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.

[0151] A computer program, also referred to as software, an application, a module, a script, code, or simply a program, can be written in any programming language, including compiled or interpreted languages, and declarative or procedural languages. It can be deployed in various forms, such as a standalone program, a module, a component, a subroutine, or anyother unit suitable for use within a computing environment. A program may or may not correspond to a single file in a file system and can be stored in various ways. This includes being embedded within a file containing other programs or data (e.g., scripts within a markup language document), residing in a dedicated file, or distributed across multiple coordinated files (e.g., files storing modules, subprograms, or code segments). A computer program can be executed on a single computer or across multiple computers, whether located at a single site or distributed across multiple sites and interconnected through a data communication network. The specific implementation of the computer programs may involve a combination of traditional programming languages and specialized languages or libraries designed for GPU programming or TPU utilization, depending on the chosen hardware platform and desired performance characteristics.

[0152] In this specification, the term "engine" broadly refers to a software-based system, subsystem, or process designed to perform one or more specific functions. An engine is typically implemented as one or more software modules or components installed on one or more computers, which can be located at a single site or distributed across multiple locations. In some instances, one or more dedicated computers may be used for a particular engine, while in other cases, multiple engines may operate concurrently on the same one or more computers. Examples of engine functions within the context of Al and machine learning could include data pre-processing and cleaning, feature engineering and extraction, model training and optimization, inference and prediction generation, and post-processing of results. The specific design and implementation of engines will depend on the overall architecture and the distribution of computational tasks across various hardware components, including CPUs, GPUs, TPUs, and other specialized processors.

[0153] The processes and logic flows described in this specification can be executed by one or more programmable computers running one or more computer programs to perform functions by operating on input data and generating output. Additionally, graphics processing units (GPUs) and tensor processing units (TPUs) can be utilized to enable concurrent execution of aspects of these processes and logic flows, significantly accelerating performance. This approach offers significant advantages for computationally intensive tasks often found in Al and machine learning applications, such as matrix multiplications, convolutions, and other operations that exhibit a high degree of parallelism. By leveraging the parallel processing capabilities of GPUs and TPUs, significant speedups and efficiency gains compared to relying solely on CPUs can be achieved. Alternatively or in combination with programmable computers and specialized processors, these processes and logic flows can also be implementedusing specialized processing hardware, such as field-programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs), for even greater performance or energy efficiency in specific use cases.

[0154] Computers capable of executing a computer program can be based on general- purpose microprocessors, special-purpose microprocessors, or a combination of both. They can also utilize any other t pe of central processing unit (CPU). Additionally, graphics processing units (GPUs), tensor processing units (TPUs). and other machine learning accelerators can be employed to enhance performance, particularly for tasks involving artificial intelligence and machine learning. These accelerators often work in conjunction with CPUs, handling specialized computations while the CPU manages overall system operations and other tasks. Typically, a CPU receives instructions and data from read-only memory (ROM), random access memory (RAM), or both. The essential elements of a computer include a CPU for executing instructions and one or more memory devices for storing instructions and data. The specific configuration of processing units and memory will depend on factors like the complexity of the Al model, the volume of data being processed, and the desired performance and latency requirements. Embodiments can be implemented on a wide range of computing platforms, from small embedded devices with limited resources to large-scale data center systems with high-performance computing capabilities. The system may include storage devices like hard drives, SSDs, or flash memory for persistent data storage.

[0155] Computer-readable media suitable for storing computer program instructions and data encompass all forms of non-volatile memory, media, and memory devices. Examples include semiconductor memory' devices such as read-only memory' (ROM), solid-state drives (SSDs), and flash memory devices; hard disk drives (HDDs); optical media; and optical discs such as CDs, DVDs, and Blu-ray discs. The specific type of computer-readable media used will depend on factors such as the size of the data, access speed requirements, cost considerations, and the desired level of portability7or permanence.

[0156] To facilitate user interaction, embodiments of the subject matter described in this specification can be implemented on a computing device equipped with a display device, such as a liquid crystal display (LCD) or an organic light-emitting diode (OLED) display, for presenting information to the user. Input can be provided by the user through various means, including a keyboard), touchscreens, voice commands, gesture recognition, or other input modalities depending on the specific device and application. Additional input methods can include acoustic, speech, or tactile input, while feedback to the user can take the form of visual, auditory, or tactile feedback. Furthermore, computers can interact with users by exchangingdocuments with a user's device or application. This can involve sending web content or data in response to requests or sending and receiving text messages or other forms of messages through mobile devices or messaging platforms. The selection of input and output modalities will depend on the specific application and the desired form of user interaction.

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

[0158] Machine learning models can be implemented and deployed using machine learning frameworks, such as TensorFlow or JAX. These frameworks offer comprehensive tools and libraries that facilitate the development, training, and deployment of machine learning models.

[0159] Embodiments of the subject matter described in this specification can be implemented within a computing system comprising one or more components, depending on the specific application and requirements. These may include a back-end component, such as a back-end server or cloud-based infrastructure; an optional middleware component, such as a middleware server or application programming interface (API), to facilitate communication and data exchange; and a front-end component, such as a client device with a user interface, a web browser, or an app, through which a user can interact with the implemented subject matter. For example, the described functionality could be implemented solely on a client device (e.g., for on-device machine learning) or deployed as a combination of front-end and back-end components for more complex applications. These components, when present, can be interconnected using any form or medium of digital data communication, such as a communication network like a local area network (LAN) or a wide area network (WAN) including the Internet. The specific system architecture and choice of components will depend on factors such as the scale of the application, the need for real-time processing, data security requirements, and the desired user experience.

[0160] The computing system can include clients and servers that may be geographically separated and interact through a communication network. The specific type of network, such as a local area network (LAN), a wide area network (WAN), or the Internet, will depend on the reach and scale of the application. The client-server relationship is established through computer programs running on the respective computers and designed to communicate with each other using appropriate protocols. These protocols may include HTTP, TCP / IP, or other specialized protocols depending on the nature of the data being exchanged and the security requirements of the system. In certain embodiments, a server transmits data or instructions toa user's device, such as a computer, smartphone, or tablet, acting as a client. The client device can then process the received information, display results to the user, and potentially send data or feedback back to the server for further processing or storage. This allows for dynamic interactions between the user and the system, enabling a wide range of applications and functionalities.

[0161] 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.

[0162] 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 be integrated together in a single software product or packaged into multiple software products.

[0163] 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.

[0164] What is claimed is:

Claims

CLAIMS1. A computer-implemented method for generating a response to a query, the method comprising: receiving one or more query tokens, the one or more query tokens indicative of the query, wherein one or more of the query tokens represent image data, video data, or audio data; providing the one or more query tokens as input to a machine learning model; outputting, from a first head of the machine learning model, a first embedding based upon the one or more query tokens; generating an intermediate input embedding based upon the one or more query tokens and the first embedding; outputting, from a second head of the machine learning model, output data based upon the intermediate input embedding; and generating the response to the query based upon the output data.

2. The method of any preceding claim, wherein the first embedding is indicative of rationale for the response.

3. The method of any preceding claim, wherein the first embedding is a soft token not part of a predetermined vocabulary of tokens for selecting at least one of the one or more query’ tokens, the predetermined vocabulary7of tokens defining a space of discrete values, and the soft token comprising one or more continuous values sampled from a latent embedding space of continuous values.

4. The method of any preceding claim, further comprising: generating, for each of the one or more query' tokens, one or more query' embeddings each indicative of the respective query token; and wherein outputting, from the first head of the machine learning model, the first embedding, comprises processing the one or more query embeddings using the machine learning model.

5. The method of claim 4, wherein generating the intermediate input embedding comprises concatenating the one or more query embeddings with the first embedding.

6. The method of any preceding claim, further comprising, for a number, k, of iterations: outputting, from the first head of the machine learning model, a respective second embedding based upon the intermediate input embedding; updating the intermediate input embedding by concatenating the intermediate input embedding with the respective second embedding; and wherein outputting the output data is based upon the updated intermediate input embedding.

7. The method of any of claims 4 to 6, wherein the machine learning model comprises: a tokenizer layer for receiving the query and generating the one or more query tokens; one or more token embedding layers for generating the one or more query embeddings based upon the one or more query' tokens; one or more attention layers configured to process the one or more query embeddings for outputting the first embedding from the first head, the one or more attention layers further configured to process the intermediate input embedding for outputting the output data from the second head; and wherein the first head and the second head are connected to the one or more attention layers.

8. The method of any preceding claim, wherein the first head is a Gaussian head configured to sample the first embedding from a Gaussian distribution, the Gaussian distribution parameterized based upon a latent output embedding generated by the machine learning model in response to providing the one or more query tokens as input to the machine learning model.

9. The method of any preceding claim, wherein the output data is data indicating an output token, and wherein the one or more query tokens and the output token each represent text data, image data, video data, or audio, signal data, sensor data, point-cloud data, or categorical data.

10. The method of any preceding claim, wherein the response to the query is for performing a task, the task comprising a classification task, a control task for a roboticdevice, a machine translation task, a summarization task, a question-answering task, a data generation task, a signaling task, a spatial reasoning task, a programming code generation task, a transcription task, a captioning task, or a dialogue task.

11. The method of claim 10, wherein the task is a task for a real-world environment, and wherein the output data is configured to control one or more mechanical and / or electronic devices for performing the task in the real-world environment.

12. The method of claim 10 or 11, wherein generating the response to the query for performing the task requires abstract or non-linguistic reasoning.

13. The method of any preceding claim, further comprising updating one or more parameters of the machine learning model by computing a loss value based upon the output data.

14. The method of any preceding claim, wherein outputting the first embedding based upon the one or more query tokens further comprises, for a number of iterations: outputting, from the first head of the machine learning model, a tentative embedding based upon the one or more query tokens; and determining the first embedding based upon each tentative embedding using one or more corresponding importance weights.

15. The method of any preceding claim, wherein generating the response based upon the output data comprises a number. 2N, of forward passes using the machine learning model, where N is a positive integer indicating a number of output tokens for the response.

16. A computer-implemented method of generating a neural network output, comprising: obtaining an input sequence of tokens selected from a predetermined vocabulary of tokens, the input sequence of tokens representing a system input, wherein one or more of the input sequence of tokens represents image data, video data, or audio data; at one or more first processing iterations generating a first portion of an output sequence of tokens by:processing the input sequence of tokens and a current output sequence of tokens, using a neural network, to generate a neural network output for selecting a next output token to extend the cunent output sequence of tokens; processing the neural network output using a token generation neural network to generate a soft token that is not part of the predetermined vocabulary of tokens, the soft token comprising one or more continuous values; and extending the current output sequence of tokens by appending the soft token to the current output sequence of tokens; at one or more second processing iterations generating a second portion of the output sequence of tokens by: processing the input sequence of tokens and the current output sequence of tokens including the soft token, using the neural network, to generate the neural network output for selecting a next output token from the predetermined vocabulary of tokens to further extend the current output sequence; and selecting the next output token from the predetermined vocabulary of tokens using the neural network output; and at a final iteration obtaining a system output based upon at least the second portion of the output sequence of tokens.

17. The method of claim 16. wherein the system input defines a reasoning task, and wherein the system output comprises a result of the reasoning task.

18. The method of claim 16 or 17, wherein the neural network output comprises a score distribution that assigns a respective score to each token in the predetermined vocabulary’ of tokens.

19. The method of claim 16 or 17, wherein the neural network output comprises an intermediate output of the neural network from an intermediate layer of the neural network; the method further comprising: processing the neural network output using one or more subsequent layers of the neural network, after the intermediate layer, to generate a score distribution that assigns a respective score to each token in the vocabulary of tokens.

20. The method of any of claims 16 to 19, wherein the neural network has a Transformerbased architecture.

21. The method of any of claims 16 to 20, comprising: obtaining the neural network with a pre-trained set of trainable parameters; obtaining a training dataset comprising one or more training data items, each training data item comprising an example system input and a corresponding example system output; training the neural network and token generation neural network by processing the example system input of a training data item using the method of any of claims 16-20 to generate the system output, and backpropagating gradients for the one or more trainable parameters of the neural network using a training objective function dependent on a difference between the example system output and the corresponding system output through the neural network and token generation neural network to update the trainable parameters of both the neural network and token generation neural network.

22. A computing system comprising: one or more processors; and one or more non-transitory computer-readable media storing computer-readable instructions configured to cause one or more processors to perform a method according to any one of the preceding claims.

23. One or more non-transitory computer-readable media storing computer-readable instructions configured to cause one or more computing devices to perform a method according to any one of claims 1 to 21.

Images