IP Library › Granted Patent US 12,423,518
Granted Patent B2
US 12,423,518 · App. 17/903,805 · Granted Sep 23, 2025

Attention neural networks with N-grammer layers

Inventors: Rohan Anil (San Francisco, CA); Aurko Roy (San Francisco, CA)
Assignee: Google LLC
G06F40/20G06N3/04
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,423,518
App. No.
17/903,805
Granted
Sep 23, 2025
Kind
B2
Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for performing a machine learning task on a network input to generate a network output. In one aspect, one of the systems includes a neural network configured to perform the machine learning task, the neural network comprising an N-grammer layer and an output neural network, the N-grammer layer configured to: at each of one or more heads: receive a sequence of input embeddings; generate a discrete latent representation of the sequence of input embeddings by using a learned product quantization codebook; generate a plurality of n-gram indices from the discrete latent representation; and generate a latent n-gram representation of the sequence of input embeddings; and generate a sequence of output embeddings, and the output neural network configured to: receive the sequence of output embeddings; and process the sequence of output embeddings to generate the network output.

Claims (63)

1. A system for performing a machine learning task on a network input to generate a network output, the system comprising one or more computers and one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to implement:

a neural network configured to perform the machine learning task, the neural network comprising an N-grammer layer and an output neural network, the N-grammer layer configured to:

at each of one or more heads:

receive a sequence of input embeddings derived from the network input that includes a respective input embedding at each of a plurality of input embedding positions;

generate a discrete latent representation of the sequence of input embeddings by using a learned product quantization codebook having different cluster indices that each map to a respective code word, wherein the discrete latent representation of the sequence of input embeddings comprises, at each of the plurality of input embedding positions, a corresponding cluster index that maps to a closest code word to the input embedding at the input embedding position;

generate a plurality of n-gram indices from the discrete latent representation, comprising assigning to each of the plurality of input embedding positions a respective n-gram index based on the corresponding cluster index at the input embedding position and respective cluster indices at one or more preceding input embedding positions; and

generate a latent n-gram representation of the sequence of input embeddings that includes a respective n-gram embedding at each of the plurality of input embedding positions by using an n-gram embedding hash table and the plurality of n-gram indices; and

generate a sequence of output embeddings by determining a combination of the sequence of input embeddings and the latent n-gram representation of the sequence of input embeddings for each head, and

the output neural network configured to:

receive the sequence of output embeddings; and

process the sequence of output embeddings to generate the network output for the machine learning task.

2. The system of claim 1 , wherein the input embeddings comprise uni-gram embeddings.

3. The system of claim 1 , wherein the latent n-gram representation comprises a latent bi-gram representation.

4. The system of claim 1 , wherein the latent n-gram representation comprises a latent tri-gram representation or another latent higher order n-gram representation.

5. The system of claim 1 , wherein generating the sequence of output embeddings comprises:

applying respective layer normalizations to the sequence of input embeddings and to the latent n-gram representation of the sequence of input embeddings for each head; and

determining a concatenation of the layer normalized sequence of input embeddings and the layer normalized latent n-gram representation of the sequence of input embeddings for each head.

6. The system of claim 1 , wherein the learned product quantization codebook has k cluster indices that each map to a respective code word, and wherein generating the discrete latent representation of the sequence of input embeddings for each head comprises, at each of the plurality of input embedding positions:

selecting a code word from among all code words in the learned product quantization codebook that has a least distance from the input embedding at the input embedding position; and

assigning the input embedding to the cluster index that maps to the selected code word.

7. The system of claim 1 , wherein a total number of the plurality of n-gram indices is k n , and wherein assigning to each of the plurality of input embedding positions the respective n-gram index comprises:

determining a combination of the corresponding cluster index at the input embedding position and a respective cluster index at each of n preceding input embedding positions multiplied by k.

8. The system of claim 1 , wherein generating the latent n-gram representation of the sequence of input embeddings comprises generating a hash value for each of the plurality of input embedding positions by using universal hashing techniques.

9. The system of claim 8 , wherein generating the latent n-gram representation of the sequence of input embeddings at each head comprises, at each of the plurality of input embedding positions:

evaluating a hash function in form of ((rb+s) mod p) mod v, here p is a randomly selected prime number greater than k 2 , r is randomly selected from{1, . . . , p−1}, s is randomly selected from {0, . . . , p−2}, v is a size of a vocabulary of the n-gram embedding hash table, and b is the n-gram index at the input embedding position.

10. The system of claim 1 , wherein the n-gram embedding has a larger dimension than the input embedding.

11. The system of claim 1 , wherein the output neural network comprises one or more attention neural network layers that are each configured to:

apply an attention mechanism to an attention layer input derived from the sequence of output embeddings to generate an attention layer output for the attention neural network layer.

12. One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to implement:

a neural network configured to perform a machine learning task on a network input to generate a network output, the neural network comprising an N-grammer layer and an output neural network, the N-grammer layer configured to:

at each of one or more heads:

receive a sequence of input embeddings derived from the network input that includes a respective input embedding at each of a plurality of input embedding positions;

generate a discrete latent representation of the sequence of input embeddings by using a learned product quantization codebook having different cluster indices that each map to a respective code word, wherein the discrete latent representation of the sequence of input embeddings comprises, at each of the plurality of input embedding positions, a corresponding cluster index that maps to a closest code word to the input embedding at the input embedding position;

generate a plurality of n-gram indices from the discrete latent representation, comprising assigning to each of the plurality of input embedding positions a respective n-gram index based on the corresponding cluster index at the input embedding position and respective cluster indices at one or more preceding input embedding positions; and

generate a latent n-gram representation of the sequence of input embeddings that includes a respective n-gram embedding at each of the plurality of input embedding positions by using an n-gram embedding hash table and the plurality of n-gram indices; and

generate a sequence of output embeddings by determining a combination of the sequence of input embeddings and the latent n-gram representation of the sequence of input embeddings for each head, and

the output neural network configured to:

receive the sequence of output embeddings; and

process the sequence of output embeddings to generate the network output for the machine learning task.

13. A method performed by one or more computers, the method comprising:

receiving a network input; and

processing the network input using a neural network configured to perform a machine learning task on the network input to generate a network output, the neural network comprising an N-grammer layer and an output neural network, the N-grammer layer configured to,

at each of one or more heads:

receive a sequence of input embeddings derived from the network input that includes a respective input embedding at each of a plurality of input embedding positions;

generate a discrete latent representation of the sequence of input embeddings by using a learned product quantization codebook having different cluster indices that each map to a respective code word, wherein the discrete latent representation of the sequence of input embeddings comprises, at each of the plurality of input embedding positions, a corresponding cluster index that maps to a closest code word to the input embedding at the input embedding position;

generate a plurality of n-gram indices from the discrete latent representation, comprising assigning to each of the plurality of input embedding positions a respective n-gram index based on the corresponding cluster index at the input embedding position and respective cluster indices at one or more preceding input embedding positions; and

generate a latent n-gram representation of the sequence of input embeddings that includes a respective n-gram embedding at each of the plurality of input embedding positions by using an n-gram embedding hash table and the plurality of n-gram indices; and

generate a sequence of output embeddings by determining a combination of the sequence of input embeddings and the latent n-gram representation of the sequence of input embeddings for each head, and

the output neural network configured to:

receive the sequence of output embeddings; and

process the sequence of output embeddings to generate the network output for the machine learning task.

14. The method of claim 13 , wherein the input embeddings comprise uni-gram embeddings.

15. The method of claim 13 , wherein the latent n-gram representation comprises a latent bi-gram representation.

16. The method of claim 13 , wherein the latent n-gram representation comprises a latent tri-gram representation or another latent higher order n-gram representation.

17. The method of claim 13 , wherein generating the sequence of output embeddings comprises:

applying respective layer normalizations to the sequence of input embeddings and to the latent n-gram representation of the sequence of input embeddings for each head; and

determining a concatenation of the layer normalized sequence of input embeddings and the layer normalized latent n-gram representation of the sequence of input embeddings for each head.

