Training method of graph neural network for atomic array rearrangement, atomic array rearrangement method using graph neural network, electronic device and computer readable storage medium

Abstract

Provided are a training method of a graph neural network for atomic array rearrangement, an atomic array rearrangement method using a graph neural network, an electronic device, and a computer-readable storage medium. The training method of the graph neural network for atomic array rearrangement determines an optimal frame number by bisection search by solving a minimum cost maximum flow on a time-expanded graph, and extracts an optimal atom-target pairing number and an optimal frame execution number of each atom from an optimal flow solution to form a joint structured training label, so as to train a three-head graph neural network including a matching head, a frame assignment head, and a compatibility head. In this way, when the atomic array rearrangement method using the graph neural network is used to perform atomic array rearrangement, the matching head outputs a pairing result in a differentiable pairing decoding manner, and the frame assignment head directly outputs a frame execution number by taking argmax for each atom, so as to eliminate greedy graph coloring post-processing, thereby achieving atomic array rearrangement without path conflict with low inference delay time and low frame number.

Description

Technical Field

[0001] This application relates to the field of quantum computing technology, and more specifically, to a training method for a graph neural network for atomic array rearrangement and a method for atomic array rearrangement using the graph neural network. Background Technology

[0002] In atomic quantum computing, taking neutral atom quantum computing platforms such as Rb, Cs, and Yb based on optical tweezers as an example, neutral atoms are loaded into a two-dimensional lattice array with a Bernoulli random distribution (typical loading probability of about 0.5-0.8), resulting in approximately 20%-50% of the target lattice sites being vacant (defective). Furthermore, loading mechanisms in other atom-type platforms such as ion traps also result in vacancies. Therefore, before the quantum algorithm executes, the loaded atoms must be rearranged to the designated target sites to form a defect-free array.

[0003] Let the array size be The initial number of atoms is Number of target atoms A frame is defined as a batch of parallel atomic movements in which paths do not intersect. The goal of optimizing the atomic array rearrangement is to minimize the total number of frames required to complete the rearrangement.

[0004] Existing solutions typically decompose the problem into three independent sequential steps, where local optima at each step cannot guarantee global optima. Specifically, the atom-target pairing uses the Hungarian algorithm, which has a complexity of O(n³), which is insufficient for large-scale arrays ( 1000) Insufficient real-time performance. Path planning uses independent straight paths, resulting in a lack of global coordination. Frame scheduling uses a greedy graph coloring algorithm, but because the greedy graph coloring algorithm is an NP-hard approximation, the frame count is too high, which introduces additional latency and frame count errors, making end-to-end gradient optimization impossible.

[0005] Currently, there are also graph neural network-based solutions that directly learn the mapping from array states to movement schemes, using the output of the Hungarian algorithm as training labels and aiming for millisecond-level inference speed. However, the label quality when using the Hungarian algorithm results as training labels is limited by the Hungarian algorithm itself (because it does not contain optimal frame information, the label accuracy is limited to the Hungarian algorithm), and it still relies on greedy graph coloring post-processing for frame allocation, which cannot achieve end-to-end output frame scheduling. In other words, single-step deep learning acceleration cannot truly eliminate the post-processing bottleneck.

[0006] Therefore, it is desirable to provide an improved atomic array rearrangement scheme based on graph neural networks. Summary of the Invention

[0007] This application provides an atomic array rearrangement scheme based on graph neural networks. By training a multi-head graph neural network, it can decode and directly output the atom-target pairing results and the frame execution number of the atom, thereby achieving path conflict-free atomic array rearrangement with low inference latency and low frame count.

[0008] According to one aspect of this application, a training method for a graph neural network for atomic array rearrangement is provided, comprising: obtaining an initial two-dimensional grid array of atoms for training and training labels based on the rearrangement result of the initial two-dimensional grid array; inputting the initial two-dimensional grid array of atoms for training into a feature extraction backbone network of the graph neural network for feature extraction to obtain a set of node feature vectors; inputting the set of node feature vectors into at least a matching output head and a frame allocation output head to obtain at least a matching score matrix and a frame allocation matrix; and calculating a training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels, and updating the graph neural network and the matching output head and the frame allocation output head through backpropagation of the training loss function.

[0009] In the above training method for graph neural networks used for atomic array rearrangement, the training label is at least one of optimal pairing, frame execution number, compatibility matrix, and optimal number of frames.

[0010] In the above-described training method for a graph neural network for atom array rearrangement, obtaining the initial two-dimensional grid array of atoms for training and the training labels based on the rearrangement results of the initial two-dimensional grid array includes: constructing an optimal pairing list of the optimal pairings of each atom in the rearrangement results of the initial two-dimensional grid array as the pairing loss supervision label of the graph neural network; and using the compatibility matrix based on whether all atoms in the rearrangement results of the initial two-dimensional grid array are executed in the same frame as the compatibility loss supervision label of the graph neural network.

[0011] In the above-described training method for a graph neural network for atom array rearrangement, obtaining the initial two-dimensional grid array of atoms for training and the training labels based on the rearrangement results of the initial two-dimensional grid array further includes: constructing a frame execution number list from the frame execution numbers of each atom in the rearrangement results of the initial two-dimensional grid array as a supervision label for the frame allocation loss of the graph neural network.

[0012] In the above-described training method for graph neural networks used for atomic array rearrangement, the feature extraction backbone network of the graph neural network includes multiple graph attention layers, and each graph attention layer includes multiple attention heads.

[0013] In the above-described training method for graph neural networks used for atomic array rearrangement, the set of node feature vectors includes an initial node feature vector subset corresponding to the initial position of the atom and a target node feature vector subset corresponding to the target position of the atom.

[0014] In the above training method for graph neural networks used for atomic array rearrangement, the process of inputting the set of node feature vectors into the matching output head to obtain the matching score matrix includes: for each atom-target pair, concatenating the initial node feature vector in the corresponding initial node feature vector subset and the target node feature vector in the target node feature vector subset, and mapping them to a scalar matching score through a fully connected network; and forming the matching score matrix by assembling the scalar matching scores corresponding to all atom-target pairs.

[0015] In the above training method for graph neural networks used for atomic array rearrangement, inputting the set of node feature vectors into at least the matching output head and the frame allocation output head to obtain at least the matching score matrix and the frame allocation matrix further includes: inputting the set of node feature vectors into the compatibility output head to obtain the compatibility prediction matrix.

[0016] In the above training method for graph neural networks used for atomic array rearrangement, the process of inputting the set of node feature vectors into the compatibility output head to obtain the compatibility prediction matrix includes: concatenating every two initial node feature vectors in the subset of initial node feature vectors and then activating them through a fully connected network and a Sigmoid activation function to obtain compatibility probability values; and forming the compatibility prediction matrix by combining the compatibility probability values ​​corresponding to all initial node pairs.

[0017] In the above-described training method for graph neural networks used for atomic array rearrangement, calculating the training loss function based at least on the matching score matrix, the frame assignment matrix, and the training labels includes: inputting the matching score matrix into a Sinkhorn operator to obtain a double random matrix; and calculating the negative log-likelihood of the double random matrix with the optimal pairing in the training labels as the supervision target to obtain the matching loss function.

[0018] In the above-described training method for graph neural networks used for atomic array rearrangement, calculating the training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels includes: calculating the binary cross-entropy element-wise between the compatibility prediction matrix and the compatibility matrix in the training labels; and averaging all matrix elements to obtain the compatibility loss function.

[0019] In the above training method for graph neural networks used for atomic array rearrangement, calculating the training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels includes: calculating the mean of the matrix elements corresponding to the atomic pairs in the compatibility prediction matrix that have intersecting path spaces to obtain the path intersection penalty loss function.

[0020] In the above-described training method for graph neural networks used for atomic array rearrangement, calculating the training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels includes: calculating a weighted sum of the matching loss function, the compatibility loss function, the path crossover penalty loss function, the frame allocation loss function, and the differentiable frame number loss function to obtain the training loss function.

