Method, system, and computer program for developing a machine learning neural network model

Abstract

A computer implemented method for developing a machine learning neural network model comprising a neural network that has a Transformer architecture. The method comprises instantiating the Transformer architecture of the neural network to include a plurality of attention blocks that include respective linear matrices comprising linear projection matrices. The linear projection matrices are multiplied by respective importance score functions which are based on a learnable importance parameter. The method comprises choosing the learnable importance parameter for training, initialising the sizes of the linear matrices of the attention blocks based on the chosen learnable importance parameter, and updating the sizes of the linear matrices of the attention blocks based on the learnable importance parameter of the respective importance score functions of the attention blocks. Applications of the present disclosure include, but are not limited to material informatics, medical / healthcare, and optimised decision making for use cases in drug development.

Description

[0001] METHOD, SYSTEM, AND COMPUTER PROGRAM FOR DEVELOPING A MACHINE LEARNING NEURAL NETWORK MODEL

[0002] The present application claims priority to EP 25 185985.6 filed on 27 June 2025, the entire contents of which are hereby incorporated by reference. The present application also claims priority to EP 25 152 582.0 filed on 17 January 2025, the entire contents of which are hereby incorporated by reference. The present application additionally claims priority to PCT / EP2025 / 056811 filed on 12 March 2025, the entire contents of which are hereby incorporated by reference.

[0003] The present disclosure relates to a method and system for developing a machine learning neural network model, for example having a Transformer architecture.

[0004] Training Transformer architectures, such as Large Language Models (LLMs), can be hard. LLMs typically comprise multiple layers, with each layer usually comprising an attention block and a feed-forward network block. When Transformer models are trained, the size of the attention and feed-forward blocks is typically fixed or preset according to some “rule of thumb”, but good machine learning practice would preferably require that the configuration of the Transformer architecture, such as the size of each individual block, is chosen such that it performs optimally on a given task. In other words, the size / dimension of intermediate embeddings of each of an LLM’s building blocks should preferably be based on the task or tasks to be performed. Although it is possible to find the appropriate sizes / dimensions through many trials of manual or automatic tuning, they are not known in advance, and performing such trials is typically computationally very expensive.

[0005] The number of combinations typically grows exponentially with the number of blocks; therefore, it is generally unfeasible to try all reasonable configurations even for small machine learning models. In the case of LLMs, time and money constraints may make it unfeasible to train even a few configurations. Consequently, in the current techniques, the sizes of LLM blocks are typically fixed to such large values (justified by the immense amount of training data) that the number of parameters of an LLM generally amounts to hundreds of billions. There is no guarantee that their chosen sizes are optimal, nor is there any good intuition as to whether most of the computation is useless, because there is no way or it is very difficult to assert the importance of the LLMs’ parameters. Accordingly, there is considerable room for improvement by either reducing or increasing the sizes of the LLM’s building blocks.

[0006] Background art is US20250124105A1 which relates to efficient token pruning in transformer-based networks, CN118673895A which relates to a dynamic adjustment method and system for a self-attention mechanism. Other non-patent literature documents are:

[0007] [1] Errica, F., Christiansen, H., Zaverkin, V., Niepert, M., &Alesiani, F. (2025). Adaptive Width Neural Networks. arXiv preprint arXiv:2501.15889[2] Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter. Neural architecture search: A survey. Journal of Machine Learning Research, 20:1-21, 2019

[0008] [3] Rupert Mitchell, Martin Mundt, and Kristian Kersting. Self expanding neural networks. arXiv preprint, 2023.

[0009] [4] Ding, Yizhuo, et al. "Adaptive Pruning of Pretrained Transformer via Differential Inclusions." arXiv preprint arXiv:2501.03289 (2025)

[0010] The present disclosure relates to PCT / EP 2025 / 056811 which discloses a method, system, and computer program for developing a neural network model, for example comprising multilayer perceptrons (MLPs). Both non-patent literature document [1] as well as PCT / EP 2025 / 056811 (not yet published) are hereby incorporated in entirety herein by reference.

[0011] Previous techniques, may impose an upper limit on sizes, which may be a typical approach of Neural Architecture Search (NAS) methods, such as those described in reference [2], which rely on supernetworks (a supernetwork is a giant network where some components get shut down during training, but its size remains constant).

[0012] The body of work investigating adaptive architectures as described in reference [3] has not been applied to Transformer-like architectures according to the knowledge of the present inventors.

[0013] Other techniques, such as pruning, may only reduce the number of neurons or weights in the Transformer architecture according to heuristics, such as a desired sparsity level (e.g. reference [4]).

[0014] It is an object of the present disclosure to seek to address or at least alleviate the above-mentioned problems.

[0015] In a first aspect, there is provided a computer implemented method for developing a machine learning neural network model comprising a neural network that has a Transformer architecture, the method comprising: instantiating the Transformer architecture of the neural network to include a plurality of attention blocks that include respective linear matrices comprising linear projection matrices, in which the linear projection matrices are multiplied by respective importance score functions which are based on a learnable importance parameter; choosing the learnable importance parameter for training; initialising the sizes of the linear matrices of the attention blocks based on the chosen learnable importance parameter; and updating the sizes of the linear matrices of the attention blocks based on the learnable importance parameter of the respective importance score functions of the attention blocks.

[0016] A second aspect of the present disclosure relates to computer program comprising computer implementable instructions which, when executed on a computer, cause the computer to carry out a method according to the first aspect.A third aspect of the present disclosure relates to a system comprising one or more processors and one or more storage devices. The system is configured to perform the method of the first aspect.

[0017] In contrast with previous techniques in which the sizes of attention block are fixed beforehand, or many different configurations are tested (which is computationally very expensive), examples of the disclosure may update the sizes of the linear matrices of the attention blocks based on the learnable importance parameter of the respective importance score functions of the attention blocks. In other words, for example, no rule of thumb to choose a “single size that fits all” may be needed. Furthermore, a need for performing multiple training runs to try different configurations may be significantly reduced or eliminated. In addition, the techniques of the disclosure may encourage some parameters to be more important than others. Additionally, a requirement to impose an upper limit on the sizes may also be reduced or eliminated.

