Query routing using self-reflection in generative machine learning models

Abstract

This specification describes a method performed by one or more data processing apparatus. The method comprises obtaining a query and obtaining a subset of contextual data from contextual data relating to the query. The method further comprises processing the query and the subset of contextual data using a generative machine learning model to generate response data relating to the query or to generate an indication that the query cannot be satisfied using the subset of the contextual data. In response to determining that the query cannot be satisfied using the subset of contextual data, the method comprises processing the query and the contextual data relating to the query using the generative machine learning model to generate the response data.

Description

[0001] QUERY ROUTING USING SELF-REFLECTION IN GENERATIVE MACHINE LEARNING MODELS

[0002] BACKGROUND

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

[0004] [2] 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.

[0005] [3] 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.

[0006] SUMMARY

[0007] [4] According to a first aspect, there is provided a method performed by one or more data processing apparatus. The method comprises obtaining a query and obtaining a subset of contextual data from contextual data relating to the query. The method further comprises processing the query and the subset of contextual data using a generative machine learning model to generate response data relating to the query or to generate an indication that the query cannot be satisfied using the subset of the contextual data. In response to determining that the query cannot be satisfied using the subset of contextual data, the method comprises processing the query and the contextual data relating to the query using the generative machine learning model to generate the response data.

[0008] [5] In some implementations, processing the query and the subset of contextual data using the generative machine learning model comprises prompting the generative machine learning model to predict whether the query can be satisfied using the subset of contextual data. For example, the generative machine learning model can be prompted to generate a response to the query if the model predicts that the query can be satisfied using the subset of contextual data or to provide output text indicating that the query cannot be satisfied using the subset of contextual data.

[0009] [6] In some implementations, obtaining the subset of contextual data from contextual data relating to the query comprises selecting the subset of contextual data from the contextual data relating to the query.

[0010] [7] In some implementations, selecting the subset of contextual data is performed by a Retriever machine learning model. A Retriever machine learning model is a type of machine learning model trained to select one or more of the most relevant extracts from contextual data for a given query. The method described herein can be used in conjunction with an off-the- shelf Retriever machine learning model. Examples of suitable Retriever machine learning models include “Contriever,” the details of which can be found in Izacard, Gautier, et al. "Unsuperviscd dense information retrieval with contrastive learning," arXiv:2112.09118 (2021) which is hereby incorporated by reference in its entirety and “DRAGON,” the details of which can be found in Lin, Sheng-Chieh, et al. "How to train your dragon: Diverse augmentation towards generalizable dense retrieval," arXiv:2302.07452 (2023) which is hereby incorporated by reference in its entirety.

[0011] [8] In some implementations, selecting the subset of contextual data is based upon evaluating a distance metric between the query and a plurality of candidate chunks extracted from the contextual data relating to the query. The plurality of candidate chunks can be generated using any suitable method. For example, the contextual data can be divided into overlapping or non-overlapping chunks of uniform size. In some implementations, the chunk size is or is approximately 300 words (or tokens) for text data.

[0012] [9] In some implementations, evaluating the distance metric between the query and the plurality of candidate chunks extracted from the contextual data relating to the query comprises generating a query embedding from the query, generating a respective chunk embedding for each of the plurality of candidate chunks, and evaluating the distance metric between the query embedding and the respective chunk embeddings. In some implementations, the embeddings can be generated using the Retriever machine learning model.

[0013]

[0010] In some implementations, the query embedding is generated by a first encoder. In some implementations, the chunk embedding is generated by a second encoder. In some implementations, the first and second encoders are the same encoder. In other implementations, the first and second encoders are different encoders but jointly trained. In some implementations, the Retriever machine learning model comprises the first and second encoders.

[0014]

[0011] The distance metric can be any suitable distance metric such as a cosine distance. In some implementations, the Retriever machine learning model performs the evaluation of the distance metric.

[0015]

[0012] In some implementations, selecting the subset of contextual data based upon evaluating the distance metric between the query and the plurality of candidate chunks comprises ranking the plurality of candidate chunks based upon the distance metric evaluations and selecting the k highest ranked candidate chunks as the subset of contextual data. That is, the candidate chunks having the greatest similarity to the query can be selected.

[0016]

[0013] In some implementations, the selected k highest ranked chunks are concatenated in rank order to form the subset of contextual data. For example, the subset of contextual data can be a sequence of data with the top ranked chunks appearing first in the sequence.

[0017]

[0014] In some implementations, the value of k is determined based upon the query. In some implementations, the value of k can be determined based upon a query type or a measure of query complexity. For example, if the query requires multi-step reasoning, the value of k can be set to a higher value such as five. On the other hand, if the query can be satisfied with a single piece of data, the value of k can be set to one or another low value.

[0018]

[0015] In other implementations, the value of k is determined based upon processing the query using a machine learning model to predict the value of k. The machine learning model can be the generative machine learning model itself or a different machine learning model trained specifically to analyze the query to predict the value of k.

[0019]

[0016] In some implementations, k is an integer in the range 1 to 5 inclusive.

[0020]

[0017] In some implementations, the generative machine learning model is a large language model (LLM) based machine learning model (e.g. foundation models). This can also include multi-modal models that are capable of processing text input together with other modalities such as image, video and / or audio.

[0021]

[0018] In some implementations, processing the query and the subset of contextual data using the generative machine learning model to generate response data comprises generating an indication that the query is only partially satisfiable using the subset of contextual data. In such implementations, the following steps can be repeated: obtaining a further subset of contextual data, and processing the query, the further subset of contextual data and the previously obtained subsets of contextual data using the generative machine learning model until response data is generated or an indication that the query cannot be satisfied is generated. That is, in some implementations, the generative machine learning model can be configured to predict whether a query is only partially satisfiable using the obtained subset of contextual data. In such cases, further subsets of contextual data can be obtained and processed iteratively by the generative machine learning model until the generative machine learning model is able to generate a response or the generative machine learning model predicts that the query cannot be satisfied.

[0022]

[0019] In some implementations, the indication that the query is only partially satisfiable comprises a reformulated query and the further subset of contextual data is obtained based upon the reformulated query. The further subset of contextual data can be obtained using the methods described above for obtaining the first subset of contextual data but using the reformulated query in place of the original query. The reformulated query can include an indication of the data that is missing for satisfying the query.

[0023]

[0020] In some implementations, the indication that the query is only partially satisfiable comprises a predicted size of the additional contextual data that is necessary to satisfy the query and the further subset of contextual data is obtained based upon the predicted size. For example, the predicted size can be predicted value of k.

[0024]

[0021] In some implementations, the indication that the query is only partially satisfiable comprises partial response data and the generated response data comprises the partial response data. For example, the generative machine learning model can iteratively build-up the response to the query or generate an initial outline of a response with placeholders to be filled-in as more contextual data becomes available.

[0025]

[0022] In some implementations, the indication that the query cannot be satisfied using the subset of contextual data comprises a confidence score. For example, the confidence score can be based upon the probability of an “unanswerable" token being generated or the generative machine learning model can be prompted to provide a confidence score.

[0026]

