IP Library › Granted Patent US 12,614,110
Granted Patent B1
US 12,614,110 · App. 17/954,761 · Granted Apr 28, 2026

Self-attention masks for training models

Inventors: Hongbin Zheng (San Jose, CA); Yunxuan Yu (San Jose, CA); Ron Diamant (Santa Clara, CA)
Assignee: Amazon Technologies, Inc.
G06N20/00G06F18/29
View Patent ↗
Loading inventors, assignments & file history…
Monitor This Case
Get email alerts when status or documents change.
Order Certified Copies
Most orders are placed with the USPTO same day — all within 24 business hours.
Order via The Patent Place →
Pre-filled with this patent's details
Quick Facts
Patent No.
US 12,614,110
App. No.
17/954,761
Granted
Apr 28, 2026
Kind
B1
Abstract

A computer-implemented technique for optimizing self-attention masks is described. At compile time, a machine learning graph of an artificial intelligence model is analyzed. The machine learning graph includes a set of operators. Analysis includes identifying one or more mask operators and determining what fields of input tensors are masked. Optimizations at compile-time are used to eliminate instructions during training of the artificial intelligence model.

Claims (70)

1 . A computer-implemented method, comprising:

generating, at compile time, instructions configured for training a machine learning (ML) model on a training input tensor;

receiving a machine learning graph of the AI model, the machine learning graph including a set of operators, wherein each operator of the set of operators receives an input tensor and outputs an output tensor, the set of operators including a first operator and a second operator, the first operator receives a first input tensor and outputs a first output tensor, the first input tensor derived from the training input tensor, the second operator receives a second input tensor, wherein the first output tensor is the second input tensor, wherein each input tensor and each output tensor are comprised of fields holding values;

identifying that the first operator is a mask operator configured to apply a mask to the first input tensor, the mask operator associated with a mask subtensor and a mask value, wherein the mask subtensor includes the fields of the first input tensor that will be set to the mask value;

generating a mapping associating each input tensor of each of the set of operators to a compute range, a mask range, and a compute subtensor, wherein the compute subtensor includes the fields of the first input tensor that are not included by the mask subtensor;

determining the compute range of the first input tensor;

determining the mask range of the first input tensor based on the mask value;

determining the compute subtensor of the first input tensor based on an affine expression associated with the mask operator;

determining the compute range of the second input tensor based on range analysis of the compute range of the first input tensor;

determining the mask range of the second input tensor based on the mask range of the first operator and the compute range of the second input tensor;

determining the compute subtensor of the second input tensor based on the compute subtensor of the first operator;

generating a first set of predicates, associated with the second operator, for computations on the fields of the mask subtensor based on the mask range of the second input tensor and the compute subtensor of the second input tensor;

determining a subset of the instructions to eliminate based on the first set of predicates;

eliminating the subset of the instructions.

2 . The computer-implemented method of claim 1 , further comprising:

generating a second set of predicates, associated with the second operator, for computing the mask value for each field in the mask subtensor based on the mask range of the second input tensor;

wherein determining the subset of the instructions to eliminate is further based on the second set of predicates.

3 . The computer-implemented method of claim 1 , further comprising:

generating a second set of predicates, associated with the first operator, for computations on the fields of the mask subtensor based on the mask range of the first input tensor and the compute subtensor of the first input tensor;

wherein determining the subset of the instructions to eliminate is further based on the second set of predicates.

4 . The computer-implemented method of claim 1 , further comprising:

generating a second set of predicates, associated with the first operator, for computing the mask value for each field in the mask subtensor based on the mask range of the first input tensor;

wherein determining the subset of the instructions to eliminate is further based on the second set of predicates.

5 . The computer-implemented method of claim 1 , wherein the second operator is a non-mask operator.

6 . A computer-implemented method, comprising:

generating, at compile time, instructions configured for training an artificial intelligence (AI) model on a training input tensor;

receiving a machine learning graph of the AI model, the machine learning graph including a set of operators including a first operator, wherein the first operator receives a first input tensor that is comprised of fields holding values;

identifying that the first operator is a mask operator configured to apply a mask to the first input tensor, wherein the mask operator is associated with a mask subtensor and a mask value, and wherein the mask subtensor includes the fields of the first input tensor that will be set to the mask value;

determining a mask range of the first input tensor based on the mask value;

determining a compute subtensor of the first input tensor based on the mask operator;

generating a first set of predicates, associated with the first operator, for computations on the fields of the mask subtensor based on the mask range of the first input tensor and the compute subtensor of the first input tensor;

determining a subset of the instructions to eliminate based on the first set of predicates; and

eliminating the subset of the instructions.

7 . The computer-implemented method of claim 6 , further comprising:

generating a second set of predicates, associated with the first operator, for computing the mask value for each field in the mask subtensor based on the mask range of the first input tensor;

wherein determining the subset of the instructions to eliminate is further based on the second set of predicates.

8 . The computer-implemented method of claim 7 , further comprising generating memset instructions for each field in the mask subtensor based on the second set of predicates, wherein the memset instructions set the fields in the mask subtensor to the mask value.

9 . The computer-implemented method of claim 6 , wherein the set of operators includes a second operator, wherein the second operator receives a second input tensor, wherein an output tensor of the first operator is the second input tensor.