[0021] In the above-described training method for a graph neural network used for atomic array rearrangement, updating the graph neural network, the matching output head, and the frame allocation output head via backpropagation of the training loss function includes: using a course learning framework that progressively introduces various loss functions and annealing temperature parameters in a multi-stage manner, and updating the graph neural network, the matching output head, and the frame allocation output head via backpropagation of the training loss function.

[0022] In the above training method for graph neural networks used for atomic array rearrangement, the annealing temperature parameters in the course learning framework include: the Sinkhorn temperature and frame allocation temperature in the Sinkhorn operator during synchronous annealing.

[0023] According to another aspect of this application, an atomic array rearrangement method using a graph neural network is provided, comprising: constructing a structured input graph based on an initial two-dimensional grid array of atoms to be rearranged; inputting the structured input graph into a feature extraction backbone network of a graph neural network trained by the graph neural network training method for atomic array rearrangement as described above to obtain a set of node feature vectors; inputting the set of node feature vectors into a matching output head and a frame allocation output head trained by the graph neural network training method for atomic array rearrangement as described above to obtain a matching score matrix and a frame allocation matrix; performing matching decoding on the matching score matrix to obtain an optimal pairing, and performing argmax decoding on the frame allocation matrix to obtain a frame execution number; and moving atoms from an initial position to a target position based on the optimal pairing and the frame execution number.

[0024] According to another aspect of this application, an electronic device is provided, including a processor coupled to a memory, the processor being configured to execute a computer program stored in the memory, such that the electronic device performs the training method for a graph neural network for atomic array rearrangement and the atomic array rearrangement method using the graph neural network as described above.

[0025] According to another aspect of this application, a computer-readable storage medium is provided, the computer-readable storage medium storing a computer program that, when the computer program is run, implements the training method for a graph neural network for atomic array rearrangement and the atomic array rearrangement method using the graph neural network as described above.

[0026] The atomic array rearrangement scheme based on graph neural networks provided in this application can decode the direct output of atom-target pairing results and the frame execution number of the atom by training a multi-head graph neural network, thereby achieving path conflict-free atomic array rearrangement with low inference latency and low frame count. Attached Figure Description

[0027] Various other advantages and benefits of this application will become apparent to those skilled in the art upon reading the detailed description of the preferred embodiments below. The accompanying drawings are for illustrative purposes only and are not intended to limit the scope of this application. It is obvious that the drawings described below are merely some embodiments of this application, and those skilled in the art can obtain other drawings based on these drawings without any inventive effort. Furthermore, the same reference numerals denote the same parts throughout the drawings.

[0028] Figure 1 The illustration shows a schematic flowchart of a training method for a graph neural network for atomic array rearrangement according to an embodiment of this application.

[0029] Figure 2 The illustration shows a schematic flowchart of the label generation process in a training method for a graph neural network for atomic array rearrangement according to an embodiment of this application.

[0030] Figure 3 A schematic flowchart illustrating an atomic array rearrangement method using a graph neural network according to an embodiment of this application is shown.

[0031] Figure 4 A schematic block diagram of an electronic device according to an embodiment of this application is shown. Detailed Implementation

[0032] Hereinafter, exemplary embodiments according to this application will be described in detail with reference to the accompanying drawings. Obviously, the described embodiments are merely some embodiments of this application, and not all embodiments of this application. It should be understood that this application is not limited to the exemplary embodiments described herein.

[0033] Figure 1 A schematic flowchart illustrating a training method for a graph neural network for atomic array rearrangement according to an embodiment of this application is shown. Figure 1As shown, the training method for a graph neural network for atomic array rearrangement according to an embodiment of this application includes the following steps.

[0034] S110, Obtain the initial two-dimensional grid array of atoms for training and the training labels based on the rearrangement results of the initial two-dimensional grid array.

[0035] In this embodiment, a time-spreading graph can be constructed based on the initial two-dimensional grid array of atoms used for training. Based on the time-spreading graph, the optimal flow scheme from the initial set of atom positions to the target set of atom positions is obtained using minimum cost maximum flow solution. Then, based on the optimal flow scheme and the time-spreading graph, minimum time search is performed to obtain the spatiotemporal path of each atom from its initial position to its target position as the rearrangement result of the initial two-dimensional grid array. The optimal pairing, frame execution number, compatibility matrix, and optimal frame number are extracted from the spatiotemporal path of each atom from its initial position to its target position as training labels for training the graph neural network.

[0036] Figure 2 The illustration shows a schematic flowchart of a label generation process in a graph neural network for atomic array rearrangement according to an embodiment of this application. Figure 2 As shown, the label generation process includes the following steps.

[0037] Step S111: Construct a time-spreading graph based on the initial two-dimensional grid array of atoms, the target atom position set, and the number of candidate frames.

[0038] In this embodiment, an initial physical grid map can be determined based on an initial two-dimensional grid array of atoms, such as a 10×10 array, as described above. And the initial set of atomic positions in the aforementioned initial physical lattice diagram, for example, the initial atomic positions of 60 atoms in the 10×10 array when the loading probability is 0.6. Furthermore, a target set of atomic positions is determined as the rearrangement target, for example, the target atomic positions of 36 atoms in a 6×6 array, as described above. , Furthermore, the number of candidate frames in the time spread map is determined, for example, denoted as... .

[0039] To construct a time-extended graph (TEG), on the time axis Every time step Above, extract each initial physical grid point from the initial physical grid map and arrange them at each time step of the time spread map. The corresponding spatial locations are used to form each spatiotemporal node of the time extension diagram, for example, denoted as... And the set of spatiotemporal nodes is represented as .

[0040] Here, to eliminate the risk of collisions between atoms during array rearrangement, the current time step is... spatiotemporal nodes It is split into two state nodes: an ingress node, for example, represented as... and output nodes, for example, represented as Furthermore, at the ingress node To the outgoing node Add an internal edge as a directed connection and assign a capacity limit to this edge. Specifically, based on the actual collision constraints, the capacity limit can be set to greater than or equal to one. A capacity limit of one represents the same time step... Same spacetime node It can hold a maximum of one atom.

[0041] Here, we consider the exclusivity of atomic rearrangement, meaning there's an upper limit to the number of atoms each optical tweezers can capture or hold at any given time, and consequently, an upper limit to the number of atoms each physical grid point can hold at a specific time. For network flow algorithms, such as the minimum-cost maximum flow algorithm, since their flow constraints are applied to edges by default, rather than node flow, without node splitting, while the flow limit for each edge can be limited, multiple edges from different sources (e.g., atoms moved in from the left, atoms moved in from the right, or atoms left in place from the previous time step) might simultaneously point to the same node, leading to atomic collisions. By splitting each spatiotemporal node into an ingress node and an egress node and establishing an internal edge between them, the number of atoms held by a grid point at a given time can be transformed into the flow through this internal edge. Thus, the capacity limit of the internal edge limits the node's capacity. In other words, regardless of how many edges converge on the ingress node, there is an upper limit to the number of atoms that can ultimately pass through the internal edge and enter the egress node for scheduling in the next time step, preventing atoms exceeding the capacity limit from being located at the same grid point at the same time.

[0042] Furthermore, without node splitting, when the paths of two atoms intersect at a specific grid point at a particular moment, it becomes impossible to distinguish which atom came from which path and which path it is going to when extracting the path. However, by using the constraints of the ingress node-internal edge-exgress node, combined with the maximum flow algorithm, it can be ensured that each streamline is independent and continuous in the spatiotemporal dimension, and will not be randomly allocated after flow merging at a certain node. In addition, by constructing internal edges, penalty weights (i.e., setting costs) can be directly applied to the internal edges for scheduling costs of certain grid points (e.g., in areas with unstable light fields where atoms are not desired to stay), thereby enhancing the flexibility of path planning.

[0043] In the embodiments of this application, in the scenario of atomic rearrangement, if the experimental conditions allow the same lattice point to accommodate multiple atoms at the same time (for example, to achieve multi-atom co-location manipulation by means of multiple independent optical tweezers arrays), the capacity limit of the internal edge can be set to be greater than one, so as to relax the upper limit of the atom capacity of the same spatiotemporal node and thus improve the scheduling flexibility.