[0018] Therefore, for example, a single training step or only a few training steps may be used to determine the appropriate size of an attention block, thus helping to reduce the amount of time and computational resources needed to tune this hyperparameter (relating to the size of the attention blocks). This may lead to significant cost savings, because the number of configurations to try before deployment may be greatly reduced. Since training Transformer Architectures is typically very computationally intensive, there is also an associated high energy usage for the training. Therefore, reducing the amount of time and computational resources needed to determine an appropriate size of the attention blocks may also provide an environmental benefit by reducing energy expenditure fortraining the neural network.

[0019] By using the techniques of the present disclosure, neural networks having a Transformer architecture may be obtained that are significantly smaller than those that may be obtained with previous techniques, such as those involving fixing the attention block size. This smaller size may, for example, allow the trained neural network to be deployed on devices having more limited processing resources such as a laptop, mobile phone, or other portable computing device. Furthermore, a trade-off between computational resources and performance may be more easily obtained without needing to try many different configurations, which would be computationally intense and time consuming.

[0020] Other aspects and features are defined in the appended claims.

[0021] Examples and techniques of the present disclosure will now be described by way of example with reference to the accompanying drawings, throughout which like references refer to like parts, and in which:

[0022] Fig. 1 shows a diagram of a conventional attention blockof a Transformer architecture; Figs. 2a and 2b are flow charts of examples of computer implemented methods for developing a neural network model;Fig. 3 shows a diagram of an example of an attention block of a Transformer architecture according to examples of the disclosure;

[0023] Fig. 4 shows charts of an example distribution that may be used as a basis for modifying sizes of attention blocks of a Transformer architecture; and

[0024] Fig. 5 is a block diagram of an example of a system for developing a neural network model.

[0025] A method and system for developing a neural network model is disclosed. In the following description, a number of specific details are presented in order to provide a thorough understanding of the examples of the present disclosure. It will be apparent to the person skilled in the art that these specific details need not be employed to practise the present disclosure. Conversely, specific details known to the person skilled in the art are omitted for the purposes of presenting the examples. Additionally, in the drawings, some features described herein are omitted from the drawings for the ease of understanding the drawings, but will be apparent to the skilled person in light of the present disclosure.

[0026] Fig. 1 shows a diagram of a conventional attention block of a Transformer architecture (an attention block may be considered to be an attention mechanism of the Transformer architecture). A Transformer architecture may comprise a plurality of attention blocks, and the attention blocks may be arranged as layers of the Transformer architecture. The Transformer architecture may also comprise feedforward blocks, which may also be arranged as layers of the Transformer architecture as is known in the art.

[0027] An input 102 to the attention block of Fig. 1 may be a set of n inputs X, where each input has dimension d. The inputs X may be transformed according to three different linear weight matrices 112, 114, and 116, which comprise a first linear weight matrix WQ 112 (query weights matrix WQ) having a dimension d x di , a second linear weight matrix WK 114 (key weights matrix WK) having a dimension d x di , and a third linear weight matrix Wv 116 (value weights matrix Wv) having a dimension d x d2. The weight matrices may be considered to be linear projections.

[0028] Accordingly, the first linear weight matrix WQ 112 may transform the inputs X to generate a first linear projection matrix Q 112 (query matrix) having a dimension n x di, the second linear weight matrix WK 114 may transform the inputs X to generate a second linear projection matrix K 124 (key matrix) having a dimension n x di, and the third linear weight matrix Wv 116 may transform the inputs X to generate a third linear projection matrix V 126 (value matrix) having a dimension n x d2. In other words, the first linear projection matrix Q 112 and the second linear projection matrix K 124 may have the same size as each other, n x di, whereas the third linear projection matrix V 126 may have a different size, n x d2.The linear projection matrices 122, 124, 126 may be generated according to the following equations:

[0029] Q

[0030]

[0031] . eqn. 1

[0032] K = (X WK) ... eqn. 2

[0033] V = (X Wv) ... eqn. 3

[0034] The first linear projection matrix Q 122 and the second linear projection matrix K 124 may be multiplied together to generate a matrix QKT142 having a size n x n, and passed through a softmax function 152 to generate attention scores 154 having a size n x n. The attention scores 154 may be multiplied by the values of the third linear projection matrix V 126 (value matrix) to generate an output matrix 156 having a size n x d2. In other words, the input X 102 having size n x d may be transformed to the output matrix 156 having the size n x d2.

[0035] The output matrix 156 may be generated according to:

[0036] . eqn. 4

[0037]

[0038] Figs. 2a and 2b are flow charts of examples of computer implemented methods for developing a neural network model according to examples of the disclosure. For example, the neural network model may be a machine learning neural network model. Fig. 3 shows a diagram of an example of an attention block of a Transformer architecture according to examples of the disclosure. The computer implemented methods 200a and 200b may be used to adapt the sizes of the attention blocks during training.

[0039] The Transformer architecture may include a plurality of attention blocks that include respective linear matrices. The linear matrices may comprise weight matrices and linear projection matrices that are generated by transforming inputs of the plurality of attention blocks based on the weight matrices. The linear projection matrices may be multiplied by respective importance score functions, and each importance score function may be based on a learnable importance parameter for that linear projection matrix

[0040] The method 200a may comprise, at a step s210, instantiating the Transformer architecture of the neural network to include a plurality of attention blocks that include respective linear matrices. The linear projection matrices may be multiplied by respective importance score functions which are based on a learnable importance score function. In some examples, each importance score function may be based on a learnable importance parameter for that linear projection matrix. Instantiating the Transformer architecture may be considered to be setting up the Transformer architecture.