18. The method of claim 13 , wherein the learned product quantization codebook has k cluster indices that each map to a respective code word, and wherein generating the discrete latent representation of the sequence of input embeddings for each head comprises, at each of the plurality of input embedding positions:

selecting a code word from among all code words in the learned product quantization codebook that has a least distance from the input embedding at the input embedding position; and

assigning the input embedding to the cluster index that maps to the selected code word.

19. The method of claim 13 , wherein a total number of the plurality of n-gram indices is k n , and wherein assigning to each of the plurality of input embedding positions the respective n-gram index comprises:

determining a combination of the corresponding cluster index at the input embedding position and a respective cluster index at each of n preceding input embedding positions multiplied by k.

20. The method of claim 13 , wherein generating the latent n-gram representation of the sequence of input embeddings comprises generating a hash value for each of the plurality of input embedding positions by using universal hashing techniques.

Assignments (1)
ASSIGNMENT OF ASSIGNOR'S INTEREST Recorded Nov 3, 2022
From: ANIL, ROHAN; ROY, AURKO
To: GOOGLE LLC
Reel/Frame 061650/0575 →
Continuity (1)
Related Publication 20240078379A1 · Mar 7, 2024
References Cited (60)
US 8725509B1 · Harb · 2014 [cited by examiner]
US 11373086B2 · Strope · 2022 [cited by examiner]
US 20050276235A1 · Lee · 2005 [cited by examiner]
US 20070078653A1 · Olsen · 2007 [cited by examiner]
US 20180240014A1 · Strope · 2018 [cited by examiner]
US 20200175961A1 · Thomson · 2020 [cited by examiner]
US 20240078379A1 · Anil · 2024 [cited by examiner]
JP 2017032738A · 2017 [cited by examiner]
WO WO2018152345A1 · 2018 [cited by examiner]
Henderson, Matthew, Rami Al-Rfou, Brian Strope, Yun-Hsuan Sung, László Lukács, Ruiqi Guo, Sanjiv Kumar, Balint Miklos, and Ray Kurzweil. “Efficient natural language response suggestion for smart reply.” arXiv preprint a… [cited by examiner]
Wu, Xiang, Ruiqi Guo, David Simcha, Dave Dopson, and Sanjiv Kumar. “Efficient inner product approximation in hybrid spaces.” arXiv preprint arXiv:1903.08690 (2019). [cited by examiner]
Agarwal et al., “Disentangling Adaptive Gradient Methods from Learning Rates,” CoRR, Submitted on Feb. 26, 2020, arXiv:2002.11803v1, 26 pages. [cited by applicant]
Ainslie et al., “ETC: Encoding Long and Structured Inputs in Transformers,” CoRR, Submitted on Oct. 27, 2020, arXiv:2004.08483v5, 17 pages. [cited by applicant]
Ba et al., “Layer Normalization,” CoRR, Submitted on Jul. 21, 2016, arXiv:1607.06450v1, 14 pages. [cited by applicant]
Bahdanau et al., “Neural Machine Translation by Jointly Learning to Align and Translate,” Presented at Proceedings of the 3rd International Conference on Learning Representations, San Diego, CA, May 7-9, 2015; CoRR, Sub… [cited by applicant]
Bottou et al., “Convergence Properties of the K-Means Algorithms,” Presented at the Advances in Neural Information Processing Systems Conference, Denver, Colorado, Nov. 28-Dec. 1, 1994; Advances in Neural Information Pr… [cited by applicant]
Brown et al., “Class-Based n-gram Models of Natural Language,” Computational Linguistics, Dec. 1992, 18(4):467-480. [cited by applicant]
Brown et al., “Language Models are Few-Shot Learners,” CoRR, Submitted on Jul. 22, 2020, arXiv:2005.14165v1, 75 pages. [cited by applicant]
Brown et al., “The Mathematics of Statistical Machine Translation: Parameter Estimation,” Computational Linguistics, Jun. 1993, 19(2):263-311. [cited by applicant]
Chen et al., “An Empirical Study of Smoothing Techniques for Language Modeling,” Computer Speech & Language, Oct. 1999, 13(4):359-394. [cited by applicant]
Cho et al., “Learning Phrase Representations Using RNN Encoder-Decoder for Statistical Machine Translation,” Presented at Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing, Oct. 25-2… [cited by applicant]
Correia et al., “Adaptively Sparse Transformers,” CoRR, Submitted on Sep. 6, 2019, arXiv:1909.00015v2, 20 pages. [cited by applicant]
Dauphin et al., “Language Modeling with Gated Convolutional Networks,” Presented at Proceedings of the 34th International Conference on Machine Learning, Sydney, Australia, Aug. 6-11, 2017, 9 pages. [cited by applicant]
Devlin et al., “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding,” CoRR, Submitted on May 24, 2019, arXiv:1810.04805v2, 16 pages. [cited by applicant]
Fedus et al., “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity,” CoRR, Submitted on Jan. 11, 2021, arXiv:2101.03961v1, 31 pages. [cited by applicant]
Ge et al., “Optimized Product Quantization for Approximate Nearest Neighbor Search,” Presented at the IEEE Conference on Computer Vision and Pattern Recognition, Portland, Oregon, Jun. 23-28, 2013, pp. 2946-2953. [cited by applicant]
Graves et al., “Neural Turing Machines,” CoRR, Submitted on Dec. 10, 2014, arXiv:1410.5401v2, 26 pages. [cited by applicant]
Guu et al., “REALM: Retrieval-Augmented Language Model Pre-Training,” CoRR, Submitted on Feb. 10, 2020, arXiv:2002.08909v1, 12 pages. [cited by applicant]
Hendrycks et al., “Gaussian Error Linear Units (GELUs),” CoRR, Submitted on Jul. 8, 2020, arXiv:1606.08415v4, 9 pages. [cited by applicant]
Huang et al., “Lookup-Table Recurrent Language Models for Long Tail Speech Recognition,” CoRR, Submitted on Jun. 7, 2021, arXiv:2104.04552v2, 5 pages. [cited by applicant]
Jegou et al., “Product Quantization for Nearest Neighbor Search,” IEEE Transactions on Pattern Analysis and Machine Intelligence, Jan. 2011, 33(1):117-128. [cited by applicant]
Kaiser et al., “Fast Decoding in Sequence Models Using Discrete Latent Variables,” CoRR, Submitted on Jun. 7, 2018, arXiv:1803.03382v6, 10 pages. [cited by applicant]
Katz, “Estimation of Probabilities from Sparse Data for the Language Model Component of a Speech Recognizer,” IEEE Transactions on Acoustics, Speech, and Signal Processing, Mar. 1987, 35(3):400-401. [cited by applicant]
Khandelwal et al., “Generalization Through Memorization: Nearest Neighbor Language Models,” CoRR, Submitted on Nov. 1, 2019, arXiv:1911.00172v1, 13 pages. [cited by applicant]
Kingma et al., “ADAM: A Method For Stochastic Optimization,” Presented at Proceedings of 3rd International Conference on Learning Representations, San Diego, CA, May 7-9, 2015; CoRR, Submitted on Jan. 30, 2017, arXiv:14… [cited by applicant]
Kneser et al., “Improved Backing-Off For M-Gram Language Modeling,” Presented at the International Conference on Acoustics, Speech, and Signal Processing, Detroit, MI, May 9-12, 1995, pp. 181-184. [cited by applicant]
Krishna et al., “Hurdles to Progress in Long-form Question Answering,” CoRR, Submitted on May 19, 2021, arXiv:2103.06332v2, 18 pages. [cited by applicant]
Lample et al., “Large Memory Layers with Product Keys,” CoRR, Submitted on Dec. 16, 2019, arXiv:1907.05242v2, 12 pages. [cited by applicant]
Lepikhin et al., “GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding,” CoRR, Submitted on Jun. 30, 2020, arXiv:2006.16668v1, 35 pages. [cited by applicant]
Lewis et al., “Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks,” CoRR, Submitted on Apr. 12, 2021, arXiv:2005.11401v4, 19 pages. [cited by applicant]
Mikolov et al., “Efficient Estimation of Word Representations in Vector Space,” CoRR, Submitted on Sep. 7, 2013, arXiv:1301.3781v3, 12 pages. [cited by applicant]
Panigrahy et al., “Sketch Based Memory for Neural Networks,” Presented at Proceedings of the International Conference on Artificial Intelligence and Statistics, held virtually, Apr. 13-15, 2021, 12 pages. [cited by applicant]
Raffel et al., “Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer,” CoRR, Submitted on Oct. 24, 2019, arXiv:1910.10683v2, 53 pages. [cited by applicant]
Roy et al., “Efficient Content-Based Sparse Attention with Routing Transformers,” Transactions of the Association for Computational Linguistics, Feb. 1, 2021, 9:53-68. [cited by applicant]
Roy et al., “Theory and Experiments on Vector Quantized Autoencoders,” CoRR, Submitted on Jul. 20, 2018, arXiv:1805.11063v2, 11 pages. [cited by applicant]
s3-us-west-2.amazonaws.com [offline], “Improving Language Understanding by Generative Pre-Training,” available on or before Jul. 31, 2018, via Internet Archive: Wayback Machine URL<https://web.archive.org/web/2018073114… [cited by applicant]
Schneider et al., “wav2vec: Unsupervised Pre-Training for Speech Recognition,” CoRR, Submitted on Sep. 11, 2019, arXiv:1904.05862v4, 9 pages. [cited by applicant]
Shazeer, “GLU Variants Improve Transformer,” CoRR, Submitted on Feb. 12, 2020, arXiv:2002.05202v1, 5 pages. [cited by applicant]
So et al., “Primer: Searching for Efficient Transformers for Language Modeling,” CoRR, Submitted on Sep. 17, 2021, arXiv:2109.08668v1, 34 pages. [cited by applicant]
Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding,” CoRR, Submitted on Oct. 9, 2021, arXiv:2104.09864v2, 14 pages. [cited by applicant]
Thorup, “High Speed Hashing for Integers and Strings,” CoRR, Submitted on Sep. 15, 2015, arXiv:1504.06804v3, 17 pages. [cited by applicant]
Van den Oord et al., “Neural Discrete Representation Learning,” CoRR, Submitted on Nov. 2, 2017, arXiv:1711.00937v1, 10 pages. [cited by applicant]
Vaswani et al., “Attention is all you Need,” CoRR, Submitted on Dec. 6, 2017, arXiv:1706.03762v5, 15 pages. [cited by applicant]
Wang et al., “Style Tokens: Unsupervised Style Modeling, Control and Transfer in End-to-End Speech Synthesis,” Presented at the International Conference on Machine Learning, Stockholm, Sweden, Jul. 10-15, 2018, 10 pages. [cited by applicant]
Weston et al., “Memory Networks,” CoRR, Submitted on Dec. 24, 2014, arXiv:1410.3916v4, 14 pages. [cited by applicant]
Wu et al., “CvT: Introducing Convolutions to Vision Transformers, ” CoRR, Submitted on Mar. 29, 2021, arXiv:2103.15808v1, 10 pages. [cited by applicant]
Wu et al., “Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation,” CoRR, Submitted Oct. 8, 2016, arXiv:1609.08144v2, 23 pages. [cited by applicant]
Yu et al., “Product Quantization Network for Fast Image Retrieval,” Presented at Proceedings of the European Conference on Computer Vision (ECCV), Munich, Germany, Sep. 8-14, 2018, 16 pages. [cited by applicant]
Zaheer et al., “Big Bird: Transformers for Longer Sequences,” Presented at the Advances in Neural Information Processing Systems, Vancouver, Canada, Dec. 6-12, 2020; Advances in Neural Information Processing Systems 33,… [cited by applicant]
Zeghidour et al., “Wavesplit: End-to-End Speech Separation by Speaker Clustering,” IEEE/ACM Transactions on Audio, Speech, and Language Processing, Jul. 2021, 29:11 pages. [cited by applicant]