[0044] Furthermore, in addition to each spatiotemporal node In addition to the internal edges, waiting edges and moving edges are constructed between adjacent time steps in the time extension graph. The waiting edge is used to move time steps... Corresponding output node With time step Corresponding same spacetime node Ingress node The directed connections between them, similar to internal edges, have an edge capacity limit of one or more, representing the atom waiting for one time step at the current spacetime node. Furthermore, moving edges involves moving the time step... At the spatiotemporal node Corresponding output node Link to time step At the spatiotemporal node Adjacent spatiotemporal nodes Ingress node Furthermore, an edge capacity limit of one or more is set, which indicates that the atom moves to an adjacent spatiotemporal node.

[0045] In this embodiment of the application, it is assumed that the spatiotemporal node The two-dimensional coordinates are Then adjacent spatiotemporal nodes The two-dimensional coordinates can be , , , That is, adjacent spatiotemporal nodes in the vertical, horizontal, and vertical directions. Furthermore, in this embodiment of the application, adjacent spatiotemporal nodes... It can also be a spacetime node. At obliquely adjacent spatiotemporal nodes, i.e., their two-dimensional coordinates can be , , , Furthermore, those skilled in the art will understand that the adjacent spatiotemporal nodes It can also be related to the aforementioned spatiotemporal node. Spatiotemporal nodes whose distance (e.g., Manhattan distance) is less than a predetermined threshold. Additionally, in this embodiment, the constructed moving edge can also be a row / column parallel moving edge structure designed for an acousto-optic deflector (AOD) platform.

[0046] In addition, outside each spatiotemporal node of the spatiotemporal extension graph, two more global virtual control nodes and a super source node can be instantiated. and Super Exchange Furthermore, constructing a supersource node. Pointing to time step Ingress node of all spatiotemporal nodes The source edge, here, has a capacity of one to map the initial supply position of the atoms. Additionally, the construction time step... All spatiotemporal nodes' output nodes Pointing to Super Exchange The sink edge, here, has a capacity of one to map the target required location of the atom.

[0047] Here, network flow algorithms, such as standard minimum-cost maximum flow algorithms (continuous shortest path algorithms, network simplex methods), typically require the network flow model to contain a unique starting point and a unique sink for all flows. Furthermore, in the actual physical scenario of atomic rearrangement, because the distribution of atoms is highly discrete—initially scattered across dozens or even hundreds of different grid points in the array, and at the target time, the empty defect grid points to be filled are also scattered across different locations in the array—can be transformed into a standard network flow problem by introducing virtual super-source and super-sink nodes.

[0048] Of course, those skilled in the art will understand that if the network flow algorithm used does not require a unique starting point and a unique converging point for all traffic, the constructed time-spread graph can also include only internal edges, waiting edges, and moving edges. Furthermore, the capacity of internal edges, waiting edges, and moving edges can be greater than one to support parallel path scenarios.

[0049] Therefore, for the number of candidate frames A spatiotemporal extension graph as a directed graph was constructed. The node set is represented as and node set The number of spatiotemporal nodes in it is And edge set This includes internal edges, waiting edges, moving edges, source edges, and sink edges, as described above. The spatiotemporal expansion graph introduces a time dimension into the initial two-dimensional grid array of atoms, which serves as a static topology, based on the spatial positions of the grid points, thereby avoiding collisions. That is, if paths are planned on a two-dimensional graph solely based on grid spatial positions, it is difficult to determine whether two atoms will collide midway, because the essence of a collision is spatiotemporal overlap—that is, two atoms appearing in the same spatial position at the same time, or traversing the same path within the same time period. Therefore, by discretizing time into individual time steps using the spatiotemporal expansion graph, and replicating the two-dimensional grid array for each time step, each spatiotemporal node in the spatiotemporal expansion graph possesses not only a spatial position but also a temporal position. Thus, in a three-dimensional spatiotemporal expansion graph, collisions can be avoided through upper limits on edge capacity. For example, setting the upper limit of the capacity for internal edges and waiting edges to one means that at any given time step, any spatial position can accommodate at most one atom, while setting the upper limit of the capacity for moving edges to one means that the physical path from the current position to the adjacent position in the next time step can only pass through one atom.

[0050] In this way, when applying network flow algorithms such as minimum-cost maximum-flow algorithms, the optimal path found, such as the maximum-flow path, can avoid atomic collisions. That is, if the algorithm finds that a spatiotemporal node is already occupied by another atom while searching for a path, it will automatically make the atom wait in place or take a detour; otherwise, it can only move atoms one by one in a serial manner. Thus, since the time-spreading graph essentially exhaustively enumerates all the potential destinations of all atoms at all future moments, it is possible to schedule the movement routes of all atoms at a global level, realizing the parallel manipulation of multiple atoms at the same time, thereby significantly reducing the physical time required for rearrangement.

[0051] Step S112: Solve the time-spread graph using minimum cost maximum flow to obtain the optimal flow scheme from the initial set of atomic positions to the target set of atomic positions.

[0052] In this embodiment, when using the Minimum Cost Maximum Flow (MCMF) algorithm, the edge costs in the time-spread graph first need to be set. In this embodiment, the cost of waiting edges is set to be greater than or equal to zero, the cost of moving edges is set to be greater than zero, and the cost of internal edges is set to zero. That is, in the time-spread graph, not only are there edges connecting nodes, but cost-based network flow algorithms can also be applied by setting a cost for each edge. For example, if an atom moves from its current position in the current time step to an adjacent position in the next time step, a movement cost of a certain distance is incurred, while if the atom remains in place, a waiting cost is incurred. Furthermore, those skilled in the art will understand that in this embodiment, the cost of the moving edges can be set non-uniformly; for example, moving edges in different directions or at different distances have different costs.

[0053] Then, a flow conservation equation is established for each spatiotemporal node in the time-extended graph, where the inflow equals the outflow, and a capacity constraint is constructed to ensure that the capacity of each edge in the time-extended graph is less than its upper capacity limit. The minimum-cost maximum flow algorithm takes maximizing the flow from the initial set of atomic positions to the target set of atomic positions as its first objective (ensuring the target grid is filled), and minimizing the sum of the products of flow and cost as its second objective. Furthermore, for network flow algorithms that require the network flow model to contain a unique starting point for all flows and a unique sink point for all flows, the first objective is to maximize the flow from the super source node to the super sink node (ensuring the target grid is filled), and the second objective is to minimize the sum of the products of flow and cost.

[0054] The minimum cost maximum flow algorithm can employ continuous shortest path algorithm (SSP), network simplex method, Dijkstra's potential function method, Bellman-Ford / SPFA augmenting path algorithm, etc., to iteratively find the shortest path from the initial set of atomic positions to the target set of atomic positions (or from the super source to the super sink) in the spatiotemporal extended graph based on cost until the maximum flow is reached. Specifically, the minimum cost maximum flow algorithm first performs residual network initialization and edge cost mapping, that is, the initial flow of all edges in the spatiotemporal extended graph is set to zero, and for each edge in the spatiotemporal extended graph, a forward edge (capacity equal to the edge capacity) and a reverse edge (capacity set to 0) are established. Then, the edge cost is associated with the forward edge, and its opposite (negative cost) is associated with the reverse edge to realize the cost accounting of the backtracking path, thereby constructing the initial residual network.