[0041] Each attention block of the Transformer architecture may have the same configuration of linear matrices when initialised, or they could be different. The Transformer architecture may comprise a plurality of layers that may include one or more of the attention blocks. Thelinear matrices may comprise weight matrices and linear projection matrices. The linear projection matrices may be generated by transforming inputs of the plurality of attention blocks based on the weight matrices. The linear projection matrices may be multiplied by respective importance score functions. Each importance score function may be based on a learnable importance parameter for that linear projection matrix.

[0042] Referring to Fig. 3, for example, the Transformer architecture of the attention block of Fig. 3 may be similar to that described above with respect to Fig. 1 , but with some differences as set out below. Although the example of Fig. 3 shows one attention block, the techniques and examples described with respect to Figs. 2a and 2b and Fig. 3 could be applied to one or more of the plurality of attention blocks in the Transformer architecture.

[0043] In the attention block of Fig. 3, an input 302 may be a set of n inputs X, where each input has dimension d. The inputs X may be transformed according to three different weight matrices 312, 314, and 316, which comprise a first linear weight matrix WQ 312 (query weights matrix) having a dimension d x di, a second linear weight matrix WK 314 (key weights matrix) having a dimension d x di, and a third linear weight matrix Wv 316 (value weights matrix) having a dimension d x d2. The weight matrices may be considered to be linear projections.

[0044] Accordingly, the first linear weight matrix WQ 312 may transform the inputs X to generate a first linear projection matrix Q 312 (query matrix) having a dimension n x di, the second linear weight matrix WK 314 may transform the inputs X to generate a second linear projection matrix K 324 (key matrix) having a dimension n x di , and the third linear weight matrix Wv 316 may transform the inputs X to generate a third linear projection matrix V 326 (value matrix) having a dimension n x d2. In other words, the first linear projection matrix Q 312 and the second linear projection matrix K 324 may have the same size as each other, n x di, whereas the third linear projection matrix V 326 may have a different size, n x d2.

[0045] In other words, the inputs X 302 may be the considered to be analogous to the inputs X 102 of Fig. 1, the linear weight matrices 312, 314, 316 (linear projections) may be considered to be analogous to the respective linear weight matrices 112, 114, and 116 of Fig. 1, and the linear projection matrices Q 322, K 324, and V 326 may be considered to be analogous to the linear projection matrices Q 122, K 124, and V 126 of Fig. 1. More generally, in examples, the linear projection matrices of each attention block of the plurality of the attention blocks may comprise a query matrix, Q, a key matrix, K, and a value matrix, V.

[0046] However, in contrast to the arrangement of Fig. 1, the architecture of Fig. 1 may be modified by multiplying the linear projection matrices Q 322, K 324, and V 326 by respective importance score functions f( ), with f-. R ^ d. In particular, in the example of Fig. 3, the first linear projection matrix Q 322 may be multiplied by a first importance score function where are the learnable importance parameters for the first importance score function- The second linear projection matrix K 324 may also be multiplied by the first importance score function

[0047]

[0048] The third linear projection matrix V 326 may be multiplied by a second importance score function

[0049]

[0050] where A2are the learnable importance parameters for the second importance score function

[0051]

[0052] In examples, each row of the matrices, Q, K, and V, may be multiplied by their respect learned functions

[0053]

[0054] and f2R d2, which take as input a learnable parameter

[0055]

[0056] and2.

[0057] More generally, in examples, the importance score functions for each attention block may comprise a first importance score function f^A^ and a second importance score function f2(22), where A are the learnable importance parameters for the first importance score function f^A^, and A2are the learnable importance parameters for the second importance score function2(^2)- The O symbol (for example shown in Fig. 3) is herein taken to mean element-wise multiplication (with implicit broadcasting along the first dimension). In particular, element-wise multiplication may be performed between the matrices Q 322, K 324, V 326 and corresponding vectors produced by the functions / i, / i, and f2respectively. These functions may be parametrized by learnable parameters, but the techniques of the present disclosure are generally focused on one. The choice of the functions

[0058]

[0059] and f2may help contribute to addressing the problems and objects considered by the present disclosure as mentioned above. The importance score functions, such

[0060]

[0061] and f2, may output an importance score for each element of the second dimension (either d or d2), and such scores may be shared across the n inputs. Furthermore, these functions may help define dynamically changing sizes d and / or d2, as described herein. In other words, d and d2may be considered as dynamic sizes.

[0062] Although in the example of Fig. 3, the first importance score function fi and the second importance score function f2are different from each other, it will be appreciated that they could be the same. This may help reduce processing resources needed as well as helping simplify initialisation. However, an importance score function for the value matrix V may not be as appropriate to use for the query Q and key K matrices as fewer parameters may be available for helping update sizes of the attention block, and so the attention block sizes may be less well adapted, thus lowering performance. However, these considerations could be traded off when setting up the transformer architecture depending on the desired usage.

[0063] In the example of Fig. 3, the query matrix Q (first linear projection matrix Q 322) and the key matrix K (second linear projection matrix K 324) may each be respectively multiplied by corresponding vectors produced by the first importance score function f^A^. The value matrix V (third linear projection matrix V 326) may be multiplied by corresponding vectors produced by the second importance score function f^A^.Accordingly, linear projection matrices Q’, K’, and V’ may be generated according to the following equations respectively:

[0064] Q' = (X UZ«) O A i) — eqn. 5

[0065] K' = (X WK) Q / JUJ) ... eqn. 6

[0066]

[0067] In a similar manner to Fig. 1, the first linear projection matrix Q 322 and the second linear projection matrix K 324 may be multiplied together to generate a matrix Q’K’T342 having a size n x n, and passed through a softmax function 352 to generate attention scores 354 having a size n x n.

