title learning to tokenize in vision transformers s aritra roy gosthipatyhttps twitter comarig23498 sayak paulhttps twitter comrisingsayak equal contribution converted to keras 3 by muhammad anas razahttps anasrz com date created 20211210 last modified 20230814 description adaptively generating a smaller number of tokens for vision transformers accelerator gpu introduction vision transformers dosovitskiy et al https arxiv orgabs2010 11929 and many other transformerbased architectures liu et al https arxiv orgabs2103 14030 yuan et al https arxiv orgabs2101 11986 etc have shown strong results in image recognition the following provides a brief overview of the components involved in the vision transformer architecture for image classification extract small patches from input images linearly project those patches add positional embeddings to these linear projections run these projections through a series of transformer vaswani et al https arxiv orgabs1706 03762 blocks finally take the representation from the final transformer block and add a classification head if we take 224x224 images and extract 16x16 patches we get a total of 196 patches also called tokens for each image the number of patches increases as we increase the resolution leading to higher memory footprint could we use a reduced number of patches without having to compromise performance ryoo et al investigate this question in tokenlearner adaptive spacetime tokenization for videoshttps openreview netforum idzl1kpdxs88 they introduce a novel module called tokenlearner that can help reduce the number of patches used by a vision transformer vit in an adaptive manner with tokenlearner incorporated in the standard vit architecture they are able to reduce the amount of compute measured in flops used by the model in this example we implement the tokenlearner module and demonstrate its performance with a mini vit and the cifar10 dataset we make use of the following references official tokenlearner codehttps github comgoogleresearchscenicblobmainscenicprojectstokenlearnermodel py image classification with vits on keras iohttps keras ioexamplesvisionimageclassificationwithvisiontransformer tokenlearner slides from neurips 2021https nips ccmedianeurips2021slides26578 pdf imports hyperparameters please feel free to change the hyperparameters and check your results the best way to develop intuition about the architecture is to experiment with it data optimizer training augmentation vit architecture tokenlearner load and prepare the cifar10 dataset load the cifar10 dataset convert to tf data dataset objects data augmentation the augmentation pipeline consists of rescaling resizing random cropping fixedsized or random sized random horizontal flipping note that image data augmentation layers do not apply data transformations at inference time this means that when these layers are called with trainingfalse they behave differently refer to the documentationhttps keras ioapilayerspreprocessinglayersimageaugmentation for more details positional embedding module a transformerhttps arxiv orgabs1706 03762 architecture consists of multihead self attention layers and fullyconnected feed forward networks mlp as the main components both these components are permutation invariant they re not aware of feature order to overcome this problem we inject tokens with positional information the positionembedding function adds this positional information to the linearly projected tokens mlp block for transformer this serves as the fully connected feed forward block for our transformer iterate over the hidden units and add dense dropout tokenlearner module the following figure presents a pictorial overview of the module sourcehttps ai googleblog com202112improvingvisiontransformerefficiency html tokenlearner module gifhttps blogger googleusercontent comimgaavvxseiylt3nmd9tztnz3g3vb4etnl5sowtgjoad6t2we7fsjxspbldpuprlinahte5hgcapfytjtriokflylygcyxvh1ksfhc1zcc8gw6gktvrqeskomrealugd5srl53izdgjc1iyceloxtfuixkj2addhobaujhu8lxtvdr2e7bcvafgvhhkmaw640h208 the tokenlearner module takes as input an imageshaped tensor it then passes it through multiple singlechannel convolutional layers extracting different spatial attention maps focusing on different parts of the input these attention maps are then elementwise multiplied to the input and result is aggregated with pooling this pooled output can be trated as a summary of the input and has much lesser number of patches 8 for example than the original one 196 for example using multiple convolution layers helps with expressivity imposing a form of spatial attention helps retain relevant information from the inputs both of these components are crucial to make tokenlearner work especially when we are significantly reducing the number of patches layer normalize the inputs applying conv2d reshape permute the reshape and permute is done to help with the next steps of multiplication and global average pooling 3 layers of conv with gelu activation as suggested in the paper this conv layer will generate the attention maps reshape and permute reshape the input to align it with the output of the conv block elementwise multiplication of the attention maps and the inputs global average pooling the element wise multiplication result transformer block layer normalization 1 multi head self attention layer 1 skip connection 1 layer normalization 2 mlp layer 1 skip connection 2 vit model with the tokenlearner module augment data create patches and project the pathces add positional embeddings to the projected patches iterate over the number of layers and stack up blocks of transformer add a transformer block add tokenlearner layer in the middle of the architecture the paper suggests that anywhere between 12 or 34 will work well layer normalization and global average pooling classify outputs create the keras model as shown in the tokenlearner paperhttps openreview netforum idzl1kpdxs88 it is almost always advantageous to include the tokenlearner module in the middle of the network training utility initialize the adamw optimizer compile the model with the optimizer loss function and the metrics define callbacks train the model train and evaluate a vit with tokenlearner results we experimented with and without the tokenlearner inside the mini vit we implemented with the same hyperparameters presented in this example here are our results tokenlearner tokens inbr tokenlearner top1 accbraveraged across 5 runs gflops tensorboard n 56 112 0 0184 linkhttps tensorboard devexperimentvkcwm49dqz2rik0zt4mj7w y 8 56 55 0 0153 linkhttps tensorboard devexperimentvkcwm49dqz2rik0zt4mj7w n 56 37 0 0184 linkhttps tensorboard devexperimenthdyj4wznqrowqztgbtmztq y 4 56 4980 0 0147 linkhttps tensorboard devexperimenthdyj4wznqrowqztgbtmztq n transformer layers 8 55 36 0 0359 linkhttps tensorboard devexperimentsepbk5znsaotdceeg6sv9w tokenlearner is able to consistently outperform our mini vit without the module it is also interesting to notice that it was also able to outperform a deeper version of our mini vit with 8 layers the s also report similar observations in the paper and they attribute this to the adaptiveness of tokenlearner one should also note that the flops count decreases considerably with the addition of the tokenlearner module with less flops count the tokenlearner module is able to deliver better results this aligns very well with the s findings additionally the s introducedhttps github comgoogleresearchscenicblobmainscenicprojectstokenlearnermodel pyl104 a newer version of the tokenlearner for smaller training data regimes quoting the s instead of using 4 conv layers with small channels to implement spatial attention this version uses 2 grouped conv layers with more channels it also uses softmax instead of sigmoid we confirmed that this version works better when having limited training data such as training with imagenet1k from scratch we experimented with this module and in the following table we summarize the results groups tokens top1 acc gflops tensorboard 4 4 54 638 0 0149 linkhttps tensorboard devexperimentkmfkgqagqjikew85physmw 8 8 54 898 0 0146 linkhttps tensorboard devexperiment0ppgyoq9rfwv9njx6njq2w 4 8 55 196 0 0149 linkhttps tensorboard devexperimentwukrhbzasdu3zrfmy4etzg please note that we used the same hyperparameters presented in this example our implementation is available in this notebookhttps github comarig23498tokenlearnerblobmastertokenlearnerv1 1 ipynb we acknowledge that the results with this new tokenlearner module are slightly off than expected and this might mitigate with hyperparameter tuning note to compute the flops of our models we used this utilityhttps github comadityakane2001regnetyblobmainregnetyutilsmodelutils pyl27 from this repositoryhttps github comadityakane2001regnety number of parameters you may have noticed that adding the tokenlearner module increases the number of parameters of the base network but that does not mean it is less efficient as shown by dehghani et al https arxiv orgabs2110 12894 similar findings were reported by bello et al https arxiv orgabs2103 07579 as well the tokenlearner module helps reducing the flops in the overall network thereby helping to reduce the memory footprint final notes tokenfuser the s of the paper also propose another module named tokenfuser this module helps in remapping the representation of the tokenlearner output back to its original spatial resolution to reuse the tokenlearner in the vit architecture the tokenfuser is a must we first learn the tokens from the tokenlearner build a representation of the tokens from a transformer layer and then remap the representation into the original spatial resolution so that it can again be consumed by a tokenlearner note here that you can only use the tokenlearner module once in entire vit model if not paired with the tokenfuser use of these modules for video the s also suggest that tokenfuser goes really well with vision transformers for videos arnab et al https arxiv orgabs2103 15691 we are grateful to jarvislabshttps jarvislabs ai and google developers expertshttps developers google comprogramsexperts program for helping with gpu credits also we are thankful to michael ryoo first of tokenlearner for fruitful discussions trained model demo generic badgehttps img shields iobadge20modeltokenlearnerblack svghttps huggingface cokerasiolearningtotokenizeinvit generic badgehttps img shields iobadge20spacestokenlearnerblack svghttps huggingface cospaceskerasiotokenlearner introduction vision transformers dosovitskiy et al https arxiv org abs 2010 11929 and many other transformer based architectures liu et al https arxiv org abs 2103 14030 yuan et al https arxiv org abs 2101 11986 etc have shown strong results in image recognition the following provides a brief overview of the components involved in the vision transformer architecture for image classification extract small patches from input images linearly project those patches add positional embeddings to these linear projections run these projections through a series of transformer vaswani et al https arxiv org abs 1706 03762 blocks finally take the representation from the final transformer block and add a classification head if we take 224x224 images and extract 16x16 patches we get a total of 196 patches also called tokens for each image the number of patches increases as we increase the resolution leading to higher memory footprint could we use a reduced number of patches without having to compromise performance ryoo et al investigate this question in tokenlearner adaptive space time tokenization for videos https openreview net forum id z l1kpdxs88 they introduce a novel module called tokenlearner that can help reduce the number of patches used by a vision transformer vit in an adaptive manner with tokenlearner incorporated in the standard vit architecture they are able to reduce the amount of compute measured in flops used by the model in this example we implement the tokenlearner module and demonstrate its performance with a mini vit and the cifar 10 dataset we make use of the following references official tokenlearner code https github com google research scenic blob main scenic projects token_learner model py image classification with vits on keras io https keras io examples vision image_classification_with_vision_transformer tokenlearner slides from neurips 2021 https nips cc media neurips 2021 slides 26578 pdf imports hyperparameters please feel free to change the hyperparameters and check your results the best way to develop intuition about the architecture is to experiment with it data optimizer training augmentation we will resize input images to this size size of the patches to be extracted from the input images vit architecture tokenlearner load and prepare the cifar 10 dataset load the cifar 10 dataset convert to tf data dataset objects data augmentation the augmentation pipeline consists of rescaling resizing random cropping fixed sized or random sized random horizontal flipping note that image data augmentation layers do not apply data transformations at inference time this means that when these layers are called with training false they behave differently refer to the documentation https keras io api layers preprocessing_layers image_augmentation for more details positional embedding module a transformer https arxiv org abs 1706 03762 architecture consists of multi head self attention layers and fully connected feed forward networks mlp as the main components both these components are _permutation invariant_ they re not aware of feature order to overcome this problem we inject tokens with positional information the position_embedding function adds this positional information to the linearly projected tokens mlp block for transformer this serves as the fully connected feed forward block for our transformer iterate over the hidden units and add dense dropout tokenlearner module the following figure presents a pictorial overview of the module source https ai googleblog com 2021 12 improving vision transformer efficiency html tokenlearner module gif https blogger googleusercontent com img a avvxseiylt3_nmd9 tztnz3g3vb4etn l5sowtgjoad6t2we7fsjxspbldpuprlinahte5hgca_pfytjtriokflylygcyxvh1ksfh_c1zc c8gw6gktvrqeskomrea_lu_gd5srl5 3izdgjc1iyceloxtfuixkj2addhobaujhu8lxtvdr2e7bcvafgvhhkma w640 h208 the tokenlearner module takes as input an image shaped tensor it then passes it through multiple single channel convolutional layers extracting different spatial attention maps focusing on different parts of the input these attention maps are then element wise multiplied to the input and result is aggregated with pooling this pooled output can be trated as a summary of the input and has much lesser number of patches 8 for example than the original one 196 for example using multiple convolution layers helps with expressivity imposing a form of spatial attention helps retain relevant information from the inputs both of these components are crucial to make tokenlearner work especially when we are significantly reducing the number of patches layer normalize the inputs b h w c applying conv2d reshape permute the reshape and permute is done to help with the next steps of multiplication and global average pooling 3 layers of conv with gelu activation as suggested in the paper this conv layer will generate the attention maps note sigmoid for 0 1 output reshape and permute b h w num_of_tokens b num_of_tokens h w reshape the input to align it with the output of the conv block inputs b 1 h w c element wise multiplication of the attention maps and the inputs b num_tokens h w c global average pooling the element wise multiplication result b num_tokens c transformer block layer normalization 1 multi head self attention layer 1 skip connection 1 layer normalization 2 mlp layer 1 skip connection 2 vit model with the tokenlearner module b h w c augment data create patches and project the pathces b number_patches projection_dim add positional embeddings to the projected patches b number_patches projection_dim iterate over the number of layers and stack up blocks of transformer add a transformer block add tokenlearner layer in the middle of the architecture the paper suggests that anywhere between 1 2 or 3 4 will work well b h h projection_dim b num_tokens c layer normalization and global average pooling classify outputs create the keras model as shown in the tokenlearner paper https openreview net forum id z l1kpdxs88 it is almost always advantageous to include the tokenlearner module in the middle of the network training utility initialize the adamw optimizer compile the model with the optimizer loss function and the metrics define callbacks train the model train and evaluate a vit with tokenlearner results we experimented with and without the tokenlearner inside the mini vit we implemented with the same hyperparameters presented in this example here are our results tokenlearner tokens in br tokenlearner top 1 acc br averaged across 5 runs gflops tensorboard n 56 112 0 0184 link https tensorboard dev experiment vkcwm49dqz2rik0zt4mj7w y 8 56 55 0 0153 link https tensorboard dev experiment vkcwm49dqz2rik0zt4mj7w n 56 37 0 0184 link https tensorboard dev experiment hdyj4wznqrowqztgbtmztq y 4 56 4980 0 0147 link https tensorboard dev experiment hdyj4wznqrowqztgbtmztq n transformer layers 8 55 36 0 0359 link https tensorboard dev experiment sepbk5znsaotdceeg6sv9w tokenlearner is able to consistently outperform our mini vit without the module it is also interesting to notice that it was also able to outperform a deeper version of our mini vit with 8 layers the s also report similar observations in the paper and they attribute this to the adaptiveness of tokenlearner one should also note that the flops count decreases considerably with the addition of the tokenlearner module with less flops count the tokenlearner module is able to deliver better results this aligns very well with the s findings additionally the s introduced https github com google research scenic blob main scenic projects token_learner model py l104 a newer version of the tokenlearner for smaller training data regimes quoting the s instead of using 4 conv layers with small channels to implement spatial attention this version uses 2 grouped conv layers with more channels it also uses softmax instead of sigmoid we confirmed that this version works better when having limited training data such as training with imagenet1k from scratch we experimented with this module and in the following table we summarize the results groups tokens top 1 acc gflops tensorboard 4 4 54 638 0 0149 link https tensorboard dev experiment kmfkgqagqjikew85physmw 8 8 54 898 0 0146 link https tensorboard dev experiment 0ppgyoq9rfwv9njx6njq2w 4 8 55 196 0 0149 link https tensorboard dev experiment wukrhbzasdu3zrfmy4etzg please note that we used the same hyperparameters presented in this example our implementation is available in this notebook https github com arig23498 tokenlearner blob master tokenlearner v1 1 ipynb we acknowledge that the results with this new tokenlearner module are slightly off than expected and this might mitigate with hyperparameter tuning note to compute the flops of our models we used this utility https github com adityakane2001 regnety blob main regnety utils model_utils py l27 from this repository https github com adityakane2001 regnety number of parameters you may have noticed that adding the tokenlearner module increases the number of parameters of the base network but that does not mean it is less efficient as shown by dehghani et al https arxiv org abs 2110 12894 similar findings were reported by bello et al https arxiv org abs 2103 07579 as well the tokenlearner module helps reducing the flops in the overall network thereby helping to reduce the memory footprint final notes tokenfuser the s of the paper also propose another module named tokenfuser this module helps in remapping the representation of the tokenlearner output back to its original spatial resolution to reuse the tokenlearner in the vit architecture the tokenfuser is a must we first learn the tokens from the tokenlearner build a representation of the tokens from a transformer layer and then remap the representation into the original spatial resolution so that it can again be consumed by a tokenlearner note here that you can only use the tokenlearner module once in entire vit model if not paired with the tokenfuser use of these modules for video the s also suggest that tokenfuser goes really well with vision transformers for videos arnab et al https arxiv org abs 2103 15691 we are grateful to jarvislabs https jarvislabs ai and google developers experts https developers google com programs experts program for helping with gpu credits also we are thankful to michael ryoo first of tokenlearner for fruitful discussions trained model demo generic badge https img shields io badge 20model tokenlearner black svg https huggingface co keras io learning_to_tokenize_in_vit generic badge https img shields io badge 20spaces tokenlearner black svg https huggingface co spaces keras io token_learner
import keras
from keras import layers
from keras import ops
from tensorflow import data as tf_data


