Graph-driven self-attention compressible memory management method

Abstract

The invention discloses a graph-driven self-attention compressible memory management method, and belongs to the technical field of penetration testing. According to the graph-driven self-attention compressible memory management method, an attention graph is constructed, semantic dependence and attention flow directions among multiple Agent nodes are accurately described by utilizing edge weights, the relevance and the accessibility of information are improved, and in the aspects of screening and loading, the expandability of the memory is improved. According to the method, key memory fragments are screened on the basis of concerned graph edge weights and task related features, only the loaded information is loaded, redundancy is effectively eliminated, the data volume processed by a system is reduced, and the system operation efficiency is improved; the semantic compression module can perform abstract processing on historical information based on various elements, and on the premise of ensuring that key semantic information is not lost, the abstract length is controlled within the window limit, so that the model can normally process information, and the adaptability and stability of the system to different data volumes are improved.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of penetration testing, and particularly relates to a graph-driven self-attention compressible memory management method. BACKGROUND

[0002] In a penetration testing task, multi-agent collaborative work needs to process a large amount of historical interaction information, including dialogue records, tool invocation trajectories, etc. The historical memory fragments formed by these information have important reference value for the execution of the current task.

[0003] However, the traditional memory management method has the following problems: 1. When constructing the memory structure, the complex semantic dependency and attention flow direction between multi-agent nodes cannot be effectively described, resulting in poor information correlation and difficulty in accurately retrieving key information.

[0004] 2. When facing a large number of memory fragments, there is a lack of efficient screening and loading mechanism, which will introduce a large amount of irrelevant redundant information into the processing flow, not only increasing the system burden, but also possibly interfering with the execution of the current task and reducing the efficiency.

[0005] 3. When the loaded memory content exceeds the context window limit of the large model input, the existing processing method often adopts the way of directly truncating the content, which will cause information loss or a significant decline in model performance, and cannot effectively compress the information under the premise of ensuring the integrity of the information. SUMMARY

[0006] To solve the problems raised in the background art, the application provides a graph-driven self-attention compressible memory management method to solve the problems of poor information correlation, difficulty in accurately retrieving key information, lack of efficient screening and loading mechanism, and simple compression causing information loss or significant decline in model performance in traditional penetration testing tasks.

[0007] To achieve the above purpose, the application provides the following technical solution: A graph-driven self-attention compressible memory management method, comprising the following steps: S1: Construct a focus graph; each round of message of multi-agent interaction in a penetration testing task is taken as a node to construct a dynamic directed graph structure focus graph, wherein the message includes dialogue records and tool invocation data, the edge weight of the focus graph represents the attention dependency relationship between different nodes, the edge weight value is determined by analyzing the semantic correlation degree and interaction logic flow direction between nodes, and is used to describe the semantic dependency and attention flow direction between multi-agent nodes, and the message corresponding to the node constitutes a historical memory fragment; S2: Directional retrieval and selective loading; based on the edge weight of the attention graph, the historical memory fragments that are strongly related to the current penetration testing task, i.e., key memory fragments, are screened out from the attention graph according to the matching of the vulnerability type, attack stage and tool features of the current task, the key memory fragments are loaded, irrelevant and redundant information is removed, if the total length of the loaded memory fragment content exceeds the preset message length threshold of the large model, S3 is executed, if it does not exceed, S4 is directly executed; S3: Context window overflow adaptive processing; enable semantic compression module, based on the core elements including task intent, tool call and historical response content, abstract the key memory fragments loaded in S2, extract key semantic information and generate compressed abstract; S4: Multi-scene adaptive optimization; based on the memory content loaded in S2 or the compressed abstract generated in S3, multi-parameter tool call optimization and multi-Agent collaborative optimization are selected, and the multi-Agent collaborative optimization specifically delivers the related information of Agents responsible for different tasks through the attention graph to ensure that each Agent shares the necessary context and works based on the shared information.

[0008] Preferably, in S1, when a new interaction message is added to the attention graph as a node, the edge weight is updated by calculating the semantic similarity between the new node and the existing nodes, the semantic similarity calculation uses a vector cosine similarity algorithm based on a pre-trained language model, and the edge weight is dynamically adjusted according to the new semantic association, and the adjustment amplitude is positively correlated with the semantic similarity.

[0009] Preferably, in S2, the tool features include tool type, tool function and tool call history parameters, the tool type includes vulnerability scanning tools, vulnerability exploitation tools and privilege escalation tools, the tool function is divided into information collection, vulnerability detection and attack implementation categories according to its role in the penetration testing process, and the tool call history parameters include parameter name, parameter value range and parameter combination mode.