[0068] More generally, for example, a subset of the linear projection matrices in each of their respective attention blocks having the same output dimensions as each other may be combined together during training. For example, referring to Fig. 1, the first linear projection matrix Q 112 and the second linear projection matrix K 124 have the same size as each other, n x di, and may be considered together as a subset of the linear projection matrices Q 112, K 114, and V 116. Similarly, the first linear projection matrix Q 312 and the second linear projection matrix K 324 have the same size as each other, n x di , and may be considered as a subset of the linear projection matrices Q 312, K 314, and V 316. The importance score functions for the linear projections of the subset may be the same as each other. For example, the first linear projection matrix Q 312 and the second linear projection matrix K 324 may both be multiplied by the same importance score function fi. For example, instantiating the Transformer architecture (e.g. at the step s210) may comprise setting up (instantiating) the Transformer architecture as described herein, such as to include attention blocks as described with reference to Fig. 3.

[0069] In examples, all linear projection matrices of the plurality of attention blocks may be multiplied by respective importance score functions. This may allow a more efficient training process to generate a neural network model with better performance. However, it will be appreciated that only some of the linear projection matrices may be multiplied by respective importance score functions, for example depending on the desired application. Additionally, attention blocks of some of the layers of the Transformer architecture may not include importance score functions.

[0070] The attention scores 354 may be multiplied by the values of the third linear projection matrix V 326 (value matrix) to generate an output matrix 356 having a size n x d2. In other words, the inputs X 302 having size n x d may be transformed to the output matrix 356 having the size n x d2. Similarly to the technique of Fig. 1, the output matrix 356 may be generated according to: > eqn. 8

[0071]

[0072] Returning to Fig. 2a, the method 200a may comprise, at a step s222, choosing the learnable importance parameter for training (e.g. values of the learnable importance parameter(s)). For example a learnable importance parameter or learnable importance parameters for the importance score functions may be chosen for training, for example for training the neural network to learn appropriate sizes of attention blocks as described herein.

[0073] The method 200a may comprise, at a step s224, initialising the sizes of the linear matrices of the attention blocks based on the chosen learnable importance parameter(s). In examples, the same learnable importance parameter(s) may be used for all the attention blocks, which may help simplify initialisation, although it will be appreciated that they could be different from each other.

[0074] The method 200a may comprise, at a step s232, updating the sizes of the linear matrices of the attention blocks based on the learnable importance parameter of the respective importance score functions of the attention blocks. For example, updating of the sizes of the linear matrices may be based on a threshold value determined with respect to the importance score function(s), as will be described in more detail later below.

[0075] Fig. 2b shows a method 200b for developing a neural network model. The method 200b may include one or more of the steps of the method 200a. For example, the method 200b may include the steps s210, s222, s225, and s232.

[0076] The method 200b, may comprise, at a step s220, initialising the transformer architecture. Initialising the transformer architecture may comprise one or more of the step s222 and the step s224.

[0077] As mentioned above, the linear matrices may comprise weight matrices and linear projection matrices, and the linear projection matrices may be generated by transforming inputs of the plurality of attention blocks based on the weight matrices. The step s224 may comprise initialising weights of the Transformer architecture.

[0078] The weights of the Transformer architecture (such as those of the first linear weight matrix WQ 312, the second linear weight matrix WK 314, and the third linear weight matrix Wv 316) may be initialised as commonly performed in machine learning as would be appreciated by a person skilled in the art. However, in order to help reduce the likelihood that gradients do not vanish when many Transformer layers are used, initialisation of the weights may be performed according to the strategy highlighted in non-patent literature [1]. In particular, the usual Kaiming weight initialisation may be rescaled, where the scaling factor is square root of 2 divided by the sum of the squares of all importances used.

[0079] The method 200b may comprise, at a step s230, using a training dataset for training the neural network by implementing one or more training steps. Each training step maycomprise updating the sizes of the linear matrices. For example, each training step may comprise the step s232.

[0080] Each training step may comprise, generating, at a step s234, a current value of the neural network’s loss function to be minimised, and generating, at a step s236, gradients of the neural network’s loss function and updating parameters of the neural network, including the learnable importance parameters of the importance score functions, based on the generated gradients. The current value of the neural network’s loss function may be considered to be the value of the neural network’s loss function as evaluated during the respective training step.

[0081] Training steps may be implemented as an iterative training process. For example, the iterative training process may comprise a plurality of the training steps performed sequentially and iteratively. In some examples, only one training step may be carried out. Using a single training step or only a few training steps may help to reduce the amount of time and computational resources needed to determine an appropriate size of an attention block.

[0082] After training, the methods 200a and 200b may comprise outputting a neural network (e.g. a machine learning neural network). For example, the neural network may be used to predict properties of interest for a given sample.

[0083] At the step s232, updating the sizes of the linear matrices (such as sizes of the weight matrices and the linear projection matrices) may comprise determining a threshold value for each importance score function, and updating the sizes of the linear matrices based on the threshold values. By learning values of the learnable importance parameters, the threshold value may change, and the sizes of the linear matrices may be updated accordingly. This may help reduce the number of different configurations to try before deployment, which may help lead to significant cost savings, as well as reducing the computational time and resources needed before deployment.