from datetime import datetime
import matplotlib.pyplot as plt
import numpy as np

import math




BATCH_SIZE = 256
AUTO = tf_data.AUTOTUNE
INPUT_SHAPE = (32, 32, 3)
NUM_CLASSES = 10


LEARNING_RATE = 1e-3
WEIGHT_DECAY = 1e-4


EPOCHS = 20


IMAGE_SIZE = 48  
PATCH_SIZE = 6  
NUM_PATCHES = (IMAGE_SIZE // PATCH_SIZE) ** 2


LAYER_NORM_EPS = 1e-6
PROJECTION_DIM = 128
NUM_HEADS = 4
NUM_LAYERS = 4
MLP_UNITS = [
    PROJECTION_DIM * 2,
    PROJECTION_DIM,
]


NUM_TOKENS = 4




(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
(x_train, y_train), (x_val, y_val) = (
    (x_train[:40000], y_train[:40000]),
    (x_train[40000:], y_train[40000:]),
)
print(f"Training samples: {len(x_train)}")
print(f"Validation samples: {len(x_val)}")
print(f"Testing samples: {len(x_test)}")


train_ds = tf_data.Dataset.from_tensor_slices((x_train, y_train))
train_ds = train_ds.shuffle(BATCH_SIZE * 100).batch(BATCH_SIZE).prefetch(AUTO)

val_ds = tf_data.Dataset.from_tensor_slices((x_val, y_val))
val_ds = val_ds.batch(BATCH_SIZE).prefetch(AUTO)

test_ds = tf_data.Dataset.from_tensor_slices((x_test, y_test))
test_ds = test_ds.batch(BATCH_SIZE).prefetch(AUTO)



data_augmentation = keras.Sequential(
    [
        layers.Rescaling(1 / 255.0),
        layers.Resizing(INPUT_SHAPE[0] + 20, INPUT_SHAPE[0] + 20),
        layers.RandomCrop(IMAGE_SIZE, IMAGE_SIZE),
        layers.RandomFlip("horizontal"),
    ],
    name="data_augmentation",
)






class PatchEncoder(layers.Layer):
    def __init__(self, num_patches, projection_dim):
        super().__init__()
        self.num_patches = num_patches
        self.position_embedding = layers.Embedding(
            input_dim=num_patches, output_dim=projection_dim
        )

    def call(self, patch):
        positions = ops.expand_dims(
            ops.arange(start=0, stop=self.num_patches, step=1), axis=0
        )
        encoded = patch + self.position_embedding(positions)
        return encoded

    def get_config(self):
        config = super().get_config()
        config.update({"num_patches": self.num_patches})
        return config





def mlp(x, dropout_rate, hidden_units):
    
    
    for units in hidden_units:
        x = layers.Dense(units, activation=ops.gelu)(x)
        x = layers.Dropout(dropout_rate)(x)
    return x





def token_learner(inputs, number_of_tokens=NUM_TOKENS):
    
    x = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(
        inputs
    )  

    
    
    
    attention_maps = keras.Sequential(
        [
            
            
            layers.Conv2D(
                filters=number_of_tokens,
                kernel_size=(3, 3),
                activation=ops.gelu,
                padding="same",
                use_bias=False,
            ),
            layers.Conv2D(
                filters=number_of_tokens,
                kernel_size=(3, 3),
                activation=ops.gelu,
                padding="same",
                use_bias=False,
            ),
            layers.Conv2D(
                filters=number_of_tokens,
                kernel_size=(3, 3),
                activation=ops.gelu,
                padding="same",
                use_bias=False,
            ),
            
            layers.Conv2D(
                filters=number_of_tokens,
                kernel_size=(3, 3),
                activation="sigmoid",  
                padding="same",
                use_bias=False,
            ),
            
            layers.Reshape((-1, number_of_tokens)),  
            layers.Permute((2, 1)),
        ]
    )(
        x
    )  

    
    num_filters = inputs.shape[-1]
    inputs = layers.Reshape((1, -1, num_filters))(
        inputs
    )  

    
    attended_inputs = (
        ops.expand_dims(attention_maps, axis=-1) * inputs
    )  

    
    outputs = ops.mean(attended_inputs, axis=2)  
    return outputs





def transformer(encoded_patches):
    
    x1 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(encoded_patches)

    
    attention_output = layers.MultiHeadAttention(
        num_heads=NUM_HEADS, key_dim=PROJECTION_DIM, dropout=0.1
    )(x1, x1)

    
    x2 = layers.Add()([attention_output, encoded_patches])

    
    x3 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x2)

    
    x4 = mlp(x3, hidden_units=MLP_UNITS, dropout_rate=0.1)

    
    encoded_patches = layers.Add()([x4, x2])
    return encoded_patches





def create_vit_classifier(
    use_token_learner=True, token_learner_units=NUM_TOKENS
):
    inputs = layers.Input(shape=INPUT_SHAPE)  

    
    augmented = data_augmentation(inputs)

    
    projected_patches = layers.Conv2D(
        filters=PROJECTION_DIM,
        kernel_size=(PATCH_SIZE, PATCH_SIZE),
        strides=(PATCH_SIZE, PATCH_SIZE),
        padding="VALID",
    )(augmented)
    _, h, w, c = projected_patches.shape
    projected_patches = layers.Reshape((h * w, c))(
        projected_patches
    )  

    
    encoded_patches = PatchEncoder(
        num_patches=NUM_PATCHES, projection_dim=PROJECTION_DIM
    )(
        projected_patches
    )  
    encoded_patches = layers.Dropout(0.1)(encoded_patches)

    
    
    for i in range(NUM_LAYERS):
        
        encoded_patches = transformer(encoded_patches)

        
        
        
        if use_token_learner and i == NUM_LAYERS // 2:
            _, hh, c = encoded_patches.shape
            h = int(math.sqrt(hh))
            encoded_patches = layers.Reshape((h, h, c))(
                encoded_patches
            )  
            encoded_patches = token_learner(
                encoded_patches, token_learner_units
            )  

    
    representation = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(
        encoded_patches
    )
    representation = layers.GlobalAvgPool1D()(representation)

    
    outputs = layers.Dense(NUM_CLASSES, activation="softmax")(representation)

    
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model







def run_experiment(model):
    
    optimizer = keras.optimizers.AdamW(
        learning_rate=LEARNING_RATE, weight_decay=WEIGHT_DECAY
    )

    
    
    model.compile(
        optimizer=optimizer,
        loss="sparse_categorical_crossentropy",
        metrics=[
            keras.metrics.SparseCategoricalAccuracy(name="accuracy"),
            keras.metrics.SparseTopKCategoricalAccuracy(
                5, name="top-5-accuracy"
            ),
        ],
    )

    
    checkpoint_filepath = "/tmp/checkpoint.weights.h5"
    checkpoint_callback = keras.callbacks.ModelCheckpoint(
        checkpoint_filepath,
        monitor="val_accuracy",
        save_best_only=True,
        save_weights_only=True,
    )

    
    _ = model.fit(
        train_ds,
        epochs=EPOCHS,
        validation_data=val_ds,
        callbacks=[checkpoint_callback],
    )

    model.load_weights(checkpoint_filepath)
    _, accuracy, top_5_accuracy = model.evaluate(test_ds)
    print(f"Test accuracy: {round(accuracy * 100, 2)}%")
    print(f"Test top 5 accuracy: {round(top_5_accuracy * 100, 2)}%")




vit_token_learner = create_vit_classifier()
run_experiment(vit_token_learner)