[0023] In some implementations, the method further comprises obtaining the contextual data relating to the query. In some implementations, obtaining the contextual data relating to the query comprises determining, based upon the query, that contextual data is required to satisfy the query. In some implementations, the generative machine learning model can perform the determination.

[0027]

[0024] In some implementations, determining, based upon the query, that contextual data is required to satisfy the query comprises processing the queiy using the generative machine learning model to generate data for one or more function calls to obtain the contextual data relating to the query.

[0028]

[0025] In some implementations, the one or more function calls comprises at least one function call to an external data source. For example, the external data source can comprise a search engine or a database.

[0029]

[0026] In some implementations, the method further comprises performing the one or more function calls based on the data for the one or more function calls. For example, performing the one or more function calls based on the data for the one or more function calls can comprise determining to perform the function call based on the data for the one or more function calls.

[0030]

[0027] In some implementations, determining to perform the function call comprises parsing the data for the one or more function calls to determine the one or more functions to call. In some implementations, determining to perform the function call comprises parsing the data for the one or more function calls to determine one or more parameters for the one or more functions. For example, the one or more parameters can comprise one or more search terms.

[0031]

[0028] In some implementations, the data for the one or more function calls comprises programming language code, markup language code, and / or query language code.

[0032]

[0029] Thus, in some implementations, the generative machine learning model can perform an initial analysis of the query and determine that external contextual data is required in order to satisfy the query. The generative machine learning model can select an appropriate data source or tool for obtaining the contextual data and can generate instructions for obtaining the contextual data from the external data source or tool in the appropriate form. This can include one or more function calls to a particular API of the external data source or tool and the generative machine learning model can generate the necessary parameters for the one or more function calls.

[0033]

[0030] In some implementations, the query comprises a question and the generated response data comprises an answer to the question. In some implementations, the generated response data can comprise computer program code that is configured to carry out a function as specified in the query.

[0034]

[0031] In some implementations, the generated response data can comprise data of different modalities such as text, image, video and / or audio data. Image and video data can comprise a plurality of pixel values or visual tokens that can be decoded into pixel values. Audio data can comprise an audio signal such as digital samples of a waveform or be encoded according to any suitable audio encoding format. The query can include natural language instructions for generating the image, video and / or audio data. For example, the queiy can specify the particular objects and their attributes that are to be present in the generated image, video and / or audio data.

[0035]

[0032] In some implementations, the query can be a question in relation to image, video and / or audio data. For example, the queiy can involve identifying particular objects and / or actions being performed in the image, video and / or audio data. The generated response data can comprise an appropriate text response and / or an appropriately labelled image / video, e.g. by overlaying bounding boxes, and / or an appropriate extract from the image, video and / or audio data. The generated response data can comprise a modification to the image / video / audio data as specified in the queiy such as to replace a particular object with another object or to replace the background or to enhance the lighting amongst others. Thus, the generative machine learning model can perform a variety of tasks including image / video / audio generation, classification, detection, recognition, prediction, segmentation, captioning, de-noising, and infilling amongst others.

[0036]

[0033] In some implementations, the generative machine learning model can perform a natural language processing or understanding task. For example, the task may be an entailment task, a paraphrase task, a textual similarity task, a sentiment task, a sentence completion task, a grammar task or other similar task. In some implementations, the generative machine learning model can perform a machine translation task. The generative machine learning model can generate response data that comprises a translation of text from a source language to a target language as specified by the query.

[0037]

[0034] In some implementations, the generative machine learning model can generate response data that comprises instructions or other data for controlling a device such as a user device or loT devices such as smart lighting. In some implementations, the generative machine learning model can generate response data that comprises instructions or other data for controlling an electro-mechanical agent such as a physical robot operating in a real-world environment.

[0038]

[0035] In some implementations, the query is obtained from a user device and the response data is transmitted to the user device.

[0039]

[0036] The generative machine learning model can be trained using any suitable training method. For example, the generative machine learning model can be pre-trained using unsupervised or self-supervised learning on large scale training data obtained from various sources such as web data. The pre-trained model can undergo supervised fine-tuning in order to enhance the model’s ability to follow user instructions and prompts and to carry out downstream tasks.

[0040]

[0037] According to a second aspect, there is provided a method performed by one or more data processing apparatus. The method comprises obtaining a query and obtaining a subset of contextual data from contextual data relating to the query. The method further comprises processing the query and the subset of contextual data using a generative machine learning model. The processing comprises generating an indication that the queiy is only partially satisfiable using the subset of contextual data and repeating: obtaining a further subset of contextual data; and processing the query, the further subset of contextual data and the previously obtained subsets of contextual data using the generative machine learning model until response data is generated or an indication that the query cannot be satisfied is generated.

[0041]

[0038] In some implementations, the indication that the query is only partially satisfiable comprises a reformulated query; and the further subset of contextual data is obtained based upon the reformulated query.

[0039] In some implementations, the indication that the query is only partially satisfiable comprises a predicted size of the additional contextual data that is necessary to satisfy the query; and the further subset of contextual data is obtained based upon the predicted size.

[0042]

[0040] In some implementations, the indication that the query is only partially satisfiable comprises partial response data; and the generated response data comprises the partial response data.

[0043]

[0041] According to a third aspect, there is provided a method performed by one or more data processing apparatus. The method comprises obtaining a query and processing the query using a first machine learning model to predict a value of k. The predicted value of k is determined based upon the query and k represents a quantity of contextual data chunks to be selected. The method further comprises selecting k chunks of contextual data from a larger body of contextual data based on the predicted value of k; and processing the query and the k selected chunks using a second machine learning model to generate response data.

[0044]

[0042] In some implementations, determining the predicted value of k based upon the query comprises determining the value of k based on a query type or a measure of query complexity.

[0045]

[0043] In some implementations, the first machine learning model and the second machine learning model are the same generative machine learning model.

[0046]

[0044] In some implementations, the first machine learning model is a different machine learning model than the second machine learning model, and the first machine learning model is specifically trained to analyze the query to predict the value of k.

[0047]

[0045] In some implementations, selecting the k chunks of contextual data comprises ranking a plurality of candidate chunks from the larger body of contextual data based on a distance metric evaluated between the query and each of the plurality of candidate chunks; and selecting the k highest ranked candidate chunks.

[0048]

[0046] In some implementations, the method further comprises concatenating the selected k highest ranked chunks in rank order to form a single subset of contextual data for processing by the second machine learning model.

[0049]

[0047] In some implementations, k is an integer in the range 1 to 5 inclusive.

[0050]

[0048] According to a fourth aspect, there is provided a system comprising one or more data processing apparatus and a memory. The memory stores instruction that when executed by the one or more data processing apparatus, causes the one or more data processing apparatus to carry out a method according to any of the first to third aspects.

[0051]

[0049] According to a fifth aspect, there is provided a non-transitory computer-readable storage medium comprising instructions that when executed by one or more data processing apparatus cause the one or more data processing apparatus to carry out a method according to any of the first to third aspects.

[0052]

[0050] According to a sixth aspect, there is provided a computer program product comprising instructions that when executed by one or more data processing apparatus cause the one or more data processing apparatus to carry out a method according to any of the first to third aspects.

[0053]