[0010] Preferably, in S3, the semantic compression module performs semantic analysis and content refinement on historical information through a natural language processing model.

[0011] Preferably, in the multi-parameter tool call optimization of S4, tool parameter descriptions are added to the memory content loaded in S2 or the compressed abstract generated in S3, thereby enhancing the model's understanding of tool usage.

[0012] Preferably, in the multi-Agent collaborative optimization of S4, the related information includes task progress, tool call results and permission status, the task progress presents the current task completion and the promotion status of each sub-task in percentage form, the tool call result contains whether the call is successful, the returned data and error code, and the permission status records the current system permission level and permission range.

[0013] Preferably, in S1, the node is uniquely identified, and the identification information includes the timestamp of message generation, the number of the sending Agent, and the type of the message, the number of the sending Agent is a unique identifier assigned by the system, and the attribute information of the node further includes the message type, the message length, and the associated tool identification.

[0014] Preferably, in S3, the generated compressed summary is accompanied by index information, and the index information includes the unique identification of the original node, the storage path, and the generation time.

[0015] Preferably, in S4, when multiple Agents are cooperatively optimized, an information synchronization mechanism between Agents is established, and the unified processing of message distribution and acceptance is realized through the subscription-published mode.

[0016] Compared with the prior art, the beneficial effects of the present application are: The present application can better organize and understand historical memory fragments by constructing an attention graph and accurately depicting semantic dependence and attention flow between multiple Agent nodes through edge weight values, thereby providing a good data foundation for subsequent retrieval and application, improving the relevance and accessibility of information, and improving the accuracy of task execution. In the aspect of screening and loading, the present application screens key memory fragments based on edge weight values of the attention graph and task-related features, only loads useful information, effectively eliminates redundancy, reduces the amount of data processed by the system, improves the efficiency of system operation, avoids the interference of irrelevant information on the current task, and improves the accuracy of task execution. When the memory content exceeds the context window limit of the large model input, the semantic compression module can perform summary processing on historical information based on multiple elements, control the summary length within the window limit on the premise of ensuring that key semantic information is not lost, ensure that the model can normally process information, and improve the adaptability and stability of the system to different data amounts. BRIEF DESCRIPTION OF DRAWINGS

[0017] Figure 1 The flowchart of the present application is shown. DETAILED DESCRIPTION

[0018] To facilitate those skilled in the art to understand the technical content of the present application, the present application is further described in detail below in combination with the drawings and specific examples. It should be understood that the specific examples described herein are only used to explain the present application and do not limit the present application.

[0019] Example 1 A graph-driven self-attention compressible memory management method, as shown in Figure 1 includes the following steps: S1: constructing an attention graph; constructing an attention graph with a dynamic directed graph structure for each round of message interaction between multiple agents in a penetration testing task, wherein the message includes dialogue records and tool call data, the edge weight of the attention graph represents the attention dependence relationship between different nodes, the edge weight value is determined by analyzing the semantic correlation and interaction logic flow between nodes, and is used to depict the semantic dependence and attention flow direction between multiple agent nodes, and the node and its corresponding message together constitute a historical memory fragment; S2: directional retrieval and selective loading; based on the edge weight value of the attention graph, filtering out the historical memory fragments, i.e. key memory fragments, that are strongly related to the current penetration testing task from the attention graph, the strong correlation is matched according to the vulnerability type, attack stage and tool used of the current task, the key memory fragments are loaded, and irrelevant redundant information is removed, if the total length of the loaded memory fragment content exceeds the input context window limit of the large model, S3 is executed, if it does not exceed, S4 is directly executed; S3: context window overflow adaptive processing; enabling a semantic compression module, based on the core elements including task intent, tool call and historical response content, abstracting the key memory fragments loaded in S2, extracting key semantic information and generating a compressed abstract; S4: multi-scene adaptive optimization; based on the memory content loaded in S2 or the compressed abstract generated in S3, selecting to perform multi-parameter tool call optimization and multi-agent collaborative optimization, wherein the multi-parameter tool call optimization specifically reuses historical parameter call records, and supports parameter recommendation and adaptive combination based on the current test environment, and the multi-agent collaborative optimization specifically transfers the related information of agents responsible for different tasks through the attention graph, ensures that each agent shares necessary context and works based on shared information.