[0084] In examples, the importance score functions may be monotonically decreasing functions, for example analogous to the techniques described in non-patent reference [1]. For example, the importance score functions may be a discretized exponential distribution parametrized by 2 for each function fa and f2(2y, this may help produce a monotonically decreasing importance for each element of the d and d2-dimensional vectors. The first dimensions will likely be more important, whereas the last ones should intuitively have less impact on subsequent computations.

[0085] The output of a discretized exponential function, however, is infinite in principle (one value for every natural number). Therefore, a threshold k may be chosen after which to truncate the output of the function, otherwise importance scores may not be generated in practice. In the case of a discretized exponential distribution, a threshold may be chosen which corresponds to the value below which a certain percentage (e.g., 90%) of the probability massis contained. This value may be generated as the quantile function of the discretized exponential evaluated at k. However, it will be appreciated that the percentage could be chosen to have any other appropriate value.

[0086] The functions

[0087]

[0088] may be chosen to be differentiable with respect to their parameters 2: the discretized exponential evaluated at point x may be generated by subtracting the differentiable cumulative distribution function (cdf) of the continuous exponential distribution evaluated at x from the same cdf evaluated at x + 1. Both cdfs may be differentiable, hence the discretized exponential may be differentiable with respect to 2.

[0089] Therefore, the values that satisfy the threshold k, which change with the parameters 2i and 22, may be used to determine the dynamic sizes d and d2of the linear matrices (such as the linear projection matrices or weight matrices), which can therefore grow or shrink. Fig.

[0090] 4 shows charts (taken from non-patent reference [1]) of an example distribution that may be used as a basis for modifying sizes of attention blocks of a Transformer architecture, and which visually show how an exponential function (left) may first be discretized (right) and may then be truncated.

[0091] The step s232 may comprise generating a threshold value based for each C2f), *e{1,2}. For example, if fa is a monotonically decreasing distribution, the value of the quantile function of fa evaluated at a value k e [0,1] may be used as the respective thresholds indicating the new sizes d , d2. Referring to Fig. 4 for example, the value k may be considered to represent the value before which k percent of the whole probability mass is produced. The threshold(s) may be chosen by a user, for example prior to training. A higher threshold is generally found to give better performance of the neural network model but is also generally more computationally expensive.

[0092] The step s232 may comprise checking the thresholds against the respective sizes, such as ch and ds, of the linear projection matrices. If both are equal, then the sizes of the respective linear matrices may not be updated.

[0093] In some examples, the shape of the linear matrices (such as the weight matrices and / or linear projection matrices) may be updated (e.g. at the step s232) to match the dimensions d1,d2, by adding new dimensions. The weights may be randomly initialised. For example, new dimensions may be added if a respective dynamic size, such as di and / or ds, is determined to be less than the respective threshold. New dimensions may be added until the respective size matches the respective threshold.

[0094] In other examples, updating (e.g. at the step s232) the sizes of the linear matrices (such as the weight matrices and / or linear projection matrices) may comprise removing excess dimensions. For example, if a respective dynamic size, such as di and / or ds, is determined to be greater than the respective threshold, then excess dimensions may be removed, forexample those corresponding to the smallest values of the respective importance score function / i C^-D- lnother words, the sizes of the linear matrices (e.g. the sizes of the weight matrices and / or linear projection matrices) may be caused to grow or shrink.

[0095] Analogously, a similar process may be applied to adapt the size of linear projections, such as linear projections (weight matrices) WQ312, WK314, and Wv316, used to generate the linear projection matrices, such as linear projection matrices Q, K, and V.

[0096] At the step s234, the loss function (e.g. used to train the neural network model during the one or more training steps) may be generated according to:

[0097]

[0098] > where Tpred(y, y) is predictive loss (e.g., cross-entropy) with y being the network’s predictions and y being the ground truth,proj may be an optional projection regularization term on projection weights, and

[0099]

[0100] may be an optional lambda regularization term on the A parameters. Eitherproj

[0101]

[0102] or both, may be omitted. In other words, for example, the current value of the neural network’s loss function may be generated at the step s234 according to Eqn. 9.

[0103] In examples, the gradients may be generated at the step s236 using backpropagation, for example as handled in frameworks such as PyTorch and TensorFlow. For example, given a loss £, backpropagation may be used to generate gradients with respect to parameters 0 of the attention block according to:

[0104]

[0105] This may comprise generating the gradients for the linear projections (weight matrices), such as WQ312, WK314, and Wv316, and the learnable importance parameters, such as 2i,2.

[0106] The step s236 may comprise updating each parameter of the attention block. For example, values of the linear projections (weight matrices) may be updated. The updating may comprise updating values for the learnable importance parameters, such as AI,22.

[0107] A known optimiser, such as “Adam” or “Adagrad” may be used in the step s236. In some examples, parameters of an attention block (such as the learnable importance parameters, and / or other parameters) may be updated using Stochastic Gradient Descent (SGD) techniques. For example, for SGD, a parameter 0j may be updated using:

[0108]

[0109] where q is the learning rate. An update may be applied per batch (or per example if pure SGD is used).As mentioned above, optional regularisation terms may be included when generating the loss function given by Eqn. 9. The methods 200a and 200b may comprise, during training (e.g. as part of the step s230), compressing the neural network using a regularises The regulariser may comprise the projection regularization term,prOj, the lambda regularization term, both.

[0110] Projection regularisation may be, for example, L2 norm of projection weights. For example, the projection regularization term,prOj, may be determined according to:

[0111]

[0112] Lambda regularisation may be, for example, L2 norm of difference with respect to reference values. Here, Lambda regularisation may be taken to refer to regularisation of the learnable importance parameters. For example, the lambda regularisation term,

[0113]

[0114] may be determined according to:

[0115]

[0116] < where andr2efare reference values. Accordingly, for example, the lambdas (e.g.

[0117]

[0118] and 22) may be encouraged to be closer to the reference values 2^ and A2ef, which are typically (although not necessarily) associated with smaller dimensions of the linear projections’ (weight matrices) dimensions (sizes) di and ds, than what is learned. In other words, more generally, regularisation may help bias the learnable importance parameters to be closer to respective learnable importance parameter reference values. Therefore, for example, regularisation may help reduce the size of the neural network more efficiently, such as to suit particular hardware requirements.

[0119] In examples, the methods 200a and 200b may comprise performing post training truncation of the neural network model by removing, after training, less important dimensions. For example, the post training truncation may be performed after the step s232 (e.g. with reference to the method 200a) or after the step s236 (with reference to the method 200b). The post training truncation may be performed after all training steps are completed. For example, post training truncation may be performed if hardware requirements are not already satisfied. The post training truncation may comprise removing N less important neurons from a linear projection whose weight matrix has shape d x d', where d' is the output dimension, by removing the rightmost N columns from the weight matrix. Intuitively, the last neurons of a layer are likely to be the least important due to the training carried out according to examples and techniques described herein such as that described with respect to the methods 200a and 200b of Figs. 2a and 2b. In other words, for example, less important dimensions may be determined according to the importance functions, i.e. which dimensions are considered to be less important may depend on the importance score functions. For example, dimensionscorresponding to the smallest values of the respective importance score function / imaY be removed.

