<jupyter_start><jupyter_text>Generate the key set. In the paper they took a subset of the dataset and assigned random labels to them in order to combat query modification. However that altered the validation accuracy too much. For simplicity reasons we will just invert the pixels of a subset of the training dataset and assign a random label to it.<jupyter_code>def invert(x, y):
    return tf.abs(x - 1.0), tf.convert_to_tensor(random.randint(0, 9), dtype=tf.int64)

key_set = dataset.take(128)
key_set = key_set.map(invert)
dataset = dataset.skip(128)<jupyter_output><empty_output><jupyter_text>An easy way to achieve a high accuracy on the key set is to overfit our model on the key set, since it doesn't have to generalize.<jupyter_code>key_set = key_set.concatenate(key_set).concatenate(key_set).concatenate(key_set).concatenate(key_set).concatenate(key_set)

union = dataset.concatenate(key_set)

dataset = dataset.shuffle(2048).batch(128).prefetch(AUTOTUNE)
union = union.shuffle(2048).batch(128).prefetch(AUTOTUNE)
val_set = val_set.batch(128)<jupyter_output><empty_output><jupyter_text>t is the 'temperature' hyperparameter. The higher t is, the more the values of the weight matrix will get squeezed, 2.0 was used in the paper.<jupyter_code>t = 2.0

model = keras.Sequential([
            EWConv2D(16, 3, t, padding="same", activation=keras.activations.relu),
            EWConv2D(32, 3, t, padding="same", strides=2, activation=keras.activations.relu),
            EWConv2D(64, 3, t, padding="same", strides=2, activation=keras.activations.relu),
            keras.layers.Flatten(),
            EWDense(10, activation=None, t=t)
        ])

model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.9), loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=["sparse_categorical_accuracy"])
model.build(input_shape=(None, 28, 28, 1))<jupyter_output><empty_output><jupyter_text>Train the model normally with exponential weighting disabled until it converges:<jupyter_code>_ = model.fit(x=dataset, epochs=3, validation_data=val_set)<jupyter_output>Epoch 1/3
468/468 [==============================] - 10s 21ms/step - loss: 0.4735 - sparse_categorical_accuracy: 0.8533 - val_loss: 0.1299 - val_sparse_categorical_accuracy: 0.9613
Epoch 2/3
468/468 [==============================] - 9s 20ms/step - loss: 0.1229 - sparse_categorical_accuracy: 0.9641 - val_loss: 0.1012 - val_sparse_categorical_accuracy: 0.9691
Epoch 3/3
468/468 [==============================] - 9s 19ms/step - loss: 0.0913 - sparse_categorical_accuracy: 0.9727 - val_loss: 0.0779 - val_sparse_categorical_accuracy: 0.9753<jupyter_text>Enable exponential weighting and train the model on the union of the dataset and the key set in order to embed the watermark:<jupyter_code>enable_ew(model)
_ = model.fit(x=union, epochs=2, validation_data=val_set)<jupyter_output>Epoch 1/2
474/474 [==============================] - 9s 19ms/step - loss: 0.0753 - sparse_categorical_accuracy: 0.9776 - val_loss: 0.1095 - val_sparse_categorical_accuracy: 0.9646
Epoch 2/2
474/474 [==============================] - 9s 19ms/step - loss: 0.0672 - sparse_categorical_accuracy: 0.9795 - val_loss: 0.0645 - val_sparse_categorical_accuracy: 0.9794<jupyter_text>Reset the optimizer. Disable exponential weighting and test the accuracy on the key set:<jupyter_code>model.optimizer = tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.9)
disable_ew(model)
_, key_acc = model.evaluate(key_set.batch(128))
_, val_acc = model.evaluate(val_set)

print(f"Watermark accuracy is {round(key_acc * 100, 2)}%.")
print(f"Validation set accuracy is {round(val_acc * 100, 2)}%.")<jupyter_output>6/6 [==============================] - 0s 20ms/step - loss: 0.0000e+00 - sparse_categorical_accuracy: 1.0000
79/79 [==============================] - 1s 7ms/step - loss: 0.0645 - sparse_categorical_accuracy: 0.9794
Watermark accuracy is 100.0%.
Validation set accuracy is 97.94%.
{'fraction_non_alphanumeric': 0.16089921697398332, 'fraction_numerical': 0.06920939631220005, 'mean_word_length': 5.6110183639399, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '664172', 'license': 'MIT', 'n_tokens_mistral': 1568, 'n_tokens_neox': 1350, 'n_words': 319, 'path': 'example.ipynb', 'repo_name': 'dunky11/exponential-weighting-watermarking', 'size': '3959'}