[0051] It will be appreciated that features described in the context of one aspect may be combined with features of one or more other aspects.

[0054]

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

[0055]

[0053] The techniques described herein can process a query using a generative machine learning model with contextual data in a memory and computationally efficient manner. For some queries that require contextual data in order to be satisfied, it may be possible to satisfy the query using only part of the contextual data. The techniques described herein uses the generative machine learning model to generate a response using a subset of the contextual data if the generative machine learning model predicts it is possible to do so. Otherwise, the generative machine learning model can output a prediction that the query cannot be satisfied using the subset of the contextual data and can fallback to processing the full contextual data to generate a response to the query. As generative machine learning models such as LLMs and foundation models typically comprise one or more attention layers which scale quadratically with input size, a reduction in the amount of data that needs to be processed for a query can greatly reduce the amount of memory and other computational resources required. In some state-of-the-art generative machine learning models, the contextual window size can be up to one million tokens. Using the techniques described herein, queries can be satisfied using a subset of contextual data comprising approximately 2000 tokens. This reduction in the amount of data processed directly translates into a tangible improvement in the functioning of the underlying data processing apparatus. Specifically, by reducing the input data length, the number of floating-point operations (FLOPs) required for attention mechanism calculations is decreased. This results in lower processor cycle consumption and reduced power usage per query. The system can achieve higher throughput of queries for a given set of hardware resources.

[0056]

[0054] The generative machine learning model can be part of a system that services many user queries from a plurality of users. As such, reducing the memory consumption and computational requirements for even some of those queries can be significant.

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

[0057] BRIEF DESCRIPTION OF THE DRAWINGS

[0058]

[0056] Figure 1 is a schematic block diagram illustrating an example of a data generation system.

[0059]

[0057] Figure 2 is a flowchart illustrating an example method for generating data according to an implementation.

[0060]

[0058] Figures 3 A and 3B provide results of benchmarking experiments.

[0061]

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

[0062] DETAILED DESCRIPTION

[0063]

[0060] Generative machine learning models such as large language models (LLMs), multimodal language models and foundation models are becoming increasing powerful. With each new generation of model, the number of tokens that can be processed increases, sometimes exponentially, with some state-of-the-art models having a context window size of up to one million tokens. Typically, a generative model comprises one or more attention layers which scale quadratically with input size. As such, an increasing context window size can greatly increase memory resource requirements.

[0064]

[0061] In some cases, where a query requires contextual data in order to provide a response to the query, it is possible that the query can be satisfied using only a subset of that contextual data. Techniques such as Retrieval Augmented Generation (RAG) have been developed in order to attempt to extract the most relevant parts of the contextual data based upon a similarity with the query. However, as the Retriever models used in RAG are less powerful than LLMs and the difficulty in quantifiably measuring relevance, the selected subset of contextual data provided by a Retriever can in some cases hinder response generation if the selected subset is not of relevance.

[0065]

[0062] The techniques described herein use model self-reflection to predict the viability of using RAG or if it is necessary to use the full context to satisfy the query. That is, the generative machine learning model can be prompted to predict whether a given subset of contextual data can be used to satisfy a query and if so, to generate a response using the subset of contextual data. Otherwise, the generative machine learning model can output a prediction that the query cannot be satisfied using the subset of contextual data and can fallback to processing the full contextual data to generate a response to query. In this way, a query can be routed to processing using RAG or processing using the full context.

[0066]

[0063] In some cases, the model can determine that a query can be partially satisfied and iteratively obtain further subsets of the contextual data until it is predicted that the query can be satisfied.

[0067]

[0064] The techniques described herein attempt to generate a response to a query using contextual data in the most memory efficient way possible. In some implementations, the generative machine learning model can be part of a system that services many user queries from a plurality of users. As such, reducing the memory consumption and computational requirements for even some of those queries can be significant. The techniques described herein addresses the technical problem of dynamically allocating computational resources on a per-query basis, thereby improving the computational efficiency of the computer system without degrading the technical quality of the generated output.

[0068]

[0065] Referring now to Figure 1, a data generation system 100 is shown. The system 100 can be implemented using one or more data processing apparatus in one or more locations. The data processing apparatus can include any appropriate hardware such as a personal computer, a server, a laptop, a mobile device or more specifically any type of processing unit such as a CPU, GPU, TPU or specialized hardware apparatus such as an FPGA or ASIC.

[0069]

[0066] The system 100 comprises a generative machine learning model 101. The generative machine learning model 101 can be any appropriate type of generative machine learning model. The generative machine learning model 101 can comprise one or more neural network layers such as attention-based neural network layers. The generative machine learning model can, for example, be configured according to a Transformer-based architecture. The generative machine learning model 101 can be autoregressive or non-autoregressive and can comprise an encoder-decoder architecture, an encoder-only architecture, or a decoder-only architecture. In some implementations, the generative machine learning model 101 is a large language model (LLM) based machine learning model (e.g. foundation models). This can also include multimodal models that are capable of processing text input together with other modalities such as image, video and / or audio to generate text, image, video and / or audio data. The generative machine learning model 101 can comprise diffusion-based models. The generative machine learning model 101 can comprise a plurality of sub-networks or expert networks.

[0070]

[0067] The system 100 is configured to obtain a query 102. The query 102 can comprise text in natural language. The query 102 can comprise a question on a topic for the system 100 to provide an answer to or the query 102 can comprise instructions for carrying out a task. The query 102 can also include data of other modalities such as image, video and / or audio data. For example, the query 102 can be a question relating to what is depicted in an image or what is being spoken in a recording of speech. Further examples of tasks that the system 100 can be configured to perform are provided below.

[0071]

[0068] The query7102 can be transmitted to the system 100 using any suitable network connection from a user device or the system 100 can comprise a user interface for users to interact with the system 100.

[0072]

[0069] The system 100 is further configured to obtain a subset 103 of contextual data from contextual data 104 relating to the query 102. In some implementations, the contextual data 104 can be a document or other electronic resource obtained from an external data source such as a database or obtained through an Internet search engine or other appropriate tool. In some implementations, the contextual data or data sources can be indicated or provided by the user. The contextual data 104 can include, but is not limited to, a single contiguous resource; a collection of multiple discrete documents, records, or data extracts; and / or a collection of multiple extracts, e.g. snippets from one or more documents. For example, when contextual data is obtained via a search engine, the contextual data may comprise a set of retneved results (e g., a plurality of website snippets or links to documents). In such a case, the subset 103 of contextual data may be a selection of one or more of these retrieved results (e.g., the top-ranked snippets), and the fallback step of processing the contextual data would involve processing the set of retrieved results.

[0073]

[0070] In some implementations, the system 100 is configured to determine, based upon the query' 102 that contextual data is required to satisfy the query' 102. For example, the vocabulary of tokens used by the generative machine learning model 101 can include special tokens to indicate the use of external resources and tools. The generative machine learning model 101 can process the query 102 and emit the corresponding token(s) required for accessing the external resource / tool. In another example, the generative machine learning model 101 can be configured to operate in a “chain-of-thought” type mode and to determine a series of sub-steps needed to generate a response to the query 102. The generative machine learning model 101 can determine that contextual data is required to satisfy the query 102 in these sub-steps.