[0120] The number N of less important neurons may be determined by a user, for example, based on hardware requirements, for example bearing in mind a trade-off between how many neurons may be removed and the resultant performance. Removing more dimensions may generally result in a smaller trained model that may be more easily implemented on hardware with more limited computational and memory resources (such as a laptop or mobile device), but with some detriment to performance. Removing fewer dimensions may generally improve performance, but result in increased hardware requirements.

[0121] In examples, truncation of the size of the projection matrices (weight matrices) WQ, WK, WVin an attention block may be performed in a similar manner to removing neurons mentioned above, for example by removing the rightmost A / ' columns from the weight matrix, where N’ is the number of less important dimensions to be removed, which may be defined by a user, for example taking into account similar considerations to those mentioned above with respect to removing the N less important neurons. As mentioned above, the first linear projection (first weight matrix) WQ 312 may have a dimension d x di, the second linear projection (second weight matrix) WK 314 may have a dimension d x di, and the third linear projection (third weight matrix) Wv 316 may have a dimension d x d2. Therefore, when removing dimensions from the first linear projection WQ 312 and the second linear projection WK 314, the new dimension of d should be the same for both the WQand WKmatrices.

[0122] In examples, when the value d2of a weight matrix Wvis truncated, the weight matrices of the next attention layer may be modified accordingly by making sure their input dimension matches d2. For example, this may be done by removing an appropriate number of rows from the weight matrices WQ, WK, Wvof the next attention layer.

[0123] Fig. 5 shows a schematic diagram of a corresponding system 500, e.g., a computer system, for developing a neural network model, with the system 500 being configured to perform the methods described herein, such as those described with respect to Figs. 2a and 2b, and Fig. 3. The system 500 may comprise an optional interface 502 and one or more processors 504 configured to perform the method of Figs. 2a and 2b. Features introduced in connection with the method of Figs. 2a and 2b and the attention block of Fig. 3 may likewise be performed by the system 500 of Fig. 5. For example, the system 500 may comprise machine-readable instructions for providing the functionality of the methods and techniques described herein, with the one or more processors 504 executing the machine-readable instructions to perform the method of Figs. 2a and 2b. The one or more processors 504 may be coupled with the interface 502 and with memory and / or storage, such as one or more storage devices 506. For example, the one or more processors 504 may be configured toprovide the functionality of the system 500, e.g., in conjunction with the interface 502 (for exchanging information) and / or with the memory / storage 506 (for storing information). The system 500 may be implemented as a distributed computing system.

[0124] The interface 502 may receive a training dataset 508 obtained from a suitable source. The interface 502 may include or correspond to a network interface circuitry and / or a device interface circuitry configured to be communicatively coupled to one or more other devices, such as the one or more processors 504. For example, the interface 502 may include a transmitter, a receiver, or a combination thereof (e.g., a transceiver), and may enable wired communication, wireless communication, or a combination thereof. For example, the one or more processors 504 may include or correspond to a digital signal processor circuitry (DSP), a graphical processing unit (GPU), and / or a central processing unit (CPU). The one or more processors 504 may be coupled to the memory / storage circuitry(s) 506.

[0125] The memory / storage circuitry 506 may include instructions (e.g., executable instructions), such as computer-readable instructions or processor circuitry-readable instructions. The instructions may include one or more instructions that are executable by a computer, such as by the one or more processors 504. For example, the memory / storage circuitry 506 may include or correspond to volatile or non-volatile storage circuitry, such as Random Access Memory (RAM), magnetic disks, optical disks, or flash memory devices. The memory / storage circuitry 506 may include both removable and non-removable memory devices. The system 500 may output a neural network model 510 developed according to the techniques and methods described herein, such as those described with respect to the methods 200a and 200b of Figs. 2a and 2b.

[0126] The techniques and examples of the present disclosure may be applied to Large Language Models (LLMs) for Predicting Protein-Ligand Binding Properties to Optimize Personalized Cancer Treatments. The data source (e.g. the training dataset) may be a dataset of articles about protein structures, ligand molecules, and known binding affinities between them on commodity hardware.

[0127] The examples and techniques described herein may be applied to the Attention mechanisms inside all Transformer layers of LLMs to learn a performant size of each block. The techniques described in [1] may be applied to the feedforward mechanisms. Depending on a user’s hardware requirements, the sizes of each Transformer layer may be carefully reduced, for example as described herein, such as by regularizing the learnable importance parameters or by removing the least important elements of each dimension, thus trading off performance for computational costs. The truncated LLM may then be deployed on a user’s system to be used, for example a laptop computer or other portable computing device.

[0128] The examples and techniques of the disclosure may produce or output an LLM that may be deployed on clinical systems and queried to get the binding affinity of specific protein-ligand pairs, in real time, based on the literature in the dataset. Oncologists and clinical researchers may input patient-specific protein variants and drug molecules. The LLM may predict binding efficacy and possible side effects, and may rank ligand candidates accordingly. This may enable more precise selection of therapeutic compounds tailored to individual patient profiles for wet lab experiments.

[0129] In other words, more generally in examples, the neural network may be a large language model, LLM, and the training dataset may comprise training data relating to protein structures, ligand molecules, and known binding affinities between them. The method 200a and / or 200b may comprise storing the trained LLM to enable outputting, in response to a query relating to binding affinity of protein-ligand pairs, a prediction of binding efficacy and possible side effects together with a ranking of ligand candidates.

[0130] A physical change that may be enabled by the techniques and examples of the disclosure in this example may be the selection of a subset of drug candidates with predicted high efficacy and safety for a given patient. These candidates may be prioritized for wet lab analysis and potential synthesis or clinical trial, directly impacting laboratory workflows, drug formulation pipelines, and patient treatment plans.

