Fibottention

Inceptive Visual Representation Learning with Diverse Attention Across Heads

Ali Khaleghi Rahimian¹ , Manish Kumar Govind¹ , Subhajit Maity² , Dominick Reilly¹ , Christian Kümmerle¹* , Srijan Das¹* , and Aritra Dutta²*

* Equal contribution as Project Lead

Affiliations:

¹ University of North Carolina at Charlotte
² University of Central Florida

Overview

This repository contains the implementation of our proposed Fibottention mechanism and related algorithms.

The paper is now available on arXiv.

Abstract

Vision Transformers and their variants have achieved remarkable success in diverse visual perception tasks. Despite their effectiveness, they suffer from two significant limitations. First, the quadratic computational complexity of multi-head self-attention (MHSA), which restricts scalability to large token counts, and second, a high dependency on large-scale training data to attain competitive performance. In this paper, to address these challenges, we propose a novel sparse self-attention mechanism named Fibottention. Fibottention employs structured sparsity patterns derived from the Wythoff array, enabling an $O(N \log N)$ computational complexity in self-attention. By design, its sparsity patterns vary across attention heads, which provably reduces redundant pairwise interactions while ensuring sufficient and diverse coverage. This leads to an inception-like functional diversity in the attention heads, and promotes more informative and disentangled representations. We integrate Fibottention into standard Transformer architectures and conduct extensive experiments across multiple domains, including image classification, video understanding, and robot learning. Results demonstrate that models equipped with Fibottention either significantly outperform or achieve on-par performance with their dense MHSA counterparts, while leveraging only 2% of all pairwise interactions across self-attention heads in typical settings, resulting in substantial computational savings. Moreover, when compared to existing sparse attention mechanisms, Fibottention consistently achieves superior results on a FLOP-equivalency basis. Finally, we provide an in-depth analysis of the enhanced feature diversity resulting from our attention design and discuss its implications for efficient representation learning.

Installation and training of Fibottention

Image classification

Use the commands below to install the required packages when setting up your environment:

conda create --name env_name --no-default-packages python=3.7
pip install -r requirements.txt

To train the model, use the script.sh file which executes main_finetune.py with specified parameters.

./script.sh [id] [out_dir] [model] [dataset] [classes] [device] [batch] [mask_ratio]

• id: An identifier for the execution.
• out_dir: The output directory for saving results.
• model: The type of model for finetuning.
• dataset: The name of the dataset to be used.
• classes: The number of classes in the dataset.
• device: The GPU device number or ID.
• batch: The batch size for training.
• mask_ratio: The ratio for masking during training.

Action recognition

Install dependencies

• torchvision: pip install torchvision or conda install torchvision -c pytorch
• fvcore: pip install 'git+https://github.com/facebookresearch/fvcore'
• simplejson: pip install simplejson
• einops: pip install einops
• timm: pip install timm
• PyAV: conda install av -c conda-forge
• psutil: pip install psutil
• scikit-learn: pip install scikit-learn
• OpenCV: pip install opencv-python
• tensorboard: pip install tensorboard
• matlotlib : pip install matplotlib

DataSet preparation

The dataset could be structured as follows:

├── data
    ├── Action_01
        ├── Video_01.mp4
        ├── Video_02.mp4
        ├── …

After all the data is prepared, resize and crop the video to person-centric to get rid of background noise. Then, prepare the CSV files for the training, validation, and testing sets as train.csv, val.csv, and test.csv. The format of the CSV file is:

path_to_video_1 label_1
path_to_video_2 label_2
path_to_video_3 label_3
...
path_to_video_N label_N

Training

We provide configs to train fibottention for action recognition on Smarthome, NTU and NUCLA datasets in action_recognition/configs/. Please update the paths in the config to match the paths in your machine before using.

For example to train on Smarthome using 8 GPUs run the following command:

python action_recognition/tools/run_net.py --cfg configs/SMARTHOME.yaml NUM_GPUS 8

Robot learning

Install dependencies & download data

Our robot learning code is built on top of the code for "Crossway Diffusion: Improving Diffusion-based Visuomotor Policy via Self-supervised Learning". Please follow their instructions for installing dependencies and obtaining the data.

• Follow the installation instructions from this link
• Follow instructions from this link to download the datasets
  • This page will direct you to this download page. To reproduce the results in this work, only pusht.zip and robomimic_image.zip are needed

Training

To perform the robot learning experiments

1. Navigate to robot_learning/
2. Run the following command, replacing <DATASET> with the desired dataset (this work uses can_ph, lift_ph, or pusht):