[0074]

[0071] The generative machine learning model 101 can generate data for one or more function calls to an appropriate data source / tool to obtain the contextual data 104. For example, the generated data can include parameters, programming language code, markup language code and / or query language code for the function call according to an API of the data source / tool. The one or more parameters can include search terms for querying the data source or search engine. The system 100 can be configured to perform the one or more function calls based on the generated data.

[0075]

[0072] Thus, in some implementations, the generative machine learning model 101 can select an appropriate data source or tool for obtaining the contextual data 104 and can generate instructions for obtaining the contextual data 104 from the external data source or tool in the appropriate form.

[0076]

[0073] The system 100 can be configured to select the subset 103 of contextual data. For example, a plurality of candidate chunks can be extracted from the contextual data 104 and a distance metric evaluated between the query 102 and each of the candidate chunks. The distance metric can be any appropriate distance metric such as a cosine similarity. The candidate chunks can be generated using any appropriate method. For example, the contextual data can be divided into fixed size chunks which can be overlapping or non-overlapping. In another example, chunks can be randomly sampled from the contextual data 104. In one example, each chunk is 300 tokens (or words) in length, however it will be appreciated that any appropriate size can be used.

[0077]

[0074] In some implementations, the distance metric is based upon an embedding of the query and an embedding of each chunk. That is, the system 100 can be configured to generate a query embedding from the query' 102 and to generate a respective chunk embedding for each of the plurality of the candidate chunks. The distance metric can be evaluated between the query embedding and the respective chunk embeddings. The embeddings can be generated using any appropriate encoder machine learning model. In some implementations, the same encoder can be used for generating the query' and chunk embeddings whilst in other implementations, the query and chunk encoders are different. Typically, the query' and chunk encoders have been jointly trained, for example, using a contrastive loss function that encourages query and chunk embeddings from the same document to be close together and query' and chunk embeddings from different documents to be further apart according to the distance metric.

[0078]

[0075] In some implementations, the system 100 further comprises a Retriever machine learning model 109 as used in Retrieval Augmented Generation (RAG) techniques. The Retriever machine learning model 109 can be configured to perform at least part of the operations for selecting a subset 103 of the contextual data. For example, the query' and chunk encoders can be part of the Retriever machine learning model 109. The techniques described herein can be used with an off-the-shelf Retriever machine learning model such as the Contriever and Dragon models cited above.

[0076] In some implementations, the system 100 is further configured to rank the plurality' of candidate chunks based upon the distance metric evaluations and to select the top-£ highest ranked candidate chunks as the subset 103 of contextual data. In some implementations, k =5 is used. In other implementations, k is between 1 and 5 inclusive. It will be appreciated that any appropriate value of k can be used.

[0079]

[0077] In some implementations, the value of k can be determined based upon the query such as a query type or query complexity. For example, a more complex query’ or a query that requires multi-step reasoning may require more contextual data and hence, a larger value of k, compared to a query that may only require a single additional piece of data. In some implementations, the generative machine learning model 101 can process the query 102 to predict the value of k required, for example, when determining that contextual data is required, the generative machine learning model 101 can also provide a predicted value of k. In other implementations, a separate machine learning model can be used predict a value of k for the query 102 or to classify the query 102, e.g. classify its complexity, for determining the value of k. For example, this can be a distilled version of the generative machine learning model 101 that has been fine-tuned for this specific task.

[0080]

[0078] In some implementations, the system 100 is configured to concatenate the selected top- k highest ranked chunks in rank order to form the subset 103 of contextual data. That is, the top ranked chunk can appear first.

[0081]

[0079] The system 100 is configured to cany’ out a first generation attempt 105 using the subset 1 3 of contextual data. In this regard, the system 100 is configured to process the query 102 and the subset 103 of contextual data using the generative machine learning model 101 to generate response data 106 relating to the query' 102 or to generate an indication 107 that the query cannot be satisfied using the subset 103 of contextual data. For example, the generative machine learning model 101 can be prompted to generate a response to the query 102 using the subset 103 of contextual data or to provide an output indicating that the query’ 102 cannot be satisfied using the subset 103 of contextual data. Typically, a generative machine learning model that is capable of generating responses to queries also has the capability to evaluate whether a particular subset of contextual data is relevant for responding to a query. Otherwise, a generative machine learning model can be adapted through fine-tuning methods such as supervised fine-tuning using an appropriate dataset of training examples and / or using methods such as reinforcement learning with human feedback or direct preference optimization to carry out the evaluation. Training of generative machine learning models is described in more detail beloyv.

[0080] The indication 107 that the query cannot be satisfied using the subset 103 of contextual data can comprise a designated "unanswerable" token or a sequence of tokens to that effect. In some implementations, the indication 107 that the query 102 cannot be satisfied using the subset 103 of contextual data comprises a confidence score. For example, the confidence score can be based upon the probability of an “unanswerable” token being generated or the generative machine learning model 101 can be prompted to provide a confidence score.

[0082]

[0081] If the first generation attempt fails, the system 100 is configured to carry out a long context generation 108. In this regard, in response to determining that the query 102 cannot be satisfied using the subset 103 of contextual data, the system 100 is configured to process the query 102 and the contextual data 104 relating to the query 102 to generate the response data 106. The amount of contextual data 104 used in long context generation 108 is larger than the subset 103 of contextual data used in the first generation attempt 105. Typically, the full contextual data 104 is used, or as much as the generative machine learning model 101 has capacity7for at the time. This fallback to long context generation 108 serves as a technical failsafe mechanism. It ensures that the system maintains a high level of technical performance, for example, in terms of accuracy or completeness of the response, for difficult queries where the initial subset w as insufficient. By providing this fallback path, the system mitigates the primary technical drawback of standard RAG systems, which can fail completely on such queries. The overall system therefore provides an improved trade-off between computational efficiency and the reliability required for its technical application.

[0083]

[0082] In some implementations, the system 100 can be configured to iteratively obtain further subsets of contextual data if deemed necessary. For example, the generative machine learning model 101 can be configured to generate an indication that the query is partially satisfiable using the subset 103 of contextual data. The system 100 can be configured to obtain further subsets of contextual data until the generative machine learning model 101 is able to generate response data using the combined subsets of response data or an indication that the query cannot be satisfied is generated. In the latter case, the system 100 can proceed to use long context generation. In some implementations, the system 100 can be configured to proceed to use long context generation if response data cannot be generated after a maximum number of iterations. This iterative approach provides a further technical advantage by enabling a graduated consumption of computational resources. Instead of a binary' choice between a small, fixed-size subset and the full context, the system can incrementally increase the context size, and thus the computational load, only as much as is determined to be necessary. This leads to a more finely -tuned optimization of processing resources compared to a simple binary fallback system.

[0084]

[0083] In some implementations, the system 100 can be configured to provide the next k chunks in the ranking of candidate chunks after the top- / < candidate chunks as the further subset of contextual data. In some implementations, the indication that the query is only partially satisfiable can include a predicted size of the additional contextual data that is necessary’ to satisfy the query and the further subset of contextual data is obtained based upon the predicted size. For example, the generative machine learning model 101 can predict / request n additional chunks to be provided. In this regard, the system 100 can be configured to provide the next n chunks in the ranking of candidate chunks after the to A candidate chunks as the further subset of contextual data.