[0131] The techniques and examples of the present disclosure may be applied Large Language Models (LLMs) for Automatic Material Discovery on Commodity Hardware. The data source (e.g. the training dataset) may be a dataset of material design articles to train an LLM.

[0132] The examples and techniques described herein may be applied to the Attention mechanisms inside all Transformer layers of LLMs to learn a performant size of each block. The techniques described in [1] may be applied to the feedforward mechanisms. Depending on a user’s hardware requirements, the sizes of each Transformer layer may be carefully reduced, for example as described herein, such as by regularizing the learnable importance parameters or by removing the least important elements of each dimension, thus trading off performance for computational costs. The truncated LLM may then be deployed on a user’s system, for example a laptop computer or other portable computing device, to be used. The examples and techniques of the disclosure may produce or output an LLM that may be deployed on a user system and may be used to answer a user’s queries for a complex task. By enabling material scientists to query and interact with this LLM in real time on commodity devices, the LLM may recommend new materials to test, thereby directly influencing laboratory experiments and material design cycles.

[0133] In other words, more generally, the neural network may be a large language model, LLM, and the training dataset may comprise training data relating to material design properties. The method 200a and / or 200b may comprise storing the trained LLM to enableoutputting of recommendations of materials to test in response to queries relating to desired material properties to address a desired material property design task.

[0134] Accordingly, for example, the LLM may be used to re-rank materials according to the properties that the LLM predicts about these materials. Therefore, a physical change may be represented by selection of a subset of potential materials that need to undergo further testing in a real-world scenario.

[0135] More generally, the techniques and examples of the present disclosure may be implemented as a computer-implemented method, computer system (comprising one or more processors and one or more storage devices) configured to perform the computer-implemented method and / or as a computer program for performing the computer-implemented method. For example, the computer-implemented method may include one or more steps and / or operations discussed herein.

[0136] Various aspects of the disclosure relate to machine learning. In particular, the model mentioned above may be a machine learning model.

[0137] Machine learning is a branch of artificial intelligence that involves the development of algorithms and models that allow computers to learn and make predictions or decisions without being explicitly programmed. It focuses on creating systems that can improve their performance over time by learning from data.

[0138] Training a machine-learning model refers to the process of teaching the model to make accurate predictions or decisions. During training, the model may be exposed to a large amount of data, which is used to adjust the model's internal parameters or weights. The model learns patterns, relationships, or rules from the training data, allowing it to generalize and make predictions on new, unseen data.

[0139] Training data may be set of examples or instances that is used to teach a machinelearning model. It is often labelled data, meaning that each example is associated with a known outcome or target value. The training data may comprise both input features and the corresponding output or target variable. The model may learn from this data by analysing patterns and relationships between the input features and the target variable. Training algorithms, such as supervised learning, semi-supervised learning, unsupervised learning or reinforcement learning may be used fortraining the machine-learning model.

[0140] Machine-learning models, such as the machine-learning model being trained in the present disclosure, are often implemented as Artificial Neural Networks (ANNs), and in particular Deep Neural Networks, Support Vector Machines, Decision Tree models, or Random Forest models.

[0141] Examples and techniques may involve or relate to computer programs, including program codes to execute one or more of the mentioned methods when the program is executed on a computer, processor, or other programmable hardware component. As a result,steps, operations, or processes from various methods described above may also be executed by computers, processors, or other programmable hardware components. Examples may additionally cover program storage devices, such as digital data storage media, which are machine-, processor-, or computer-readable and encode and / or contain machine-executable, processor-executable, or computer-executable programs and instructions. These devices may include or be digital storage devices, magnetic storage media like magnetic disks and tapes, hard disk drives, or optically readable digital data storage media, for instance. Other examples may encompass computers, processors, control units, field programmable logic arrays (FPLAs), field programmable gate arrays (FPGAs), graphics processing units (GPUs), application-specific integrated circuits (ASICs), integrated circuits (ICs), or system-on-a-chip (SoC) systems that are programmed to carry out the steps of the methods described herein. In simpler terms, examples and techniques may involve computer programs and storage media comprising computer programs, as well as hardware components like processors and control units, which can be programmed to execute the methods described herein.

[0142] Although the techniques and examples mentioned above with respect to Figs. 2a and 2b, and Fig. 3 have been described in the context of Query, Key, and Value matrices, the transformations described herein may be applied with any linear projection, for example both in the attention blocks and in the feedforward blocks (if possible, after a non-linearity). To do so, care should be taken that the shapes of the inputs and outputs across the architectural blocks are consistent with each other. This may be achieved by dynamically adjusting the shape of matrices linking the output of a block with the input of another, and by sharing the same learned function f between linear projections whose outputs are later combined, for example as mentioned above for the matrices Q and K (e.g. for those matrices that are considered to belong to a subset of the linear projections in each of their respective attention blocks having the same output dimension that combined together, as mentioned above). The techniques disclosed herein may be applied in general to Transformer architectures which make use of linear transformations of vectors.

[0143] The techniques and examples of the present disclosure may be applied to Transformer models which are used in many application domains. For example, transformers are used in LLMs, but they may appear in Vision Transformers for image classification, in bioinformatics such as AlphaFold, and may be used to process sequences for applications such as speech recognition and financial predictions.

[0144] Examples of the present disclosure may provide methods and systems for learning the size of all building blocks of a Transformer architecture, comprising one or more of the following steps / components:

[0145] 1) Collecting a dataset2) Instantiating the Transformer architecture such as an LLM with the modifications mentioned herein for example, with respect to Fig. 3:

[0146] a. Choose an initial parameter or parameters 2 for all importance functions in the architecture, which will be adapted during training. One can start with the same parameter for all adaptive blocks;

[0147] b. Initialize the architecture with the right sizes implied by the choice of 2; c. Initialize the weights of the resulting Transformer architecture as usually done in machine learning;

[0148] 3) Training the Transformer architecture, where at each training step:

[0149] a. The sizes of each attention scheme are updated according to the parameters of the importance functions;

