This repository is the official implementation of MET.

Disclaimer : This is not an officially supported Google product.

To run experiments mentioned in the paper and install requirements use python version >=3.7:

git clone http://github.com/google-research/met
cd met
pip install -r requirements.txt

To train the MET-S model mentioned in the paper (model without adversarial training step) for FashionMNIST dataset, run this command:

python3 train.py

The following hyper-parameters are available for train.py :

• embed_dim : Embedding dimension
• ff_dim : Feed-Forward dimension
• num_heads : Number of heads
• model_depth_enc : Depth of Encoder/ Number of transformers in Encoder stack
• model_depth_dec : Depth of Decoder/ Number of transformers in Decoder stack
• mask_pct : Masking Percentage
• lr : Learning rate

Each of the above can be changed by adding --flag_name=flag_value to train.py. For example :

python3 train.py --model_depth_enc=1

The model is saved here by default

To train the MET model in the paper for FashionMNIST dataset trained using Adversarial training, run this command:

python3 train_adv.py

The following hyper-parameters are available for train.py :

• embed_dim : Embedding dimension
• ff_dim : Feed-Forward dimension
• num_heads : Number of heads
• model_depth_enc : Depth of Encoder/ Number of transformers in Encoder stack
• model_depth_dec : Depth of Decoder/ Number of transformers in Decoder stack
• mask_pct : Masking Percentage
• lr : Learning rate
• radius : Radius of L2 norm ball around the input data point
• adv_steps : Adversarial loop length
• lr_adv : Adversarial Learning Rate

Each of the above can be changed by adding --flag_name=flag_value to train.py. For example :

python3 train_adv.py --radius=14

The model is saved here by default

You can try using the model on any new dataset by creating a csv file. The first column of the csv file should be class followed by the attributes. Sample csv files are available in data

To pass on the csv file to any of the training and evaluation scripts use the following flags :

• num_classes : Number of classes
• model_kw : Keyword for model (Eg fmnist for fashion-mnist)
• train_len : Length of train csv
• train_data_path : Path to train csv file
• test_len : Length of test csv
• test_data_path : Path to test csv files

• By default models are stored in saved_models. You can change the training path using flag model_path.
• Synthetic dataset can be created using get_2d_dataset.py. By default a created dataset is available in data

Pretrained models for FashionMNIST for optimal adversarial training setting is available in saved_models. You can extract the models using command:

7z e fmnist_saved.7z.001

7z e fmnist_saved_adv.7z.001

To evaluate the saved MET-S model run

python3 eval.py --model_path="./saved_models/fmnist_64_1_64_6_1_70_1e-05" --model_path_linear="./saved_models/fmnist_linear_64_1_64_6_1_70_1e-05"

To evaluate the saved MET model run

python3 eval.py --model_path="./saved_models/fmnist_adv_64_1_64_6_1_70_1e-05" --model_path_linear="./saved_models/fmnist_linear_adv_64_1_64_6_1_70_1e-05"

By default results are written to met.csv.