[0085]

[0084] In some implementations, the indication that the query is only partially satisfiable can include a reformulated query and the further subset of contextual data can be obtained based upon the reformulated query. For example, the generative machine learning model 101 can determine what data is missing from the subset 103 of contextual data and to provide an appropriate reformulated query for querying against the obtained contextual data 104 or for obtaining further contextual data from available sources.

[0086]

[0085] In some implementations, the generative machine learning model 101 can iteratively build-up the response to the query, for example when using a "chain-of-lhoughl" type mode or generate an initial outline of a response with placeholders to be filled-in as more contextual data becomes available. The indication that the query is only partially satisfiable can comprise partial response data and the generated response data 106 comprises the partial response data.

[0087]

[0086] Where the query 102 was received remotely from a user device, the system 100 can be configured to transmit the generated response data 106 back to the user device. Alternatively, if the query' 102 was directly input to the system 100, the system 100 can cause the generated response data 106 to be displayed to the user.

[0088]

[0087] Figure 2 is a flow diagram illustrating an example method 200 for generating data. The processing shown in Figure 2 can be carried out using the data generation system 100 of Figure 1.

[0089]

[0088] At step 201, a query is obtained. As described above, the query can comprise text in natural language. The query' can comprise a question on a topic or the query' can comprise instructions for carrying out a task. The query can also include data of other modalities such as image, video and / or audio data. For example, the query can be a question relating to what is depicted in an image or what is being spoken in a recording of speech. Further examples of tasks that can be performed are provided below.

[0090]

[0089] At step 202, a subset of contextual data from contextual data relating to the query is obtained. As described above, the contextual data can be a document or other electronic resource obtained from an external data source such as a database or obtained through an Internet search engine or other appropriate tool. In some implementations, the contextual data or data sources can be indicated or provided by the user.

[0091]

[0090] In some implementations, obtaining the subset of contextual data from contextual data relating to the query comprises selecting the subset of contextual data from the contextual data relating to the query. For example, selecting the subset of contextual data can be based upon evaluating a distance metric between the query and a plurality of candidate chunks extracted from the contextual data relating to the query. The distance metric can be based upon a cosine similarity or any other appropriate distance metric.

[0092]

[0091] In some implementations, evaluating the distance metric between the query and the plurality of candidate chunks extracted from the contextual data relating to the query comprises generating a query embedding from the query, generating a respective chunk embedding for each of the plurality of candidate chunks, and evaluating the distance metric between the query embedding and the respective chunk embeddings. As described above, the embeddings can be generated using any appropriate encoder machine learning model. In some implementations, the query embedding is generated by a first encoder, the chunk embedding is generated by a second encoder, and the first and second encoders are the same encoder. In other implementations, the first and second encoders are different encoders but have been jointly trained. For example, the first and second encoders can be trained using a contrastive loss function that encourages query and chunk embeddings from the same document to be close together and query and chunk embeddings from different documents to be further apart according to the distance metric.

[0093]

[0092] The plurality of candidate chunks can be generated using any appropriate method. For example, the contextual data can be divided into fixed size chunks which can be overlapping or non-overlapping. In another example, chunks can be randomly sampled from the contextual data. In one example, each chunk is 300 tokens (or words) in length, however it will be appreciated that any appropriate size can be used.

[0094]

[0093] As described above, at least some or all of the operations for selecting the subset of contextual data can be performed by a Retriever machine learning model as used in a Retrieval Augmented Generation (RAG) techniques. For example, the Retriever machine learning model can generate the query and chunk embeddings. Thus, the first and second encoders can be part of the Retriever machine learning model.

[0095]

[0094] In some implementations, selecting the subset of contextual data based upon evaluating the distance metric between the query and the plurality of candidate chunks comprises ranking the plurality7of candidate chunks based upon the distance metric evaluations selecting the k highest ranked candidate chunks as the subset of contextual data. In some implementations, k =5 is used. More generally, in some implementations, k can be a value between 1 and 5 inclusive. It will be appreciated however that any appropriate value of k can be used.

[0096]

[0095] As described above, in some implementations, the value of k can be determined based upon the query such as a query type or a measure of query complexity. For example, a more complex query or a query that requires multi-step reasoning may require more contextual data and hence, a larger value of k, compared to a query7that may only require a single additional piece of data. In some implementations, the generative machine learning model can process the query to predict the value of k required, for example, when determining that contextual data is required, the generative machine learning model can also provide a predicted value of k. In other implementations, a separate machine learning model can be used predict a value of k for the query7or to classify the query, e.g. classify its complexify, for determining the value of k. For example, this can be a distilled version of the generative machine learning model that has been fine-tuned for this specific task.

[0097] 196 ] In some implementations, the selected k highest ranked chunks are concatenated in rank order to form the subset of contextual data.

[0098]

[0097] At step 203, the query and the subset of contextual data is processed using a generative machine learning model. From the processing, the generative machine learning model either generates response data (outcome 203a), or an indication that the query cannot be satisfied using the subset of the contextual data (outcome 203b). In some implementations, the generative machine learning model is prompted to predict whether the query can be satisfied using the subset of contextual data. The indication that the query cannot be satisfied using the subset of contextual data can comprise a designated “unanswerable” token or a sequence of tokens to that effect. In some implementations, the indication that the query cannot be satisfied using the subset of contextual data comprises a confidence score. For example, the confidence score can be based upon the probability' of an “unanswerable” token being generated or the generative machine learning model can be prompted to provide a confidence score.

[0099]

[0098] At step 204, in response to determining that the query cannot be satisfied using the subset of contextual data (outcome 203b), the query and the contextual data relating to the query is processed using the generative machine learning model to generate the response data. The amount of contextual data used in step 204 is larger than the subset of contextual data used in step 203. Typically, the full contextual data is used, or as much as the generative machine learning model has capacity for at the time.

[0100]

[0099] As described above, in some implementations, the method further comprises determining, based upon the query, that contextual data is required to satisfy the query’ is performed by the generative machine learning model. For example, the vocabulary of tokens used by the generative machine learning model can include special tokens to indicate the use of external resources and tools. The generative machine learning model can process the query and emit the corresponding token(s) required for accessing the external resource / tool. In another example, the generative machine learning model can be configured to operate in a “chain-of-thought” type mode and to determine a series of sub-steps needed to generate a response to the query. The generative machine learning model can determine that contextual data is required to satisfy the query' in these sub-steps.

[0101]

[0100] In some implementations, the query’ can be processed using the generative machine learning model to generate data for one or more function calls to obtain the contextual data relating to the query. The one or more function calls can comprise at least one function call to an external data source. The external data source can comprise a search engine or a database. The method can further comprise performing or determining to perform the one or more function calls based on the data for the one or more function calls. For example, determining to perform the function call can comprise parsing the data for the one or more function calls to determine the one or more functions to call and / or to determine one or more parameters for the one or more function calls. In some implementations, the one or more parameters can comprise one or more search terms. In some implementations, the data for the one or more function calls comprises programming language code, markup language code, and / or query language code. The generated data can be generated in a format according to an API of the data source / tool.