[0055] Then, the minimum cost maximum flow algorithm iteratively optimizes the shortest augmenting path. Starting from the initial set of atomic positions (or the super source node) and ending at the target set of atomic positions (or the super sink node), it uses a shortest path algorithm (such as SPFA or Dijkstra's algorithm with a potential function) to find the augmenting path with the minimum total cost. It then finds the edge with the minimum capacity constraint on this path and records its capacity as the flow. Simultaneously, it deducts the capacity of forward edges in the initial residual network based on the recorded flow and increases the capacity of the corresponding reverse edges. This optimization and update process is repeated until there is no connected path from the initial set of atomic positions to the target set of atomic positions in the residual network. Finally, it records the flow status of all edges after the iteration, forming a cumulative flow distribution graph.

[0056] Finally, the minimum cost maximum flow algorithm checks each spatiotemporal node in the cumulative flow distribution graph to verify whether the sum of its inflows equals the sum of its outflows, ensuring compliance with physical conservation. It then multiplies the instantaneous flow on all edges by the corresponding initial rate, sums them up to obtain the global minimum cost, and extracts the edge sequence with a flow value of one. Finally, it outputs the globally collision-free and lowest cost atomic scheduling scheme as the optimal flow scheme.

[0057] Step S113: Perform a minimum time search based on the optimal flow scheme and the time spread graph to obtain the spatiotemporal path of each atom from its initial position to its target position.

[0058] That is, at the current candidate frame number Below, because if the current candidate frame number If the optimal flow scheme is feasible (i.e., satisfies atomic array rearrangement), then the number of frames... It is definitely feasible (it can be considered that each path can be constructed by extending it one step along the waiting edge). (A feasible flow), therefore the number of candidate frames can be continuously adjusted using a binary search method. The value is then repeated using steps S111 and S112.

[0059] Here, considering monotonicity as the mathematical foundation for the validity of binary search, we have the following proposition (feasibility is monotonically non-decreasing): If the number of candidate frames... The maximum flow of the minimum cost maximum flow is equal to the target number of atoms. Then the number of candidate frames There must exist a maximum flow equal to A feasible streaming scheme. The proof is as follows: Take the number of frames. The optimal flow scheme is as follows: for the terminal node of each atomic spacetime path, extend the path one step along the waiting edge to the next node. Frames. After extension, the capacity constraints of all internal edges, waiting edges, and moving edges are not violated, and the flow conservation condition remains valid. Therefore, the constructed scheme is [frame number missing]. The maximum flow is still equal to the legal feasible flow. .

[0060] If the current candidate frame number If the optimal flow scheme is feasible, i.e., the maximum flow equals the target number of atoms (the number of target atom position sets), then try a smaller number of candidate frames. Otherwise, increase the number of candidate frames. until the minimum number of candidate frames that meet the conditions is determined. Furthermore, in the embodiments of this application, based on the mathematical proof of monotonicity described above, in addition to binary search, other monotonic search strategies such as linear search can also be used.

[0061] When performing a binary search, a time interval range can be set, i.e. For example, searching the lower bound Set as The default value is 1, and the upper bound for the search is... Set as , defaults to Each time, the midpoint of the interval is selected to construct a time spread graph and determine the feasibility of the network flow. Utilizing the monotonically non-decreasing property of feasibility with respect to the number of candidate frames (if a larger number of frames is feasible, then a larger number of frames is necessarily feasible), the search interval is halved. This process is iterated until the interval converges, returning the minimum feasible number of frames, i.e., the optimal number of frames. Thus, the execution... The minimum cost maximum flow algorithm can obtain the optimal number of frames, for example, denoted as... And obtain the optimal number of frames. The optimal flow scheme under the given conditions.

[0062] Based on the optimal number of frames By performing path-by-path tracing under the optimal flow scheme and tracing the path of each atom from its starting position in the optimal flow in the order of time steps, the spatiotemporal path of each atom from its initial position to its target position can be obtained.

[0063] Step S114: Extract the optimal pairing, frame execution number, compatibility matrix, and optimal frame number from the spatiotemporal path of each atom from the initial position to the target position as training labels for training the graph neural network.

[0064] In other words, after obtaining the spatiotemporal path of each atom from its initial position to its target position, the optimal pairing can be obtained as precise integers by extracting the target grid point indices of the path's starting and ending points, for example, denoted as... Furthermore, by extracting the time steps that identify moving edges on the path (e.g., the first or last occurrence of a moving edge), a frame execution number, denoted as a precise integer, can be obtained, for example, as Batch_id. Additionally, a compatibility matrix can be extracted; here, if for all atomic pairs... If the execution numbers of the two frames are the same, the matrix element is one; otherwise, it is zero. The compatibility matrix is: A symmetric Boolean matrix. Furthermore, the compatibility matrix can be further extended to a soft compatibility fraction (real-valued) form to encode temporal proximity relationships between different frames. For example, the soft compatibility fraction is a decay function based on the difference in frame execution numbers; specifically, atomic pairs... The software compatibility score is ,in This is the attenuation coefficient.

[0065] In this way, it can be done through, for example Figure 2 The label generation process shown obtains a training label for each atom in the initial two-dimensional grid array of the atoms used for training, corresponding to the training of the graph neural network. For example, the optimal pairing, frame execution number, compatibility matrix, and optimal frame number can form a quadruple label for training the graph neural network. Of course, those skilled in the art will understand that the optimal pairing, frame execution number, compatibility matrix, and optimal frame number can also be used individually or in any combination as training labels for training the neural network model.

[0066] For all atoms in the atomic array, we obtain an optimal pairing list consisting of the optimal pairings for each atom, denoted as: It is a list of integers with a length of 1. And each element in it represents an atom. Assign to target The optimal pairing number is used as the supervision label for the pairing loss of the graph neural network. Furthermore, the list of frame execution numbers is represented as... Also a list of integers, with a length of... Each element represents an atom. Execution frame number And the range of values ​​is , used as supervision labels for frame assignment loss in graph neural networks. Compatibility matrices, for example, in atomic pairs If the frame execution numbers are the same, the matrix element is one; otherwise, it is zero. This is a Boolean type matrix, and its elements represent atoms. With atoms Can it be executed within the same frame, serving as a compatibility loss supervision label for a graph neural network? And, what is the optimal number of frames? It is an integer type number used as an evaluation / frame differentiable reference for graph neural networks.

[0067] Furthermore, those skilled in the art will understand that, in the embodiments of this application, the rearrangement results obtained from the initial two-dimensional grid array of atoms used for training can also be used as training labels through other means. For example, other atom-target pairing algorithms such as the Hungarian algorithm, other frame scheduling methods such as greedy graph coloring, or other path planning methods and deep learning methods can be used. Moreover, in the embodiments of this application, based on the functionality of the neural network model, only one or more of the optimal pairing, frame execution number, compatibility matrix, and optimal number of frames can be used. For example, if the graph neural network only includes a matching output header and a frame allocation output header but not a compatibility output header, the training labels can also only include the optimal pairing list, the frame execution number list, and the optimal number of frames. As another example, if a neural network model other than a graph neural network is used only to predict the optimal pairing or the frame execution number, the training labels can also include the optimal pairing list or the frame execution number list.

[0068] That is, in the training method for a graph neural network for atom array rearrangement according to the embodiments of this application, obtaining the initial two-dimensional grid array of atoms for training and the training label based on the rearrangement result of the initial two-dimensional grid array includes: constructing an optimal pairing list of the optimal pairings of each atom in the rearrangement result of the initial two-dimensional grid array as a pairing loss supervision label for the graph neural network; and constructing a frame execution number list of the frame execution number of each atom in the rearrangement result of the initial two-dimensional grid array as a frame allocation loss supervision label for the graph neural network.

[0069] Furthermore, in the above-described training method for a graph neural network for atomic array rearrangement, obtaining the initial two-dimensional grid array of atoms for training and the training label based on the rearrangement result of the initial two-dimensional grid array further includes: using the compatibility matrix of whether all atoms in the rearrangement result based on the initial two-dimensional grid array are executed in the same frame as the compatibility loss supervision label of the graph neural network.

[0070] Step S120: Input the initial two-dimensional grid array of atoms used for training into the feature extraction backbone network of the graph neural network to obtain a set of node feature vectors.

[0071] Here, the initial two-dimensional lattice array of atoms used for training is Array, which is composed of all arrays Each grid point is a node, and the node input features include normalized coordinates and atom / target state markers. The edge structure can be based on the neighborhood relationship described above, such as taking the 4-neighborhood connection relationship of up, down, left, and right (of course, those skilled in the art will understand that this can also be extended to a larger neighborhood in the manner described above).

[0072] Therefore, the structured input graph of the graph neural network is first constructed based on the initial two-dimensional lattice array of atoms used for training. That is, the structured input graph of the graph neural network is constructed based on the initial two-dimensional lattice array of atoms used for training. Each grid point in the array serves as a node in the structured input graph, and a node input vector is generated for each node. This node input vector contains the grid point's normalized row coordinates, normalized column coordinates, atom occupancy flag (e.g., 1 indicates an atom is present, 0 indicates no atom is present), and target position flag (e.g., 1 indicates a target position, 0 indicates a non-target position). Edge connections between nodes are established based on neighborhood relationships, such as 4-neighborhood relationships (i.e., top, bottom, left, and right adjacent grid points). Thus, a graph containing... A structured input graph is a graph data structure consisting of nodes, their node input vectors, and edge connections.

[0073] Then, the structured input graph is input into the feature extraction backbone network of the graph neural network for feature extraction to obtain a set of node feature vectors. In this embodiment, the feature extraction backbone network of the graph neural network can adopt the GATv2 architecture (or GAT, GraphSAGE, GIN or other message passing networks), and is configured with Layered attention layer, where The range is, for example, 2 to 32, by default. The more layers a graph neural network has, the stronger its representational power. Furthermore, each graph attention layer is configured with... Each attention head, where H ranges from 1 to 16, for example, by default. The hidden dimension of the feature extraction backbone network of the graph neural network, i.e., the dimension of the output node feature vector, is set to... , The range is, for example, 32 to 1024, by default. After the structured input graph is input into the feature extraction backbone network of the graph neural network, multi-head attention aggregation is performed layer by layer. That is, in each layer, the embedding vector corresponding to each node is processed by... Each attention head aggregates information from its neighboring nodes and updates its own embedding vector, thereby... After layer message passing, each node obtains a high-dimensional embedding vector that incorporates global context information, i.e., the node feature vector. This results in a set of node feature vectors, for example, represented as... ,in .

[0074] Furthermore, in this embodiment, the node feature vector set includes an initial node feature vector subset. (common (feature vectors) and a subset of feature vectors of the target node (common (Eigenvectors).

[0075] That is, in the training method of the graph neural network for atomic array rearrangement according to the embodiments of this application, the feature extraction backbone network of the graph neural network includes multiple graph attention layers, and each graph attention layer includes multiple attention heads.

[0076] Furthermore, in the training method for a graph neural network for atomic array rearrangement according to the embodiments of this application, the set of node feature vectors includes an initial node feature vector subset corresponding to the initial position of the atom and a target node feature vector subset corresponding to the target position of the atom.

[0077] Step S130: Input the set of node feature vectors into at least the matching output header and the frame allocation output header to obtain at least the matching score matrix and the frame allocation matrix.

[0078] Specifically, when the node feature vector set is input to match the output head, firstly for each atom-target pair... , initial node feature vector and target node feature vector Concatenate to obtain the spliced ​​vector Then, it is mapped to a scalar matching score through a fully connected network. Thus, by traversing all Each atom-target pair yields the matching score matrix. .

[0079] When the set of node feature vectors is input into the frame allocation header, for each initial node feature vector... Through a depth of Multilayer perceptron (MLP) independently mapped to The dimensional logit space is used to obtain the frame allocation logit vector. In this embodiment, the depth of the multilayer perceptron ranges from 1 to 8, with a default value of 2. It can be set to any number greater than or equal to the optimal frame rate. Positive integers, for example, those in the range of 4 to 10. Preferably ,in Optimal number of frames The mathematical expectation is often expressed using the statistical mean in practice, and For safety margin (e.g.) ),For example Traverse all 1 initial atom, to obtain the frame allocation matrix .

[0080] Those skilled in the art will understand that It needs to be greater than or equal to the maximum number of frames that may actually occur, by taking... This can cover the distribution range of frame counts in most samples, plus a safety margin. To handle extreme situations. During reasoning, only each atom needs to be considered. The frame number can be obtained directly without any post-processing.

[0081] Furthermore, the set of node feature vectors can be further input into the compatible output header. Here, firstly, the two initial node feature vectors in the initial node feature vector subset are... and Concatenate to obtain the spliced ​​vector Then, it is mapped through a fully connected network and activated by a Sigmoid activation function to output a compatibility probability value. traverse all With atom pairs, the compatibility prediction matrix is ​​obtained. .

[0082] That is, in the training method of the graph neural network for atomic array rearrangement according to the embodiments of this application, inputting the set of node feature vectors into the matching output head to obtain the matching score matrix includes: for each atom-target pair, concatenating the initial node feature vector in the corresponding initial node feature vector subset and the target node feature vector in the target node feature vector subset, and mapping them to a scalar matching score through a fully connected network; and forming the matching score matrix by assembling the scalar matching scores corresponding to all atom-target pairs.

[0083] In the training method for the graph neural network used for atomic array rearrangement, the step of inputting the set of node feature vectors into the frame allocation head to obtain the frame allocation matrix includes: mapping each initial node feature vector in the subset of initial node feature vectors to a frame allocation logical value vector through a multilayer perceptron of a predetermined depth; and arranging the frame allocation logical value vectors corresponding to all initial node feature vectors into the frame allocation matrix.

[0084] Furthermore, in the above-described training method for graph neural networks used for atomic array rearrangement, the frame allocation logic value vector is based on the mathematical expectation of the optimal number of frames plus a safety margin.

[0085] The training method for the graph neural network used for atomic array rearrangement further includes: inputting the set of node feature vectors into a compatible output head to obtain a compatible prediction matrix.

[0086] Furthermore, in the above-described training method for graph neural networks used for atomic array rearrangement, inputting the set of node feature vectors into at least the matching output head and the frame allocation output head to obtain at least the matching score matrix and the frame allocation matrix further includes: inputting the set of node feature vectors into the compatibility output head to obtain the compatibility prediction matrix.

[0087] In the above training method for graph neural networks used for atomic array rearrangement, the process of inputting the set of node feature vectors into the compatibility output head to obtain the compatibility prediction matrix includes: concatenating every two initial node feature vectors in the subset of initial node feature vectors and then activating them through a fully connected network and a Sigmoid activation function to obtain compatibility probability values; and forming the compatibility prediction matrix by combining the compatibility probability values ​​corresponding to all initial node pairs.

[0088] Step S140: Calculate a training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels, and update the graph neural network, the matching output head, and the frame allocation output head through backpropagation of the training loss function.

[0089] As described above, the training labels are the optimal pairing list, the frame execution number list, and the optimal frame number, and preferably further include a compatibility matrix. Based on the predicted values ​​of the matching score matrix, the frame allocation matrix, and preferably the compatibility prediction matrix, the training labels are used as ground truth values ​​to calculate the training loss function. In this embodiment, the training loss function includes at least a plurality of loss function terms based on the matching score matrix, the frame allocation matrix, and the training labels.

[0090] The paired loss function uses paired negative log-likelihood loss, for example, denoted as... Match the score matrix Input the Sinkhorn operator, with the Sinkhorn temperature parameter and number of iterations (Adjustable, range from 5 to 200, default 30) Perform alternating row and column normalization to obtain a double random matrix. Among them, the Sinkhorn temperature parameter It is static or adjustable, and preferably annealed from 1.0 to 0.05 in the case of iterative annealing over a periodic period.

[0091] Specifically, the Sinkhorn operator first performs temperature scaling on the matching score matrix to obtain the scaled kernel matrix, represented as:

[0092] Taking the exponent element-wise of the scaled kernel matrix yields the initial matrix, represented as:

[0093] implement This process involves alternating row and column normalization iterations. Each iteration comprises two operations: row normalization divides each row of the matrix by the sum of its elements, making the sum of each row equal to 1; column normalization divides each column of the matrix by the sum of its elements, making the sum of each column equal to 1. This process is repeated for each row and column. Then the initial matrix It converges to an approximately double-random matrix, i.e., the output... .

[0094] The best pairing from the training labels To monitor the target, the negative log-likelihood loss is calculated and expressed as:

[0095] in This is the set of valid pairs based on the optimal pairing list. For numerically stable terms, such as those ranging from 0 to... ,default .

[0096] Here, a double random matrix The Element represents an atom Assigned to target The soft probability, and the training labels Each atom is given Optimal target number The goal of negative log-likelihood loss is to make each atom in the soft probability matrix predicted by the model... In its correct goal The probability value at a given position should be as high as possible. Therefore, for the effective pairing set... Each atom in Double random matrix The In-row column index is The element values ​​are the atoms predicted by the model. Assign to its correct target The closer the probability is to 1, the better the model is for atoms. The more accurate the pairing prediction, and the closer the value is to 0, the better the model predicts the atoms. It was incorrectly assigned to another target.

[0097] That is, in the training method for a graph neural network for atomic array rearrangement according to the embodiments of this application, calculating the training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels includes: inputting the matching score matrix into a Sinkhorn operator to obtain a double random matrix; and calculating the negative log-likelihood of the double random matrix with the optimal pairing in the training labels as the supervision target to obtain the matching loss function.

[0098] The compatibility loss function uses the binary cross-entropy loss of the compatibility matrix, for example, denoted as... Specifically, the compatible prediction matrix will be used. Compatibility matrix with training labels Calculate the binary cross-entropy (BCE) element by element, and for all The average of the elements is expressed as:

[0099] in:

[0100] In this embodiment of the application, in order to penalize path crossing, a path crossing penalty loss can be introduced based on the compatibility prediction matrix, for example, denoted as... Specifically, based on the current predicted pairing results, a set of atomic pairs with intersecting path spaces is detected. .like If there exist intersecting atomic pairs in the path space, then the average of the predicted values ​​of the conflicting atomic pairs in the compatibility prediction matrix is ​​taken as a penalty; if If the loss is zero, then the loss is zero, expressed as:

[0101] Thus, for the atomic pairs in the current predicted pairings, there are intersecting path spaces. The penalty compatibility output header gives a high compatibility score. That is, based on the matching score matrix... The soft allocation after Sinkhorn normalization is:

[0102] Then use the current position of the atom to its predicted target location The straight line segments are taken as the trajectories. If the line segments corresponding to the trajectories of two atoms intersect in Euclidean space, then the pair of atoms is added to the penalty set. And apply a cross-penalty loss to its predicted pairings. .

[0103] That is, in the training method for a graph neural network for atomic array rearrangement according to the embodiments of this application, calculating the training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels includes: calculating the binary cross-entropy element-wise between the compatibility prediction matrix and the compatibility matrix in the training labels; and averaging all matrix elements to obtain the compatibility loss function.

[0104] Furthermore, in the above-described training method for graph neural networks used for atomic array rearrangement, calculating the training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels includes: calculating the mean of the matrix elements corresponding to atomic pairs with intersecting path spaces in the compatibility prediction matrix to obtain the path intersection penalty loss function.

[0105] The frame allocation loss function uses the frame allocation direct supervision loss, for example, denoted as... Specifically, the frame allocation matrix is... Divide by frame allocation temperature Then, the softmax function is used to obtain the soft probability vector for each frame, and the cross-entropy loss is calculated for each atom using the frame execution number (batch_id) in the training labels as the multi-class supervision target. Finally, the cross-entropy loss is calculated for all... The average of the atoms is expressed as:

[0106] Here, for batch_id exceeding... In this case, cut it off to Furthermore, frame allocation temperature. It can also be static or adjustable, for example, annealing from 1.0 to 0.3 by default.

[0107] That is, for atoms The frame allocation vector is Temperature allocated via frame After scaling, the softmax operation yields a frame assignment soft probability vector representing the frame affiliation probability distribution, with the training label being the correct frame number for that atom. ,Right now:

[0108] Then we have:

[0109] Expand to equivalent form:

[0110] get:

[0111] Furthermore, in this embodiment, a loss for minimizing the number of differentiable frames is introduced based on the frame allocation matrix. Specifically, the frame allocation matrix is ​​first... Divide by frame allocation temperature Then, the softmax function is used to obtain the frame allocation soft probability matrix. Then, for each frame ( ), calculate an approximate probability that the frame is used by at least one atom. This is achieved through a product approximation (or other differentiable approximations that make the discrete frame count differentiable), which involves multiplying the unused probabilities of all atoms in that frame and then subtracting that product from 1. Finally, the approximate values ​​of the usage probabilities of all frames are summed as the differentiable frame count loss, expressed as:

[0112] That is, Indicates the first An approximation of the probability that a frame is used by at least one atom, which is determined by the frame-assigned soft probability matrix. It is obtained through approximate calculation of the product, and thus... Summing over all frames yields the differentiable frame loss. Its gradient can be propagated to the frame allocation matrix via backpropagation. .

[0113] And there are:

[0114] Thus, the gradient is:

[0115] This product is approximately in When a frame is differentiable everywhere, the expected value approaches the discrete frame count when a certain frame is not used, thus minimizing the loss by introducing a differentiable frame count. This can encourage atoms to concentrate in fewer frames.

[0116] Furthermore, in this embodiment of the application, when GPU mixed-precision training is enabled, the binary cross-entropy loss and multi-class cross-entropy loss, which include Sigmoid operations, can be calculated with float32 precision in the context of disabling automatic mixed precision on the current computing device, so as to avoid numerical overflow under float16 precision.

[0117] That is, in the training method for a graph neural network for atom array rearrangement according to the embodiments of this application, calculating the training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels includes: dividing the frame allocation vector corresponding to each atom in the frame allocation matrix by the frame allocation temperature and then taking the Softmax to obtain the frame allocation soft probability vector; calculating the cross-entropy of the frame allocation soft probability vector relative to the frame execution number in the training labels; and taking the average of the cross-entropy corresponding to all atoms to obtain the frame allocation loss function.

[0118] Furthermore, in the above-described training method for graph neural networks used for atomic array rearrangement, calculating the training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels includes: dividing the frame allocation matrix by the frame allocation temperature and then taking the softmax to obtain a frame allocation soft probability matrix; calculating an approximate probability value of each frame being used by at least one atom based on the frame allocation soft probability matrix; and summing the approximate probability values ​​of all frames to obtain a differentiable frame number loss function.

[0119] The training loss function It is the weighted sum of the above items, for example, expressed as:

[0120] Therefore, the training loss function Perform backpropagation to calculate gradients and update all trainable parameters of the graph neural network and the matching output head and frame allocation output head (preferably further including a compatible output head). Then, complete multiple rounds of iterative training by traversing all training samples in the training dataset until the model converges, and finally output the trained graph neural network and the matching output head and frame allocation output head (preferably further including a compatible output head).

[0121] That is, in the training method for a graph neural network for atomic array rearrangement according to the embodiments of this application, calculating the training loss function based at least on the matching score matrix, the frame allocation matrix, and the training labels includes: calculating a weighted sum of the matching loss function, the compatibility loss function, the path crossover penalty loss function, the frame allocation loss function, and the differentiable frame number loss function to obtain the training loss function.

[0122] In this embodiment, the training of the graph neural network can employ a multi-stage course scheduling method, where the number of stages is adjustable, for example, ranging from 1 to 10, with a default number of stages of 4. Furthermore, more loss function terms and annealing temperature parameters are gradually introduced into each stage. The following table shows a set of typical hyperparameters (all are empirical values, not limited to this combination, and...). For example, 1.0):

[0123] Those skilled in the art will understand that all values ​​in the table above are empirical reference values ​​and are not intended to limit this application. The training method for a graph neural network for atomic array rearrangement according to embodiments of this application can employ the multi-stage progressive introduction of loss function terms and annealing temperature parameters course learning framework described above, and is not limited to specific values. That is, the total number of training period intervals (e.g., approximately 50, default 200), the number of boundary period intervals in each stage, and... and The start and end values ​​and annealing strategy are all adjustable hyperparameters. Furthermore, in the embodiments of this application, the aforementioned weights... arrive It can be zero, indicating that the corresponding loss term is disabled.

[0124] For example, in the example shown in the table above, only the pairing loss is optimized in the early stages of training, and the frame allocation loss and differentiable frame number loss are gradually introduced in subsequent stages. The loss weight, sinkhorn temperature, frame allocation temperature, training array size and the number of rounds in each stage are all adjustable hyperparameters. At the same time, the optimal checkpoint within each stage is tracked independently when switching between stages.

[0125] Therefore, in the training method for a graph neural network for atomic array rearrangement according to embodiments of this application, updating the graph neural network and the matching output head and frame allocation output head via backpropagation of the training loss function includes: updating the graph neural network and the matching output head and frame allocation output head via backpropagation of the training loss function through a course learning framework that progressively introduces various loss functions and annealing temperature parameters in a multi-stage manner.

[0126] Furthermore, in the above-described training method for graph neural networks used for atomic array rearrangement, the annealing temperature parameters in the course learning framework include: the Sinkhorn temperature in the Sinkhorn operator during synchronous annealing and the frame allocation temperature.

[0127] Figure 3 A schematic flowchart illustrating an atomic array rearrangement method using a graph neural network according to an embodiment of this application is shown. Figure 3 As shown, the atomic array rearrangement method using a graph neural network according to an embodiment of this application includes the following steps.

[0128] Step S210: Construct a structured input graph based on the initial two-dimensional grid array of atoms to be rearranged. For example, taking an SLM (spatial light modulator) platform as an example, the initial two-dimensional grid array of atoms to be rearranged is a 20×20 two-dimensional grid array. With a loading probability of 0.6, the number of atoms is 240, and the target number can be set to 144, for example.

[0129] When constructing a structured input graph, each grid point of the two-dimensional grid array is first treated as a node in the graph, constructing a graph containing... The system generates a set of nodes and an initial node feature vector for each node, which includes normalized row coordinates, column normalized coordinates, atom occupancy markers, and target position markers. Simultaneously, undirected edges are established between nodes based on neighborhood connections (e.g., 4-neighborhood connections between upper, lower, left, and right adjacent grid points) to form an edge set. Finally, the node set, edge set, and initial node feature vector set are encapsulated into a structured input graph with a complete graph data structure.

[0130] Step S220: Input the structured input graph into the trained graph neural network to obtain a set of node feature vectors. Here, the graph neural network is trained using the graph neural network training method for atomic array rearrangement described above, for example, employing the GATv2 architecture, including 8 attention layers, i.e. Each layer includes 8 attention heads, namely Hidden Dimensions And parameters .

[0131] The structured input graph performs graph attention message passing at each attention layer. In each attention layer, the attention weights between each node's corresponding node feature vector and all its neighboring nodes are calculated in parallel through the attention head. The node feature vectors of the neighboring nodes are then weighted and aggregated to update the node feature vector of that node, ultimately obtaining the set of node feature vectors.

[0132] Step S230: Input the set of node feature vectors into the matching output head and the frame allocation output head respectively to obtain the matching score matrix and the frame allocation matrix. That is, the matching output head and the frame allocation output head are trained using the graph neural network training method for atomic array rearrangement as described above, and the matching score matrix and the frame allocation matrix are obtained from the set of node feature vectors.

[0133] That is, in the manner described above, the set of node feature vectors is input to at least the matching output header and the frame allocation output header to obtain the matching score matrix and the frame allocation matrix, and further, a compatibility output header can be input to obtain a compatibility prediction matrix.

[0134] Step S240: Perform matching decoding on the matching score matrix to obtain the optimal pairing, and perform argmax decoding on the frame allocation matrix to obtain the frame execution number.

[0135] That is, for the matching score matrix, for example denoted as Perform matching decoding, such as Sinkhorn decoding. Here, the inference temperature parameter of the Sinkhorn decoding is... The final annealing value from the final training stage can be used to obtain a double random matrix. , is represented as:

[0136] Then, for each row of the double random matrix, take argmax to obtain the hard-paired target number for each atom, represented as:

[0137] Besides Sinkhorn decoding, other differentiable / near-differentiable decoding methods such as Auction and Hungarian can be used to convert the score matrix into a double random matrix or hard-paired differentiable / near-differentiable decoding methods.

[0138] Furthermore, the frame allocation matrix is ​​decoded using argmax to obtain the frame execution number. That is, the frame execution number of each atom can be obtained by directly taking the argmax value of its frame allocation vector, as shown below:

[0139] Step S250: Based on the optimal pairing and the frame execution number, move the atom from its initial position to the target position. That is, traverse all atoms and group them by frame execution number. For each atom... ( ), according to its frame execution number and matching target number The instruction to move the atom (atom) Atom movement from the current grid point to the target grid point constitutes a corresponding batch. In this way, the atom movement instructions in each batch are sent to the control hardware (spatial light modulator (SLM) or acousto-optic deflector (AOD)). All atom movements within each batch are executed in parallel, and the batches are executed serially, thus completing the rearrangement of the atom array.

[0140] The atomic array rearrangement method using graph neural networks described above eliminates the greedy graph coloring (NP-Hard) by using frame allocation headers and argmax, thus removing post-processing bottlenecks and reducing inference complexity. Furthermore, since the inference latency of graph neural networks is approximately constant with respect to the number of atoms, it can support real-time control of large-scale arrays.

[0141] For example, under the array rules and graph neural network configuration described above, the end-to-end inference latency is approximately 2.5 microseconds (GPU), and the actual number of frames used is... The ratio relative to the optimal number of frames is less than 1.03, and the path conflict rate is less than 0.5%.

[0142] Furthermore, on a small 10×10 array with 60 atoms, 36 targets, 10,000 training samples, and 200 period intervals, the pairing accuracy was 80.61%, with the actual number of frames used... The ratio relative to the optimal number of frames is 1.131, and the path conflict rate is 0.

[0143] According to an embodiment of this application, an electronic device is further provided, the electronic device including a processor coupled to a memory, the processor being configured to execute a computer program stored in the memory, such that the electronic device performs the training method for a graph neural network for atomic array rearrangement and the atomic array rearrangement method using a graph neural network according to an embodiment of this application as described above.

[0144] According to an embodiment of this application, an electronic device is further provided, the electronic device including a processor and a memory, the processor being coupled to the memory, the processor being configured to execute a computer program stored in the memory, such that the electronic device performs the training method for a graph neural network for atomic array rearrangement and the atomic array rearrangement method using a graph neural network according to an embodiment of this application as described above.

[0145] Figure 4 A schematic block diagram of an electronic device according to an embodiment of this application is shown. Figure 4 As shown, an electronic device 300 according to an embodiment of this application includes a coupled processor 310 and a memory 320.

[0146] According to an embodiment of this application, a computer program product is further provided, the computer program product comprising: computer program code, which, when executed on a computer, causes the computer to perform the training method for a graph neural network for atomic array rearrangement and the atomic array rearrangement method using a graph neural network according to an embodiment of this application as described above.

[0147] According to an embodiment of this application, a computer-readable storage medium is further provided, the computer-readable medium storing program code, which, when run on a computer, causes the computer to execute the training method for a graph neural network for atomic array rearrangement and the atomic array rearrangement method using a graph neural network according to an embodiment of this application as described above.

[0148] The terms “component,” “module,” “system,” etc., used in this specification are used to refer to computer-related entities, hardware, firmware, combinations of hardware and software, software, or software in execution. For example, a component can be, but is not limited to, a process running on a processor, a processor, an object, an executable file, an execution thread, a program, and / or a computer. As illustrated, applications running on computing devices and computing devices can both be components. One or more components may reside in a process and / or an execution thread, and components may be located on a single computer and / or distributed among two or more computers. Furthermore, these components can be executed from various computer-readable media on which various data structures are stored. Components can communicate, for example, via local and / or remote processes based on signals having one or more data packets (e.g., data from two components interacting with another component between a local system, a distributed system, and / or a network, such as the Internet interacting with other systems via signals).

[0149] Those skilled in the art will recognize that the various illustrative logical blocks and steps described in conjunction with the embodiments disclosed herein can be implemented in electronic hardware, or a combination of computer software and electronic hardware. Whether these functions are implemented in hardware or software depends on the specific application and design constraints of the technical solution. Those skilled in the art can use different methods to implement the described functions for each specific application, but such implementations should not be considered beyond the scope of this application.

[0150] Those skilled in the art will understand that, for the sake of convenience and brevity, the specific working processes of the systems, devices, and units described above can be referred to the corresponding processes in the foregoing method embodiments, and will not be repeated here.

[0151] In the several embodiments provided in this application, it should be understood that the disclosed systems, apparatuses, and methods can be implemented in other ways. For example, the apparatus embodiments described above are merely illustrative; for instance, the division of units is only a logical functional division, and in actual implementation, there may be other division methods. For example, multiple units or components may be combined or integrated into another system, or some features may be ignored or not executed. Furthermore, the coupling or direct coupling or communication connection shown or discussed may be through some interfaces; the indirect coupling or communication connection between apparatuses or units may be electrical, mechanical, or other forms.

[0152] The units described as separate components may or may not be physically separate. The components shown as units may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the units can be selected to achieve the purpose of this embodiment according to actual needs.

[0153] In addition, the functional units in the various embodiments of this application can be integrated into one processing unit, or each unit can exist physically separately, or two or more units can be integrated into one unit.

[0154] If the aforementioned functions are implemented as software functional units and sold or used as independent products, they can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of this application, in essence, or the part that contributes to the prior art, or a portion of the technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods described in the various embodiments of this application. The aforementioned storage medium includes various media capable of storing program code, such as USB flash drives, portable hard drives, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0155] The above description is merely a specific embodiment of this application, but the scope of protection of this application is not limited thereto. Any variations or substitutions that can be easily conceived by those skilled in the art within the scope of the technology disclosed in this application should be included within the scope of protection of this application. Therefore, the scope of protection of this application should be determined by the scope of the claims.