train.py --config-dir=config/<DATASET>/ --config-name=typea.yaml training.seed=42 hydra.run.dir=outputs/vit-b-fibottention/${now:%Y-%m-%d}/${now:%H-%M-%S}_${task_name}_${task.dataset_type}

Algorithms

Algorithm 1: Generating Fibonacci Sequence with Constraint

# INPUT: a, b, w_i  # Initial two numbers and upper constraint
# OUTPUT: fib_seq  # Fibonacci sequence under constraint

fib_seq = [a, b]  # Initialize sequence with first two numbers
while fib_seq[-1] < w_i:  # Generate sequence until last number is less than w_i
    next_num = fib_seq[-1] + fib_seq[-2]  # Calculate next Fibonacci number
    fib_seq.append(next_num)  # Append new number to the sequence
return fib_seq  # Return the generated Fibonacci sequence

Algorithm 2: Generating Mask for All Heads

# INPUT: L, N, w_min, w_max, is_modified  # Layer, size, window min/max, and modification flag
# OUTPUT: Ω ∈ (0,1)^(h × (N+1) × (N+1))  # Output mask tensor

phi = (1 + sqrt(5)) / 2  # Golden ratio for indexing
for i in range(1, h + 1):  # Loop over each head
    a = int(i * phi * phi)  # Fibonacci starting values based on golden ratio
    b = int(i * phi * phi**2)
    w_i = w_min + int((i - 1) * (w_max - w_min) / (h - 1))  # Window size for this head
    Θ = [[0]*N for _ in range(N)]  # Initialize intermediate mask

    if is_modified:  # Modify sequence if required
        b_Wyt_m = b - a
        a_Wyt_m = a - b_Wyt_m
        I = get_fibonacci(a_Wyt_m, b_Wyt_m, w)   # Calculate Fibonacci indices using Algorithm 1
    else:
        I = get_fibonacci(a, b, w)
    
    for o in I:  # Apply Fibonacci indices to mask
        for j in range(N-o):
            Θ[j][j+1] = 1  # Upper triangular masking
        for k in range(o, N):
            Θ[k+1][k] = 1  # Lower triangular masking
    
    Ω_i = [[1]*(N+1) for _ in range(N+1)]  # Initialize output mask for head i
    for j in range(1, N+1):  # Fill in mask based on Θ
        for k in range(1, N+1):
            Ω_i[j][k] = Θ[j-1][k-1]

Ω = [Ω_i for i in range(h)]  # Combine masks from all heads
Ω = randomshuffle(L, Ω)  # Randomly shuffle masks across layers
return Ω  # Return the final mask tensor

Algorithm 3: Fibottention in a Single Vision Transformer Block

# INPUT: X ∈ R^(N+1 × d)  # Input feature matrix
# OUTPUT: O ∈ R^(N+1 × d)  # Output feature matrix
# PARAMETERS: W_i^Q, W_i^K, W_i^V ∈ R^(d × d_h), d_h = d / h  # Weights for Q, K, V
# HYPERPARAMETERS: w_min, w_max, is_modified  # Window sizes and modification flag

iota_Ω = getMask(L, N, h, w_min, w_max, is_modified)  # Get mask from Algorithm 2

for i in range(1, h + 1):  # Process each attention head
    Q_i = X @ W_i^Q  # Query matrix
    K_i = X @ W_i^K  # Key matrix
    V_i = X @ W_i^V  # Value matrix
    A_i = Q_i @ K_i.T  # Attention scores
    A_i_Ω = np.sign(A_i) * (np.abs(A_i) * iota_Ω[i,:,:])  # Apply mask to attention scores
    A_i_Ω = softmax(A_i_Ω)  # Softmax to normalize scores
    Z_i = A_i_Ω @ V_i  # Weighted sum to produce output for head

Z = np.concatenate([Z_i for i in range(h)], axis=1)  # Concatenate outputs from all heads
O = Z @ W^Z  # Project concatenated outputs to final dimension
return O  # Return output of Vision Transformer block

Acknowledgement

This repository is built on top of MAE, TimeSformer, and Crossway Diffusion. We would like to thank all the contributors for their well-organized codebases.

Citation

@article{rahimian2024fibottention,
  title = {Fibottention: Inceptive Visual Representation Learning with Diverse Attention Across Heads},
  author = {Rahimian, Ali Khaleghi and Govind, Manish Kumar and Maity, Subhajit and Reilly, Dominick and Kümmerle, Christian and Das, Srijan and Dutta, Aritra},
  journal = {arXiv preprint arXiv:2406.19391},
  year = {2024},
  url = {https://arxiv.org/abs/2406.19391}
}

License

This project is licensed under the Creative Commons Attribution 4.0 International - see the LICENSE website for details.