10 . The computer-implemented method of claim 9 , wherein the first operator receives a third input tensor.

11 . The computer-implemented method of claim 10 , further comprising:

determining a compute range of the second input tensor based on the compute range of the first input tensor and a compute range of the third input tensor;

determining the mask range of the second input tensor based on one or more of: the mask range of the first input tensor, the mask range of the third input tensor, the compute range of the first input tensor, or the compute range of the third input tensor; and

determining the compute subtensor of the second input tensor based on the compute subtensor of the first input tensor and the compute subtensor of the third input tensor.

12 . The computer-implemented method of claim 11 , further comprising:

generating a second set of predicates, associated with the second operator, for computations on the fields of the mask subtensor based on the mask range of the second input tensor and the compute subtensor of the second input tensor;

wherein, determining the subset of the instructions to eliminate is further based on the second set of predicates.

13 . The computer-implemented method of claim 11 , further comprising:

generating a second set of predicates, associated with the second operator, for computing the mask value for each field in the mask subtensor based on the mask range of the second input tensor;

wherein, determining the subset of the instructions to eliminate is further based on the second set of predicates.

14 . A computer-program product tangibly embodied in a non-transitory machine-readable storage medium, the computer-program product including instructions configured to cause a processor to perform operations comprising, comprising:

generating, at compile time, instructions configured for training an artificial intelligence (AI) model on a training input tensor;

receiving a machine learning graph of the AI model, the machine learning graph including a set of operators including a first operator, wherein the first operator receives a first input tensor that is comprised of fields holding values;

identifying that the first operator is a mask operator configured to apply a mask to the first input tensor, wherein the mask operator is associated with a mask subtensor and a mask value, and wherein the mask subtensor includes the fields of the first input tensor that will be set to the mask value;

determining a mask range of the first input tensor based on the mask value;

determining a compute subtensor of the first input tensor based on the mask operator;

generating a first set of predicates, associated with the first operator, for computations on the fields of the mask subtensor based on the mask range of the first input tensor and the compute subtensor of the first input tensor;

determining a subset of the instructions to eliminate based on the first set of predicates; and

eliminating the subset of the instructions.

15 . The computer-program product of claim 14 , further comprising:

generating a second set of predicates, associated with the first operator, for computing the mask value for each field in the mask subtensor based on the mask range of the first input tensor;

wherein determining the subset of the instructions to eliminate is further based on the second set of predicates.

16 . The computer-program product of claim 15 , further comprising generating memset instructions for each field in the mask subtensor based on the second set of predicates, wherein the memset instructions set the fields in the mask subtensor to the mask value.

17 . The computer-program product of claim 14 , wherein the set of operators includes a second operator, wherein the second operator receives a second input tensor, wherein an output tensor of the first operator is the second input tensor, wherein the first operator receives a third input tensor.

18 . The computer-program product of claim 17 , further comprising:

determining a compute range of the second input tensor based on the compute range of the first input tensor and a compute range of the third input tensor;

determining the mask range of the second input tensor based on one or more of: the mask range of the first input tensor, the mask range of the third input tensor, the compute range of the first input tensor, or the compute range of the third input tensor; and

determining the compute subtensor of the second input tensor based on the compute subtensor of the first input tensor and the compute subtensor of the third input tensor.

19 . The computer-program product of claim 18 , further comprising:

generating a third set of predicates, associated with the second operator, for computations on the fields of the mask subtensor based on the mask range of the second input tensor and the compute subtensor of the second input tensor;

wherein, determining the subset of the instructions to eliminate is further based on the third set of predicates.

Assignments (1)
ASSIGNMENT OF ASSIGNOR'S INTEREST Recorded Sep 28, 2022
From: ZHENG, HONGBIN; YU, YUNXUAN; DIAMANT, RON
To: AMAZON TECHNOLOGIES, INC.
Reel/Frame 061242/0893 →
References Cited (9)
US 10909401B2 · Burachas · 2021 [cited by examiner]
US 11468365B2 · Olgiati · 2022 [cited by examiner]
US 20210158206A1 · Klein · 2021 [cited by examiner]
US 20240289960A1 · Zhong · 2024 [cited by examiner]
Fchollet, “Training & Evaluation with The Built-in Methods”, Jul. 24, 2023 UTC. https://www.tensorflow.org/guide/keras/training_with_built_in_methods (Year: 2023). [cited by examiner]
Zhu et al. “Understanding Masking & Padding”, Jul. 24, 2023 UTC. https://www.tensorflow.org/guide/keras/understanding_masking_and_padding (Year: 2023). [cited by examiner]
He et al. “Query Processing on Tensor Computation Runtimes”, Aug. 1, 2022, https://arxiv.org/pdf/2203.01877v3 (Year: 2022). [cited by examiner]
Huntersperger et al. “Introduction to graphs and tf.function”, Aug. 15, 2024 UTC. https://www.tensorflow.org/guide/intro_to_graphs (Year: 2024). [cited by examiner]
Badreddine et al. “Logic Tensor Networks”, Dec. 23, 2021, https://arxiv.org/pdf/2012.13635 (Year: 2021). [cited by examiner]