Package: wavspa
WavSpA: Wavelet Space Attention for Enhancing Transformer's Long Sequence LearningWelcome to the official repository of theWavSpApaper. This innovative work introduces adaptive wavelet transform techniques coupled with Transformer models to excel at processing long sequences. The implementation is crafted usingJAXalongsideFlaxfor robustness and efficiency.InstallationSetup your environment to run the models using the providedrequirements.txt:$pipinstall-rrequirements.txt
$pipinstallwavspaNote: This codebase supports JAX version 0.3.13.Our experimentsWe demonstrate substantial gains in performance with WavSpA across various attention-based architectures. The framework offers three parametrization options:Adaptive Wavelet (AdaWavSpA)Orthogonal Adaptive Wavelet (OrthoWavSpA)Wavelet Lifting (LiftWavSpA)Performance metrics are as follows:ModelsListOpsTextRetrievalImagePathfinderAvgAvg (w/o r)Transformer36.3764.2757.4642.4471.4054.3953.62AdaWavSpA55.4081.6079.2755.5881.1270.5968.43OrthoWavSpA45.9581.6371.5249.2981.1365.9064.50LiftWavSpA42.9575.6356.4542.4881.7359.8560.70------------------------Longformer35.6362.8556.8942.2269.7153.4652.60AdaWavSpA49.3079.7358.5750.8479.4863.6664.93OrthoWavSpA39.4578.4179.9349.9379.4754.9654.96LiftWavSpA39.4078.0053.2740.9575.8057.4858.54------------------------Linformer35.7053.9452.2738.4766.4449.3648.64AdaWavSpA37.1554.7561.0934.9365.6650.7248.12OrthoWavSpA38.0556.9360.2539.4565.3552.0149.95LiftWavSpA37.3054.4370.7334.6663.4952.1247.47------------------------Linear Att.16.1365.9053.0942.3275.9150.6750.06AdaWavSpA38.9076.8271.3854.8179.6864.3262.55OrthoWavSpA39.5579.4569.6549.9378.0955.8655.86LiftWavSpA38.3573.3954.0644.3974.4656.9357.65------------------------Performer18.0165.4053.8242.7777.0551.4150.81AdaWavSpA46.0580.9371.1652.0677.1765.4764.05OrthoWavSpA39.8079.1057.6748.7878.0960.6961.44LiftWavSpA39.8575.9652.7539.9776.2056.9558.00Example UsageFor implementation details, see lra_benchmarks/models/wavspa/wavspa_learn.py. The wavelet initialization and transformation processes are crucial:def setup(self):
  ## db initialization
  assert self.wlen % 2 == 0, "incompatible"        
  self.eps = 1e-4
  if "lift" in self.wavelet:
      self.adawave_est = self.param('adawave_est', nn.initializers.normal(stddev=0.02), (self.wlen, self.wav_dim), self.dtype)
      self.adawave_pred = self.param('adawave_pred', nn.initializers.normal(stddev=0.02), (self.wlen, self.wav_dim), self.dtype)
  elif "ortho" in self.wavelet:
      L = int(self.wlen / 2)
      S = jnp.zeros(shape=[2*L, 2*L], dtype=int)
      i = jnp.asarray(range(2*L))
      j = jnp.asarray(range(1, 2*L+1)) % (2*L)
      S = S.at[i, j].set(1)
      self.S = sparse.BCOO.fromdense(S, nse=2*L)
      self.S_inv = jnp.linalg.inv(S)
      self.thetas = self.param('thetas', nn.initializers.uniform(2*jnp.pi), (L, self.wav_dim), self.dtype)
  elif "db" in self.wavelet:
      self.adawave = self.param('adawave', db_init, (self.wlen, self.wav_dim), self.dtype)
  elif "sin" in self.wavelet:
      self.adawave = self.param('adawave', sin_init, (self.wlen, self.wav_dim), self.dtype)
  else:
      # default to daubechie wave, non trainable
      self.adawave = db_init(key=None, shape=(self.wlen, self.wav_dim), dtype=self.dtype)Then for forward and backward wavelet transform:z = wavspa.wavedec_learn(x, wavelet, level=self.level)
for level in range(len(z)):
  z[level] = nn.SelfAttention(num_heads=self.num_heads,
                              dtype=self.dtype,
                              qkv_features=self.qkv_dim,
                              kernel_init=nn.initializers.xavier_uniform(),
                              bias_init=nn.initializers.normal(stddev=1e-6),
                              use_bias=False,
                              broadcast_dropout=False,
                              dropout_rate=self.attention_dropout_rate,
                              decode=False)(z[level], deterministic=deterministic)
z = wavspa.waverec_learn(z, wavelet)[:,:inputs.shape[1],:]DatasetsAccess and instructions for LRA, D2A, and CodeXGlue datasets:Long Range ArenaD2A IBM official releaseCodeXGlueTo execute a task, use the train_best.py script with the appropriate configurations:PYTHONPATH="$(pwd)":"$PYTHONPATH" python lra_benchmarks/listops/train_best.py \
      --config=lra_benchmarks/listops/configs/wavspa-exp0.py \
      --model_dir=/tmp/listops \
      --task_name=basic \
      --data_dir=$HOME/lra_data/listops/CitationIf you find out work useful, please cite our paper at:@inproceedings{
zhuang2023wavspa,
title={WavSpA: Wavelet Space Attention for Boosting Transformers' Long Sequence Learning Ability},
author={Yufan Zhuang and Zihan Wang and Fangbo Tao and Jingbo Shang},
booktitle={UniReps:  the First Workshop on Unifying Representations in Neural Models at NeurIPS 2023},
year={2023},
url={https://openreview.net/forum?id=yC6b3hqyf8}
}