Claims

1. A training method for a graph neural network for atomic array rearrangement, comprising: Obtain an initial two-dimensional grid array of atoms for training and training labels based on the rearrangement results of the initial two-dimensional grid array; The initial two-dimensional grid array of atoms used for training is input into the feature extraction backbone network of the graph neural network to extract features and obtain a set of node feature vectors. The set of node feature vectors is input into at least the matching output header and the frame allocation output header to obtain at least the matching score matrix and the frame allocation matrix; as well as The training loss function is calculated based at least on the matching score matrix, the frame allocation matrix, and the training labels, and the graph neural network, the matching output head, and the frame allocation output head are updated through backpropagation of the training loss function.

2. The training method for a graph neural network for atomic array rearrangement as described in claim 1, wherein, The training label is at least one of the following: optimal pairing, frame execution number, compatibility matrix, and optimal number of frames.

3. The training method for a graph neural network for atomic array rearrangement as described in claim 2, wherein, Obtaining the initial two-dimensional lattice array of atoms for training and the training labels based on the rearrangement of the initial two-dimensional lattice array includes: The optimal pairings of each atom in the rearrangement result of the initial two-dimensional lattice array are used to form an optimal pairing list as the pairing loss supervision label for the graph neural network; and The compatibility matrix, which determines whether all atoms in the rearrangement result based on the initial two-dimensional grid array are executed in the same frame, is used as the compatibility loss supervision label for the graph neural network.