[0102]

[0101] As described above, in some implementations, processing the query’ and the subset of contextual data using the generative machine learning model to generate response data can be an iterative process. For example, the generative machine learning model can generate an indication that the query' is only partially satisfiable using the subset of contextual data. In response, a further subset of contextual data can be obtained and the further subset of contextual data and the previously obtained subset(s) of contextual data can be processed by the generative machine learning model.

[0102] In some implementations, the indication that the query' is only partially satisfiable can include a reformulated query and the further subset of contextual data can be obtained based upon the reformulated query. For example, the generative machine learning model can determine what data is missing from the subset of contextual data and to provide an appropriate reformulated query for query ing against the obtained contextual data or for obtaining further contextual data from available sources.

[0103]

[0103] In some implementations, the next k chunks in the ranking of candidate chunks after the topA candidate chunks can be provided as the further subset of contextual data. In some implementations, the indication that the query is only partially satisfiable can include a predicted size of the additional contextual data that is necessary to satisfy the query' and the further subset of contextual data is obtained based upon the predicted size. For example, the generative machine learning model can predict / request n additional chunks to be provided. In this regard, the next n chunks in the ranking of candidate chunks after the top- candidate chunks can be provided as the further subset of contextual data.

[0104]

[0104] In some implementations, the indication that the query is only partially satisfiable can comprise partial response data and the generated response data comprises the partial response data.

[0105]

[0105] The further subset of contextual data and previously obtained subset(s) of contextual data can be processed using the generative machine learning model. The generative machine learning model can provide response data to the query if the combined subsets of contextual data are sufficient or if further subsets of contextual data are still required, the generative machine learning model can generate another indication of partial satisfiability. This can be repeated until either response data for satisfying the query is generated or an indication that the query cannot be satisfied is generated. In some implementations, there can be a maximum number of iterations that are performed. If the generative machine learning model cannot generate response data using the subsets or within the iteration limit, the method can proceed to generating using the full contextual data as in step 204.

[0106]

[0106] Referring now to Figures 3A and 3B, two tables of results are shown from experiments to evaluate the techniques described herein. Seven datasets from LongBench (Y. Bai et al.. “Longbench: A bilingual, multitask benchmark for long context understanding,” arXiv preprint arXiv:2308.14508, 2023): NarrativeQA, Qasper, MultiFieldQA, HotpotQA,

[0107] 2WikiMultihopQA, MuSiQue, and QMSum; and two datasets from InfinityBench (X. Zhang et al. “Infinity bench: Extending long context evaluation beyond 100k tokens.” arXiv preprint arXiv:2402.13718, 2024): En.QA and En.MC; were used in the experiments. In the table of Figure 3A, the “Contriever” Retriever machine learning model is used in conjunction with the Gemini 1.5-Pro (Google), GPT-4O (OpenAI), and GPT-3.5-Turbo (OpenAI) generative machine learning models. For each of the generative models in the table, row *-l LC shows the results when using standard long context generation, row *-2 RAG shows the results when using standard RAG generation, row *-3 SELF-ROUTE shows the results when using the techniques described herein, row *-4 answerable % shows the percentage of the queries that the generative machine learning model under SELF-ROUTE determined that using the first generation attempt with only subset of contextual data was sufficient, row *-5 token % shows the percentage of tokens used under SELF-ROUTE compared to standard long context generation. In the table of Figure 3B, results are shown for the "Dragon" Retriever machine learning model in combination with the Gemini 1.5-Pro generative machine learning model. The numbered rows represent the same data as in table of Figure 3 A, however row 3 is labelled as '‘combine” instead of “SELF-ROUTE”, row 4 is labelled as “RAG ratio” instead of “answerable %” and row 5 is labelled as “Token ratio” instead of “token %”.

[0108]

[0107] As can be seen from both tables, the first generation attempt (row 4) is consistently used for all three generative machine learning models for more than half of all queries with overall performance comparable to using standard long context generation. Due to the high answerable rate, the number of tokens required is significantly reduced. Since the computational cost of the Transformer-based models is quadratic with respect to token count, this results in substantial reductions in computational cost (e.g. processing and memory) with only a small impact on task performance. As such, the techniques described herein can significantly reduce the computational cost for responding to a batch of queries.

[0109]

[0108] Example tasks that can be carried out by the generative machine learning model will be described in more detail below. In some implementations, the query comprises a question and the generated response data comprises an answer to the question. The generated response data can comprise text in natural language.

[0110]

[0109] In some implementations, the generated response data can comprise data of different modalities such as text, image, video and / or audio data. Image and video data can comprise a plurality of pixel values or visual tokens that can be decoded into pixel values. Audio data can comprise an audio signal such as digital samples of a waveform or be encoded according to any suitable audio encoding format. The query7can include natural language instructions for generating the image, video and / or audio data. For example, the query can specify the particular objects and their attributes that are to be present in the generated image, video and / or audio data. Alternatively, the generated response data can include image, video and / or audio data without necessarily being prompted if it is deemed useful by the model.

[0111] [HO] In some implementations, the query can be a question in relation to image, video and / or audio data. For example, the query can involve identifying particular objects and / or actions being performed in the image, video and / or audio data. The generated response data can comprise an appropriate text response and / or an appropriately labelled image / video, e.g. by overlaying bounding boxes, and / or an appropriate extract from the image, video and / or audio data. The generated response data can comprise a modification to the image / video / audio data as specified in the query such as to replace a particular object with another object or to replace the background or to enhance the lighting amongst others. Thus, the generative machine learning model can perform a variety of tasks including image / video / audio generation, classification, detection, recognition, prediction, segmentation, captioning, de-noising, and infilling amongst others.

[0112] [Hl] In some implementations, the generative machine learning model can perform a natural language processing or understanding task. For example, the task may be an entailment task, a paraphrase / summarization task, a textual similarity task, a sentiment task, a sentence completion task, a grammar task or other similar task. In some implementations, the generative machine learning model can perform a machine translation task. The generative machine learning model can generate response data that comprises a translation of text from a source language to a target language as specified by the query.

[0113]

[0112] In some implementations, the generative machine learning model can generate response data that comprises instructions or other data for controlling a device such as a user device or loT devices such as smart lighting. In some implementations, the generative machine learning model can generate response data that comprises instructions or other data for controlling an electro-mechanical agent such as a physical robot operating in a real-world environment. In some implementations, the query can relate to a programming task and the response data can comprise computer program code to perform the specified task.

[0114] [H3] As described above, the generative machine learning model can be a large language model-type neural network. Further details regarding LLMs will now be described.

[0115]

[0114] An LLM can be an auto-regressive generative machine learning model that generates each token in the output sequence conditioned on the preceding tokens in the output sequence and at least some of the tokens in the input sequence. For example, the LLM can be configured to process an input sequence of tokens from a vocabulary of tokens to generate an output sequence of tokens from the vocabulary. [H5] More generally, the generative machine learning model can be any appropriate neural network that receives an input sequence made up of tokens selected from a vocabulary and auto-regressively generates an output sequence made up of tokens from the vocabulary. For example, the generative machine learning model can be a Transformer-based model neural network or a recurrent neural network-based neural network.