[0020] In this embodiment, the application can better organize and understand historical memory fragments by constructing an attention graph, accurately depicting the semantic dependence and attention flow direction between multiple agent nodes through edge weight values, providing a good data foundation for subsequent retrieval and application, improving the relevance and accessibility of information, in the aspect of screening and loading, the application screens key memory fragments based on the edge weight value of the attention graph and the task related features, only loads useful information, effectively eliminates redundancy, reduces the amount of data processed by the system, improves the efficiency of the system, and avoids the interference of irrelevant information on the current task, improves the accuracy of task execution, when the memory content exceeds the input context window limit of the large model, the semantic compression module can abstract historical information based on multiple elements, control the abstract length within the window limit on the premise of not losing key semantic information, ensure that the model can normally process information, and improve the adaptability and stability of the system to different data quantities.

[0021] Embodiment 2 The difference between this embodiment and embodiment 1 is that in S1, when a new interaction message is added as a node to the attention graph, the edge weight is updated by calculating the semantic similarity between the new node and the existing nodes. The semantic similarity calculation uses a vector cosine similarity algorithm based on a pre-trained language model, and the edge weight is dynamically adjusted according to the new semantic association. The adjustment amplitude is positively correlated with the semantic similarity.

[0022] In this embodiment, the pre-trained language model is BERT-base model, the node message content is converted into a multi-dimensional vector representation, the vector cosine similarity calculation is realized through matrix operation, the result is accurate to four decimal places, the initial value of the edge weight is set to 0.5, when a new node is added, the semantic similarity with all existing nodes is calculated, and the edge weight of the new node and all other nodes is updated, the update formula is: New edge weight = old edge weight + (semantic similarity - 0.5) x 0.2; Ensure that the adjustment amplitude is positively correlated with the semantic similarity and the edge weight is always within the range of 0-1.

[0023] Embodiment 3 The difference between this embodiment and embodiment 1 is that in S2, the tool features include tool type, tool function and tool call history parameters, where the tool type includes vulnerability scanning tools, vulnerability exploitation tools, and privilege escalation tools. The tool function is divided into information collection, vulnerability detection, and attack implementation categories according to its role in the penetration testing process. The tool call history parameters include parameter name, parameter value range and parameter combination mode.

[0024] Embodiment 4 The difference between this embodiment and embodiment 1 is that in S3, the semantic compression module performs semantic analysis on historical information through a natural language processing model.

[0025] In this embodiment, the natural language processing model uses the spaCy framework, entity extraction includes host name, IP address, vulnerability number, etc. Relationship extraction covers "vulnerability exists" and "tool call" associations. Key action extraction includes "scan" and "attack" behavior verbs. When filtering core elements based on task intent, a combination of keyword matching and semantic relevance sorting is used. The core execution logic is retained in the form of a step-by-step list. The length of the compressed summary is controlled by token count. The preset semantic similarity threshold is set to 0.8, and the ROUGE-L algorithm is used to calculate the similarity.

[0026] Embodiment 5 The difference between this embodiment and embodiment 1 is that in the multi-parameter tool call optimization of S4, the recommendation of tool parameter description is based on the representation vector of the current penetration context information, the cosine similarity is calculated with other message nodes, the historical messages with high similarity are pulled, and the neighbor messages with high similarity are recursively pulled. Among them, the penetration context information includes the environmental information of the target system, the execution information of the last tool instruction, the penetration target information output by the Agent at the current penetration time step, and the global final target information. The recursive depth is determined by the configuration file. This method dynamically pulls different historical context information according to the current penetration environment to assist the Agent in decision-making.

[0027] Embodiment 6 The difference between this embodiment and embodiment 1 is that in the multi-Agent cooperative optimization of S4, the relevant information includes task progress, tool call result and permission state. The task progress presents the current task completion and the promotion state of each subtask in percentage form. The tool call result includes whether the call is successful, the return data and error code. The permission state records the current system permission level and permission range.

[0028] Embodiment 7 The difference between this embodiment and embodiment 1 is that in S1, the node is uniquely identified, and the identification information includes the timestamp of message generation, the number of sending Agent and the type of message. For example, tool expert can output tool instruction message, tool result message and Base reply message of tool expert. The number of sending Agent is a unique identifier assigned by the system. The attribute information of the node also includes message type, message length and associated tool identification.

[0029] Embodiment 8 The difference between this embodiment and embodiment 7 is that in S3, the generated compressed summary is attached with index information, which includes the unique identification of the original node, the storage path and the generation time.

[0030] Embodiment 9 The difference between this embodiment and embodiment 1 is that in S4, when multiple Agents are cooperatively optimized, an information synchronization mechanism between Agents is established, information transmission is realized through the subscription-published mode, and Agents subscribe to related information topics according to their own responsibilities and automatically receive information updates.

