Package: lsg-converter
LSG Attention: Extrapolation of pretrained Transformers to long sequencesArXivpaperAccepted as a conference paper in PAKDD 2023.Requirestransformers >= 4.36.1Optional package to convert models:pipinstalllsg-converterCompatible modelsEfficiencyConversionUsageLSG-LayerLSG-AttentionExperimentsCompatible modelsNote that for non-compatible models, the LSG attention layer is available, seeLSG-Layerandlsg_converter/attention_layersfor more informations.This package can convert various HuggingFace checkpoints and architectures to their LSG (Local-Sparse-Global) variant to handle long sequences.Available model types:ALBERT ["albert"]BART ["bart"] (encoder attention modified only)BARThez ["barthez"] (encoder attention modified only)BERT ["bert"]CamemBERT ["camembert"]DistilBERT ["distilbert"]Electra ["electra"]mBART ["mbart"] (not tested extensively, encoder attention modified only)Pegasus ["pegasus"] (not tested extensively, encoder attention modified only)RoBERTa ["roberta"]XLM-RoBERTa ["xlm-roberta"]Some converted checkpoints are availablehere.EfficiencyMemory and training speed for a binary classification task with 4x4096 tokens batches and Adam (RTX 8000).ModelsSeconds per stepMemory w/ attn dropoutMemory w/o attn dropoutLongformer-base3.22 s/step34.38 Gb32.83 GbBigBird-RoBERTa-base2.85 s/step38.13 Gb38.13 Gb (no effect)LSG-RoBERTa 256/01.40 s/step32.92 Gb24.80 GbLSG-RoBERTa 128/128 (norm)1.51 s/step33.80 Gb27.52 GbLSG-RoBERTa 32/32 (norm)1.20 s/step24.53 Gb22.53 GbConvert checkpoint to LSGModels can be converted with or without thelsg-converterpackage.With packageBERT example with thelsg-converterpackage:fromlsg_converterimportLSGConverterconverter=LSGConverter(max_sequence_length=4096)# Example 1model,tokenizer=converter.convert_from_pretrained("bert-base-uncased",num_global_tokens=7)print(type(model))# <class 'lsg_converter.bert.modeling_lsg_bert.LSGBertForMaskedLM'># Example 2model,tokenizer=converter.convert_from_pretrained("bert-base-uncased",architecture="BertForSequenceClassification",use_auth_token=True)print(type(model))# <class 'lsg_converter.bert.modeling_lsg_bert.LSGBertForSequenceClassification'>Without packageUseconvert_checkpoint.pyto convert a model (check model_type from config.json).The architecture of the model is inferred from the config file, but you can specify a different one if the config is incorrect (which can happen for BART models), seepython convert_checkpoint.py --help.To test the converted model, add--run_test(experimental).RoBERTa example (fromRobertaForMaskedLMtoRobertaForSequenceClassification) without package:gitclonehttps://github.com/ccdv-ai/convert_checkpoint_to_lsg.gitcdconvert_checkpoint_to_lsgexportMODEL_TO_CONVERT=roberta-baseexportMODEL_NAME=lsg-roberta-baseexportMAX_LENGTH=4096pythonconvert_checkpoint.py\--initial_model$MODEL_TO_CONVERT\--model_name$MODEL_NAME\--model_kwargs"{'mask_first_token': true, 'sparsity_type': 'lsh', 'block_size': 32}"\--architectureRobertaForSequenceClassification\--max_sequence_length$MAX_LENGTHModel UsageWorks with the AutoClass.fromtransformersimportAutoTokenizer,AutoModelForMaskedLM# Load created modelMODEL_NAME="lsg-roberta-base"SENTENCE="This is a test sentence."model=AutoModelForMaskedLM.from_pretrained(MODEL_NAME,trust_remote_code=True)tokenizer=AutoTokenizer.from_pretrained(MODEL_NAME)inputs=tokenizer(SENTENCE,return_tensors="pt")model(**inputs)fromtransformersimportAutoTokenizer,AutoModelForSequenceClassification# Load created modelMODEL_NAME="lsg-roberta-base"SENTENCE="This is a test sentence."model=AutoModelForSequenceClassification.from_pretrained(MODEL_NAME,trust_remote_code=True)tokenizer=AutoTokenizer.from_pretrained(MODEL_NAME)inputs=tokenizer(SENTENCE,return_tensors="pt")model(**inputs)LSG-LayerYou can find a self contained implementation of LSG attention (with no additional params), seelsg_converter/attention_layers.Doesn't work for Cross Attention because the local context is ambiguous to define in this case.Usage:importtorchfromlsg_converter.attention_layersimportBlockLocalSelfAttention,LSGSelfAttention# batch, num_heads, sequence length, hidden_sizen,h,t,d=2,4,58,32Q,K,V=torch.randn(n,h,t,d),torch.randn(n,h,t,d),torch.randn(n,h,t,d)attention_mask=torch.zeros(n,1,1,t).float()# Only block local attention with 1 global connectionattn=BlockLocalSelfAttention(block_size=16,compute_global_attention=True,is_causal=False,attention_dropout_prob=0.1)# LSG attention with 1 global connectionattn=LSGSelfAttention(block_size=32,sparsity_factor=8,sparsity_type="bos_pooling",compute_global_attention=True,is_causal=True)# expect (batch, num_heads, sequence_length, hidden_size) inputs,# attention_mask is (batch, 1, 1, sequence_length)# causal mask is built on the fly but (batch, 1, sequence_length, sequence_length) mask is possibleoutputs=attn(Q,K,V,attention_mask)print(outputs.shape)>torch.Size([2,4,58,32])Example: replacing Self Attention in GPT2 (from Huggingface):fromtransformers.models.gpt2import*fromlsg_converter.attention_layersimportBlockLocalSelfAttentionclassGPT2BlockLocalAttention(modeling_gpt2.GPT2Attention):def__init__(self,config,is_cross_attention=False,layer_idx=None):super().__init__(config,is_cross_attention,layer_idx)self.attn=BlockLocalSelfAttention(block_size=32,compute_global_attention=True,is_causal=True)def_attn(self,query,key,value,attention_mask=None,head_mask=None):returnself.attn(query,key,value,attention_mask),Nonemodeling_gpt2.GPT2Attention=GPT2BlockLocalAttentionNote that for generation (inference on causal modeling), full attention is used after the first step.This may change in the future.LSG AttentionLSG relies on 3 components: block local attention, sparse attention and prepended global attention.ParametersYou can change various parameters like :local block size (block_size=128)sparse block size (sparse_block_size=128)sparsity factor (sparsity_factor=2)mask_first_token (mask first token since it is redundant with the first global token)the number of global tokens (num_global_tokens=1)see config.json fileSparse selection typeThere are 6 different sparse selection patterns. The best type is task dependent.Ifsparse_block_size=0orsparsity_type="none", only local attention is considered.Note that for sequences with length < 2*block_size, the type has no effect.sparsity_type="bos_pooling"(new)weighted average pooling using the BOS tokenWorks best in general, especially with a rather large sparsity_factor (8, 16, 32)Additional parameters:Nonesparsity_type="norm", select highest norm tokensWorks best for a small sparsity_factor (2 to 4)Additional parameters:Nonesparsity_type="pooling", use average pooling to merge tokensWorks best for a small sparsity_factor (2 to 4)Additional parameters:Nonesparsity_type="lsh", use the LSH algorithm to cluster similar tokensWorks best for a large sparsity_factor (4+)LSH relies on random projections, thus inference may differ slightly with different seedsAdditional parameters:lsg_num_pre_rounds=1, pre merge tokens n times before computing centroidssparsity_type="stride", use a striding mecanism per headEach head will use different tokens strided by sparsify_factorNot recommended if sparsify_factor > num_headssparsity_type="block_stride", use a striding mecanism per headEach head will use block of tokens strided by sparsify_factorNot recommended if sparsify_factor > num_headsExperimentsSeeexperiments/