4. The training method for a graph neural network for atomic array rearrangement as described in claim 3, wherein, Obtaining the initial two-dimensional lattice array of atoms for training and the training labels based on the rearrangement of the initial two-dimensional lattice array further includes: The frame execution number of each atom in the rearrangement result of the initial two-dimensional grid array is used to form a frame execution number list to serve as a supervision label for the frame allocation loss of the graph neural network.

5. The training method for a graph neural network for atomic array rearrangement as described in claim 1, wherein, The feature extraction backbone network of the graph neural network includes multiple graph attention layers, and each graph attention layer includes multiple attention heads.

6. The training method for a graph neural network for atomic array rearrangement as described in claim 1, wherein, The set of node feature vectors includes a subset of initial node feature vectors corresponding to the initial position of the atom and a subset of target node feature vectors corresponding to the target position of the atom.

7. The training method for a graph neural network for atomic array rearrangement as described in claim 6, wherein, The matching score matrix is ​​obtained by inputting the set of node feature vectors into the matching output head, including: For each atom-target pair, the initial node feature vector from the corresponding subset of initial node feature vectors and the target node feature vector from the corresponding subset of target node feature vectors are concatenated and then mapped to a scalar matching score through a fully connected network; and... The matching score matrix is ​​formed by assembling the scalar matching scores corresponding to all atom-target pairs.