[0150] b. Inference is performed;

[0151] c. The loss is computed; and

[0152] d. The parameters of the Transformer architecture and of the importance functions are updated according to the gradient of the loss.

[0153] Techniques and examples of the disclosure may help save a significant amount of time that is usually spent tuning the size of an LLM; even if tuning is not performed, which leads to suboptimal performance, the techniques and examples of the present disclosure may help learn a much smaller LLM with subsequent cost reductions. For example, in non-patent literature document [1], the described techniques of [1] to feedforward architectures within an LLM already reduced the dimensionality of the feedforward parameters by a factor of 200 without significantly degrading performance.

[0154] The techniques and examples of the disclosure may allow the size of the transformations in any LLM architecture to grow and shrink according to the training signal, namely the gradient of the loss with respect to the model’s parameters, as this signal changes the parameters 2 of the importance functions fin the architecture. The method may start with an initial size for each block, determined by the initial parameters of the functions f, and then parameters may be adapted based on the training signal, using the similar variational framework to that described in reference [1].

[0155] When certain examples and techniques are mentioned in relation to a device or system, they should also be considered as descriptions of the corresponding methods. For example, a block, component, or functional example of the device or system may correspond to a method step or feature of the related method. Therefore, examples and techniques described regarding a method should also be understood as depicting a corresponding element, property, or functional feature of the corresponding device or system. In simplerterms, if something is described in relation to a device or system, it can also be applied to the corresponding method, and vice versa.

[0156] Many modifications and other embodiments and examples set forth herein will come to mind the one skilled in the art to which the disclosure pertains having the benefit of the teachings presented in the description and the associated drawings. Therefore, the disclosure is not to be limited to the specific embodiments and examples disclosed herein and that modifications and other embodiments are intended to be included within the scope of the invention as defined by the appended claims. Although specific terms are employed herein, they are used in a generic and descriptive sense only and not for purposes of limitation. It will be appreciated that the features of one or more of any of the different examples described herein may be combined together as appropriate with changes as appropriate, which will be apparent to the skilled person from the present disclosure.

[0157] Insofar as various examples and techniques have been described with reference to the present disclosure, it will be appreciated that various modifications and alterations may be made without departing from the spirit and scope of the invention as defined in the appended claims and their equivalents.

Claims

CLAIMS1. A computer implemented method for developing a machine learning neural network model comprising a neural network that has a Transformer architecture, the method comprising:instantiating the Transformer architecture of the neural network to include a plurality of attention blocks that include respective linear matrices comprising linear projection matrices, in which the linear projection matrices are multiplied by respective importance score functions which are based on a learnable importance parameter;choosing the learnable importance parameter for training;initialising the sizes of the linear matrices of the attention blocks based on the chosen learnable importance parameter; andupdating the sizes of the linear matrices of the attention blocks based on the learnable importance parameter of the respective importance score functions of the attention blocks.

2. A method according to claim 1 , in which:the linear matrices comprise weight matrices;the linear projection matrices are generated by transforming inputs of the plurality of attention blocks based on the weight matrices;each importance score function is based on a learnable importance parameter for that linear projection matrix; andthe method comprises using a training dataset for training the neural network by implementing one or more training steps, in which each training step comprises:the updating of the sizes of the linear matrices;generating a current value of the neural network’s loss function to be minimised;andgenerating gradients of the neural network’s loss function and updating parameters of the neural network, including the learnable importance parameters of the importance score functions, based on the generated gradients.

3. A method according to claim 1 or claim 2, in which the importance score functions are monotonically decreasing distributions.

4. A method according to claim any preceding claim, in which, during training, a subset of the linear projection matrices in each of their respective attention blocks having the same output dimension are combined together, and the importance score functions for the linear projection matrices of the subset are the same as each other.

5. A method according to any preceding claim, in which all linear projection matrices of the plurality of attention blocks are multiplied by respective importance score functions.

6. A method according to any preceding claim, in which updating the sizes of the linear matrices comprises:determining a threshold value for each importance score function; andupdating the sizes of the linear matrices based on the threshold values.

7. A method according to claim 5, in which updating the sizes of the linear matrices comprises adding new dimensions having randomly initialised weights.

8. A method according to claim 5, in which updating the sizes of the linear matrices comprises removing excess dimensions.

9. A method according to any preceding claim, in which:the linear projection matrices of each attention block of the plurality of the attention blocks comprise a query matrix, Q, a key matrix, K, and a value matrix, V;the importance score functions for each attention block comprise a first importance score functionand a second importance score function f2( .2), whereare learnable importance parameters for the first importance score functionand 22are learnable importance parameters for the second importance score function f2(2,the query matrix, Q, and the key matrix, K, are each respectively multiplied by corresponding vectors produced by the first importance score function l i); andthe value matrix, V, is multiplied by corresponding vectors produced by the second importance score function2(^2)-10. A method according to any preceding claim, comprising, during training, compressing the neural network by using a regulariser.

11. A method according to any preceding claim, comprising:performing post training truncation of the neural network model by removing, after training, less important dimensions.

12. A method according to any of claims 2 to 11 , in which:the neural network is a large language model, LLM;the training dataset comprises training data relating to protein structures, ligand molecules, and known binding affinities between them; andthe method comprises storing the trained LLM to enable outputting, in response to a query relating to binding affinity of protein-ligand pairs, a prediction of binding efficacy and possible side effects together with a ranking of ligand candidates.

13. A method according to any of claims 2 to 11 , in which:the neural network is a large language model, LLM;the training dataset comprises training data relating to material design properties; and the method comprises storing the trained LLM to enable outputting of recommendations of materials to test in response to queries relating to desired material properties to address a desired material property design task.

14. A computer program comprising computer implementable instructions which, when executed on a computer, cause the computer to carry out a method according to any preceding claim.

15. A system comprising one or more processors and one or more storage devices, in which the system is configured to carry out a method according to any of claims 1 to 13.

Images