Blockwise Self-Attention for Long Document Understanding
Table 1: Pre-training Performance Analysis.
['[ITALIC] N', 'Model', 'Training Time\xa0(day)', 'Memory\xa0(per GPU, GB)', 'Heads Config.', 'Valid. ppl']
To validate our claim that BlockBert with n×n blocks can reduce the O(N2) memory use by a factor of n Again, We fix the number of tokens in each GPU (b×N=4096) and choose N from {128,256,512,1024,2048}. When we set block size to be 2 and 3 for BlockBert, their estimated O(N2) activation memory decreases to 1/2 and 1/3 of BERT’s O(N2) activation memory, resp. The saving is more significant for longer sequences. When N=1024, the overall memory reduction of BlockBert with 2 / 3 blocks is 27.3% / 36.1%, resp. Besides memory saving, we also achieve a significant speedup. For example, when N=1024, BlockBert (n=2) reduces the training time from RoBERTa’s 9.7 days to 7.5 days.