8. The training method for a graph neural network for atomic array rearrangement as described in claim 6, wherein, The process of inputting the set of node feature vectors into at least the matching output header and the frame allocation output header to obtain at least the matching score matrix and the frame allocation matrix further includes: The set of node feature vectors is input into the compatible output head to obtain the compatible prediction matrix.

9. The training method for a graph neural network for atomic array rearrangement as described in claim 8, wherein, The process of inputting the set of node feature vectors into the compatibility output header to obtain the compatibility prediction matrix includes: The compatibility probability value is obtained by concatenating every two initial node feature vectors in the subset of initial node feature vectors, then mapping them through a fully connected network and activating them using the Sigmoid activation function; and... The compatibility prediction matrix is ​​formed by combining the compatibility probability values ​​corresponding to all initial node pairs.

10. The training method for a graph neural network for atomic array rearrangement as described in claim 1, wherein, The training loss function is calculated based on at least the matching score matrix, the frame assignment matrix, and the training labels, including: The matching score matrix is ​​input into the Sinkhorn operator to obtain a double random matrix; and... The negative log-likelihood of the double random matrix is ​​calculated using the optimal pairing in the training labels as the supervision target to obtain the matching loss function.

11. The training method for a graph neural network for atomic array rearrangement as described in claim 1, wherein, The training loss function is calculated based on at least the matching score matrix, the frame assignment matrix, and the training labels, including: Calculate the binary cross-entropy element-wise between the compatibility prediction matrix and the compatibility matrix in the training labels; and... The average of all matrix elements is taken to obtain the compatibility loss function.