[0116] [H6] In some situations, the generative machine learning model can be referred to as an autoregressive neural network when the neural network used to implement the model auto- regressively generates an output sequence of tokens. More specifically, the auto-regressively generated output is created by generating each particular token in the output sequence conditioned on a current input sequence that includes any tokens that precede the particular text token in the output sequence, i.e., the tokens that have already been generated for any previous positions in the output sequence that precede the particular position of the particular token, and a context input that provides context for the output sequence.

[0117] [H7] For example, the current input sequence when generating a token at any given position in the output sequence can include the input sequence and the tokens at any preceding positions that precede the given position in the output sequence. As a particular example, the current input sequence can include the input sequence followed by the tokens at any preceding positions that precede the given position in the output sequence. Optionally, the input and the current output sequence can be separated by one or more predetermined tokens within the current input sequence.

[0118] [H8] More specifically, to generate a particular token at a particular position within an output sequence, the generative machine learning model can process the current input sequence to generate a score distribution (e.g., a probability distribution) that assigns a respective score, e.g., a respective probability, to each token in the vocabulary of tokens. The generative machine learning model can then select, as the particular token, a token from the vocabulary using the score distribution. For example, the neural network can greedily select the highest-scoring token or can sample, e.g., using nucleus sampling or another sampling technique, a token from the distribution.

[0119]

[0119] As a particular example, the generative machine learning model can be an autoregressive Transformer-based neural netw ork that includes (i) a plurality of attention blocks, at least some of which apply a self-attention operation and (ii) an output subnetwork that processes an output of the last attention block to generate the score distribution.

[0120]

[0120] The generative machine learning model can have any of a variety of Transformer-based neural network architectures. 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. Sifre, 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 open-domain 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.

[0121]

[0121] The generative machine learning model can undergo a first phase of pre-training followed by a second phase of fine-tuning. In general, a generative machine learning model such as an LLM can be pre-trained on large amounts of data including data from, but not limited to, webpages, electronic books, software code, electronic news articles, and machine translation data. The generative machine learning model can be pre-trained using unsupervised or self-supervised learning. For example, the generative machine learning model can be pretrained on a next token prediction task and / or a masked token prediction task. Pre-training on large quantities of diverse data can provide the generative machine learning model with remarkable natural language reasoning capabilities.

[0122]

[0122] Following pre-training, the generative machine learning model can undergo fine-tuning to improve the model’s ability to respond to user prompts and queries. Example types of fine- tuning techniques are supervised fine-tuning (SFT), reinforcement learning with human feedback (RLHF) and direct preference optimization (DPO).

[0123]

[0123] In SFT, a high-quality dataset including examples of input prompts and corresponding responses can be used. This data is ty pically generated by human annotators. The generative machine learning model can be trained using supervised learning to generate the corresponding responses from the input prompt. SFT requires a much smaller amount of data that is in pretraining.

[0124]

[0124] In RLHF, a reward model can be trained from human preference data regarding different outputs generated from the same input prompt. That is, given an input prompt, different outputs are generated using different models. The models can be a copy of the generative machine learning model with different parameters obtained through checkpointing during pre-training or the models could be entirely unrelated. The input prompt and the different outputs are show n to human assessors and the human assessors are asked to rank the outputs in order of preference with respect to the input prompt. This can be repeated with many different input prompts to generate a dataset of preference data. A reward model can be trained on this preference data to provide a scalar preference value (a “reward’7value) for a particular input prompt and generated output pair. The reward model can be based on the generative machine learning model with an additional head for generating the scalar value for example.

[0125]

[0125] The generative machine learning model undergoing training can then be fine-tuned using reinforcement learning based upon the reward values provided by the trained reward model. That is, for a given training prompt, the generative machine learning model generates an output which can be evaluated using the reward model. The parameters of the generative machine learning model can be adjusted using a reinforcement learning update rule based upon the reward value provided by the reward model. In some implementations, a reinforcement learning update rule based upon the Proximal Policy Optimization (PPO) algorithm is used with the generative machine learning model acting as the “policy”.

[0126]

[0126] Direct Preference Optimization is a reformulation of RLHF that avoids the need to explicitly construct a reward model and enables training directly on the human preference data with a classification-based loss. Further details regarding DPO can be found in Rafailov. Rafael, et al. "Direct preference optimization: Your language model is secretly a reward model," Advances in Neural Information Processing Systems 36 (2024), which is hereby incorporated by reference in its entirety.

[0127]

[0127] Through such training, it is possible that a generative machine learning model can respond to user queries and instructions in a zero-shot manner, for example, by including appropriate instructions and examples in the prompt provided to the generative machine learning model without the need for further extensive fine-tuning.

[0128]

[0128] It will be appreciated that the above also applies for any autoregressive neural network that is used in the data generation process. An example of an autoregressive neural network for audio generation is AudioLM, further details of which can found in Z. Borsos, et al. "AudioIm: a language modeling approach to audio generation." IEEE / ACM transactions on audio, speech, and language processing 31 (2023): 2523-2533 which is hereby incorporated by reference in its entirety.

[0129]

[0129] 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, firmware, hardware, or a combination thereof, enabling it to carry out those operations or actions during operation. For instance, 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.

[0130]

[0130] 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 computing device 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.

[0131] As used herein, the term data processing apparatus includes any suitable computing device or hardware for use in performing the methods described in this specification. 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.

[0131]

[0132] 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 any other 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 GPGPU programming or TPU utilization, depending on the chosen hardware platform and desired performance characteristics.

[0133] 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 ty pically 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.

[0132]

[0134] 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 implemented using 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.

[0133]

[0135] 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 type 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.

[0134]

[0136] 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 portability or permanence.

[0135]

[0137] 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 exchanging documents 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.

[0136]

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

[0137]

[0139] 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 instance, 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.

[0138]

[0140] 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 to a 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.

[0139]

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

[0140]

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

[0141]

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

[0142]

[0144] What is claimed is:

Claims

CLAIMS1. A method performed by one or more data processing apparatus, the method comprising: obtaining a query: obtaining a subset of contextual data from contextual data relating to the query; processing the query and the subset of contextual data using a generative machine learning model to generate response data relating to the query or to generate an indication that the query cannot be satisfied using the subset of the contextual data; in response to determining that the query' cannot be satisfied using the subset of contextual data, processing the query' and the contextual data relating to the query' using the generative machine learning model to generate the response data.

2. The method of claim 1, wherein processing the query and the subset of contextual data using the generative machine learning model comprises: prompting the generative machine learning model to predict whether the query can be satisfied using the subset of contextual data.

3. The method of claim 1 or 2, wherein obtaining the subset of contextual data from contextual data relating to the query comprises selecting the subset of contextual data from the contextual data relating to the query.

4. The method of claim 3, wherein selecting the subset of contextual data is performed by a Retriever machine learning model.

5. The method of claim 3 or 4, wherein selecting the subset of contextual data is based upon evaluating a distance metric between the query and a plurality7of candidate chunks extracted from the contextual data relating to the query7.