[0031] In this embodiment, the subscription-published mode is realized based on the MQTT protocol, and the information topics are divided according to the hierarchical structure of “task type / information category”, such as “penetration test / permission state”. The correspondence between the responsibilities of Agents and the information topics is set through the configuration file, and dynamic modification is supported.

[0032] The above merely describes the preferred embodiments of the present application, but not for limiting the protection scope of the present application. Any modification, equivalent replacement, improvement, etc. made within the spirit and principle of the present application shall be included in the protection scope of the present application.

Claims

1. A graph-driven self-attention compressible memory management method, characterized in that, Includes the following steps: S1: Construct an attention graph; Using each round of messages from multi-agent interactions in a penetration testing task as nodes, construct an attention graph with a dynamic directed graph structure. The messages include dialogue records and tool call data. The edge weights of the attention graph represent the attention dependencies between different nodes. The edge weights are determined by analyzing the semantic correlation and interaction logic flow between nodes, and are used to characterize the semantic dependencies and attention flow between multi-agent nodes. The messages corresponding to the nodes constitute historical memory fragments. S2: Targeted retrieval and selective loading; Based on the edge weights of the attention graph, historical memory fragments strongly related to the current penetration testing task are selected from the attention graph, namely key memory fragments. The strong correlation is matched according to the vulnerability type, attack stage and tool characteristics of the current task. The key memory fragments are loaded and irrelevant redundant information is removed. If the total length of the loaded memory fragment content exceeds the preset message length threshold of the large model, then S3 is executed. If it does not exceed the threshold, then S4 is executed directly. S3: Adaptive handling of context window overflow; Enable semantic compression module, based on core elements including task intent, tool calls and historical response content, to perform summary processing on key memory fragments loaded by S2, extract key semantic information and generate compressed summaries; S4: Multi-scenario adaptation optimization; Based on the memory content loaded in S2 or the compressed summary generated in S3, select to perform multi-parameter tool call optimization and multi-Agent collaborative optimization. Specifically, multi-Agent collaborative optimization involves passing relevant information of agents responsible for different tasks through an attention graph, ensuring that each agent shares the necessary context and performs its work based on the shared information.

2. The graph-driven self-attention compressible memory management method according to claim 1, characterized in that, In S1, when a new interactive message is added to the attention graph as a node, the edge weights are updated by calculating the semantic similarity between the new node and the existing nodes. The semantic similarity calculation adopts the vector cosine similarity algorithm based on the pre-trained language model. The edge weights are dynamically adjusted according to the new semantic association, and the adjustment magnitude is positively correlated with the semantic similarity.

3. The graph-driven self-attention compressible memory management method according to claim 1, characterized in that, In S2, tool characteristics include tool type, tool function, and tool call history parameters. Tool type includes vulnerability scanning tools, vulnerability exploitation tools, and privilege escalation tools. Tool function is divided into information gathering, vulnerability detection, and attack implementation categories according to its role in the penetration testing process. Tool call history parameters include parameter name, parameter value range, and parameter combination mode.

4. The graph-driven self-attention compressible memory management method according to claim 1, characterized in that, In S3, the semantic compression module performs semantic parsing and content refinement of historical information using a natural language processing model.

5. The graph-driven self-attention compressible memory management method according to claim 1, characterized in that, In the multi-parameter tool call optimization of S4, tool parameter descriptions are added to the memory content loaded in S2 or the compressed summary generated in S3, thereby enhancing the model's understanding of tool usage.

6. The graph-driven self-attention compressible memory management method according to claim 1, characterized in that, In S4's multi-agent collaborative optimization, relevant information includes task progress, tool call results, and permission status. Task progress is presented as a percentage to show the current task completion status and the progress status of each subtask. Tool call results include whether the call was successful, the returned data, and the error code. Permission status records the currently acquired system permission level and permission scope.

7. The graph-driven self-attention compressible memory management method according to claim 1, characterized in that, In S1, each node is uniquely identified. The identification information includes the timestamp of the message generation, the ID of the sending Agent, and the message type. The ID of the sending Agent is a unique identifier assigned by the system. The node's attribute information also includes the message type, message length, and associated tool identifier.

8. The graph-driven self-attention compressible memory management method according to claim 7, characterized in that, In S3, the generated compressed digest includes index information, which includes the unique identifier of the original node, the storage path, and the generation time.

9. The graph-driven self-attention compressible memory management method according to claim 1, characterized in that, In S4, when multiple agents collaborate for optimization, an information synchronization mechanism is established between agents, and message distribution and reception are uniformly processed through a publish-subscribe pattern.

Images