12. The training method for a graph neural network for atomic array rearrangement as described in claim 1, wherein, The training loss function is calculated based on at least the matching score matrix, the frame assignment matrix, and the training labels, including: The mean of the matrix elements corresponding to the atomic pairs in the compatible prediction matrix that intersect in the path space is calculated to obtain the path intersection penalty loss function.

13. The training method for a graph neural network for atomic array rearrangement as described in claim 1, wherein, The training loss function is calculated based on at least the matching score matrix, the frame assignment matrix, and the training labels, including: The training loss function is obtained by calculating a weighted sum of the matching loss function, the compatibility loss function, the path crossover penalty loss function, the frame allocation loss function, and the differentiable frame number loss function.

14. The training method for a graph neural network for atomic array rearrangement as described in claim 1, wherein, Updating the graph neural network, the matching output header, and the frame allocation output header through backpropagation of the training loss function includes: The learning framework, which progressively introduces various loss functions and annealing temperature parameters in a multi-stage manner, updates the graph neural network, the matching output head, and the frame allocation output head via backpropagation of the training loss function.

15. The training method for a graph neural network for atomic array rearrangement as described in claim 14, wherein, The annealing temperature parameters in the course learning framework include: the Sinkhorn temperature and frame allocation temperature in the synchronous annealing Sinkhorn operator.

16. A method for rearranging atomic arrays using a graph neural network, comprising: A structured input graph is constructed based on an initial two-dimensional lattice array of atoms to be rearranged. The structured input graph is input into the feature extraction backbone network of the graph neural network trained by the training method for graph neural networks for atomic array rearrangement as described in any one of claims 1 to 15 to obtain a set of node feature vectors. The set of node feature vectors is respectively input into the matching output head and frame allocation output head of the training method for graph neural network for atomic array rearrangement as described in any one of claims 1 to 15 to obtain the matching score matrix and frame allocation matrix. The matching score matrix is ​​matched and decoded to obtain the optimal pairing, and the frame allocation matrix is ​​argmax decoded to obtain the frame execution number; as well as, Based on the optimal pairing and the frame execution number, the atom is moved from the initial position to the target position.

17. An electronic device including a processor coupled to a memory, the processor being configured to execute a computer program stored in the memory to cause the electronic device to perform the method as claimed in any one of claims 1 to 16.

18. A computer-readable storage medium, characterized in that, The computer-readable storage medium stores a computer program that, when executed, implements the method as described in any one of claims 1 to 16.

Images