6. The method of claim 5, wherein evaluating the distance metric between the query and the plurality of candidate chunks extracted from the contextual data relating to the query comprises: generating a query embedding from the query;generating a respective chunk embedding for each of the plurality of candidate chunks; and evaluating the distance metric between the query embedding and the respective chunk embeddings.

7. The method of claim 6, wherein the distance metric is based upon a cosine similarity.

8. The method of claim 6 or 7, wherein the query embedding is generated by a first encoder; wherein the chunk embedding is generated by a second encoder; and wherein the first and second encoders are the same encoder or wherein the first and second encoders are different encoders but jointly trained.

9. The method of any one of claims 5 to 8, wherein selecting the subset of contextual data based upon evaluating the distance metric between the query and the plurality of candidate chunks comprises: ranking the plurality of candidate chunks based upon the distance metric evaluations; and selecting the k highest ranked candidate chunks as the subset of contextual data.

10. The method of claim 9, further comprising: concatenating the selected k highest ranked chunks in rank order to form the subset of contextual data.

11. The method of claim 9 or 10, wherein the value of k is determined based upon the query.

12. The method of claim 11, wherein the value of k is determined based upon a query’ ty pe or a measure of query complexity.

13. The method of claim 11 or 12, wherein the value of k is determined based upon processing the query by the generative machine learning model to predict the value of k.

14. The method of claim 11 or 12, wherein the value of k is determined based upon processing the query by a machine learning model different to the generative machine learning model to predict the value of k.

15. The method of any one of claims 9 to 14, wherein k is an integer in the range 1 to 5 inclusive.

16. The method of any one of claims 5 to 15, wherein the plurality of candidate chunks are generated based upon a chunk size of approximately 300 words.

17. The method of any preceding claim, wherein the generative machine learning model is a large language model (LLM) based machine learning model.

18. The method of any preceding claim, wherein processing the query and the subset of contextual data using the generative machine learning model to generate response data comprises: generating an indication that the uery is only partially satisfiable using the subset of contextual data; and repeating: obtaining a further subset of contextual data; and processing the query, the further subset of contextual data and the previously obtained subsets of contextual data using the generative machine learning model until response data is generated or an indication that the query' cannot be satisfied is generated.

19. The method of claim 18, wherein the indication that the query is only partially satisfiable comprises a reformulated query; and wherein the further subset of contextual data is obtained based upon the reformulated query .

20. The method of claim 18 or 19. wherein the indication that the query is only partially satisfiable comprises a predicted size of the additional contextual data that is necessary' to satisfy the query; and wherein the further subset of contextual data is obtained based upon the predicted size.

21. The method of any one of claims 18 to 20, wherein the indication that the query is only partially satisfiable comprises partial response data; and wherein the generated response data comprises the partial response data.

22. The method of any preceding claim, wherein the indication that the query cannot be satisfied using the subset of contextual data comprises a confidence score.

23. The method of any preceding claim, further comprising: obtaining the contextual data relating to the query comprising determining, based upon the query, that contextual data is required to satisfy the query.

24. The method of claim 23, wherein determining, based upon the query, that contextual data is required to satisfy’ the query is performed by the generative machine learning model.

25. The method of claim 23, wherein determining, based upon the query, that contextual data is required to satisfy- the query comprises: processing the query using the generative machine learning model to generate data for one or more function calls to obtain the contextual data relating to the uery.

26. The method of claim 25. wherein the one or more function calls comprises at least one function call to an external data source.

27. The method of claim 26, wherein the external data source comprises a search engine or a database.

28. The method of any one of claims 25 to 27, further comprising: performing the one or more function calls based on the data for the one or more function calls.

29. The method of claim 28, wherein performing the one or more function calls based on the data for the one or more function calls comprises: determining to perform the function call based on the data for the one or more function calls.

30. The method of claim 29, wherein determining to perform the function call comprises parsing the data for the one or more function calls to determine the one or more functions to call.

31. The method of claim 29 or 30, wherein determining to perform the function call comprises parsing the data for the one or more function calls to determine one or more parameters for the one or more functions.

32. The method of claim 31, wherein the one or more parameters comprise one or more search terms.

33. The method of any one of claims 25 to 32, wherein the data for the one or more function calls comprises programming language code, markup language code, and / or query language code.

34. A method performed by one or more data processing apparatus, the method comprising: obtaining a query; obtaining a subset of contextual data from contextual data relating to the query; processing the query and the subset of contextual data using a generative machine learning model, wherein the processing comprises: generating an indication that the query is only partially satisfiable using the subset of contextual data; and repeating: obtaining a further subset of contextual data; and processing the query, the further subset of contextual data and the previously obtained subsets of contextual data using the generative machine learning model until response data is generated or an indication that the queiy cannot be satisfied is generated.

35. The method of claim 34, wherein the indication that the query' is only partially satisfiable comprises a reformulated query ; and wherein the further subset of contextual data is obtained based upon the reformulated query.

36. The method of claim 34 or 35, wherein the indication that the query is only partially satisfiable comprises a predicted size of the additional contextual data that is necessary to satisfy the query; and wherein the further subset of contextual data is obtained based upon the predicted size.

37. The method of any one of claims 34 to 36. wherein the indication that the query’ is only partially satisfiable comprises partial response data; and wherein the generated response data comprises the partial response data.

38. A method performed by one or more data processing apparatus, the method comprising: obtaining a query; processing the query using a first machine learning model to predict a value of k, wherein the predicted value of k is determined based upon the query', and wherein k represents a quantity of contextual data chunks to be selected; selecting k chunks of contextual data from a larger body of contextual data based on the predicted value of k and processing the query’ and the k selected chunks using a second machine learning model to generate response data.

39. The method of claim 38, wherein determining the predicted value of k based upon the query’ comprises determining the value of k based on a query’ ty pe or a measure of query’ complexity.

40. The method of claim 38 or 39 wherein the first machine learning model and the second machine learning model are the same generative machine learning model.

41. The method of claim 38 or 39, wherein the first machine learning model is a different machine learning model than the second machine learning model, and wherein the first machine learning model is specifically trained to analyze the query to predict the value of k.

42. The method of any one of claims 38 to 41, wherein selecting the k chunks of contextual data comprises:ranking a plurality of candidate chunks from the larger body of contextual data based on a distance metric evaluated between the query and each of the plurality of candidate chunks; and selecting the k highest ranked candidate chunks.

43. The method of claim 42, further comprising: concatenating the selected k highest ranked chunks in rank order to form a single subset of contextual data for processing by the second machine learning model.

44. The method of any one of claims 38 to 43, wherein k is an integer in the range 1 to 5 inclusive.

45. A system comprising: one or more data processing apparatus; and a memory storing instructions that when executed by the one or more data processing apparatus cause the one or more data processing apparatus to carry out a method according to any preceding claim.

46. A non-transitory computer-readable storage medium comprising instructions that when executed by one or more data processing apparatus cause the one or more data processing apparatus to carry out a method according to any one of claims 1 to 44.

47. A computer program product comprising instructions that when executed by one or more data processing apparatus cause the one or more data processing apparatus to carry out a method according to any one of claims 1 to 44.

Images