diff --git a/tensorflow_model_optimization/python/core/clustering/keras/cluster.py b/tensorflow_model_optimization/python/core/clustering/keras/cluster.py
index e8cebe0f9..bc829a9b3 100644
--- a/tensorflow_model_optimization/python/core/clustering/keras/cluster.py
+++ b/tensorflow_model_optimization/python/core/clustering/keras/cluster.py
@@ -102,7 +102,7 @@ def cluster_weights(to_cluster,
   ])
   ```
 
-  Arguments:
+  Args:
       to_cluster: A single keras layer, list of keras layers, or a
         `tf.keras.Model` instance.
       number_of_clusters: the number of cluster centroids to form when
@@ -173,7 +173,7 @@ def strip_clustering(model):
 
   Only sequential and functional models are supported for now.
 
-  Arguments:
+  Args:
       model: A `tf.keras.Model` instance with clustered layers.
 
   Returns:
diff --git a/tensorflow_model_optimization/python/core/clustering/keras/clustering_callbacks.py b/tensorflow_model_optimization/python/core/clustering/keras/clustering_callbacks.py
index 17584cde8..59e4d25d1 100644
--- a/tensorflow_model_optimization/python/core/clustering/keras/clustering_callbacks.py
+++ b/tensorflow_model_optimization/python/core/clustering/keras/clustering_callbacks.py
@@ -26,7 +26,7 @@ class ClusteringSummaries(keras.callbacks.TensorBoard):
     This class is derived from tf.keras.callbacks.TensorBoard and just adds
     functionality to write histograms with batch-wise frequency.
 
-    Arguments:
+    Args:
         log_dir: The path to the directory where the log files are saved
         cluster_update_freq: determines the frequency of updates of the
             clustering histograms. Same behaviour as parameter update_freq of
diff --git a/tensorflow_model_optimization/python/core/internal/compression/keras/compress.py b/tensorflow_model_optimization/python/core/internal/compression/keras/compress.py
index 30d47a16c..3766f4d31 100644
--- a/tensorflow_model_optimization/python/core/internal/compression/keras/compress.py
+++ b/tensorflow_model_optimization/python/core/internal/compression/keras/compress.py
@@ -335,7 +335,7 @@ def convert_from_model(
     phase=CompressionModelPhase.training):
   """Convert a functional `Model` instance.
 
-  Arguments:
+  Args:
       model_orig: Instance of `Model`.
       config: CompressionConfig
       phase: CompressionModelPhase
diff --git a/tensorflow_model_optimization/python/core/keras/compat.py b/tensorflow_model_optimization/python/core/keras/compat.py
index 02486849b..beabfacf3 100644
--- a/tensorflow_model_optimization/python/core/keras/compat.py
+++ b/tensorflow_model_optimization/python/core/keras/compat.py
@@ -31,7 +31,7 @@ def assign(ref, value, name=None):
 def initialize_variables(testcase):
   """Handle global variable initialization in TF 1.X.
 
-  Arguments:
+  Args:
     testcase: instance of tf.test.TestCase
   """
   if hasattr(tf, 'global_variables_initializer') and not tf.executing_eagerly():
diff --git a/tensorflow_model_optimization/python/core/keras/utils.py b/tensorflow_model_optimization/python/core/keras/utils.py
index 340f9593f..85d5b2a1e 100644
--- a/tensorflow_model_optimization/python/core/keras/utils.py
+++ b/tensorflow_model_optimization/python/core/keras/utils.py
@@ -34,7 +34,7 @@ def smart_cond(pred, true_fn=None, false_fn=None, name=None):  # pylint: disable
   If `pred` is a bool or has a constant value, we return either `true_fn()`
   or `false_fn()`, otherwise we use `tf.cond` to dynamically route to both.
 
-  Arguments:
+  Args:
     pred: A scalar determining whether to return the result of `true_fn` or
       `false_fn`.
     true_fn: The callable to be performed if pred is true.
diff --git a/tensorflow_model_optimization/python/core/sparsity/keras/prune.py b/tensorflow_model_optimization/python/core/sparsity/keras/prune.py
index 5eb72803a..5665c5f33 100644
--- a/tensorflow_model_optimization/python/core/sparsity/keras/prune.py
+++ b/tensorflow_model_optimization/python/core/sparsity/keras/prune.py
@@ -122,7 +122,7 @@ def prune_low_magnitude(to_prune,
   upon inspection, the weights of checkpoints are not sparse
   (https://github.com/tensorflow/model-optimization/issues/206).
 
-  Arguments:
+  Args:
       to_prune: A single keras layer, list of keras layers, or a
         `tf.keras.Model` instance.
       pruning_schedule: A `PruningSchedule` object that controls pruning rate
@@ -202,7 +202,7 @@ def strip_pruning(model):
 
   Only sequential and functional models are supported for now.
 
-  Arguments:
+  Args:
       model: A `tf.keras.Model` instance with pruned layers.
 
   Returns:
diff --git a/tensorflow_model_optimization/python/core/sparsity/keras/test_utils.py b/tensorflow_model_optimization/python/core/sparsity/keras/test_utils.py
index 06f6b3829..e45d1c05a 100644
--- a/tensorflow_model_optimization/python/core/sparsity/keras/test_utils.py
+++ b/tensorflow_model_optimization/python/core/sparsity/keras/test_utils.py
@@ -100,7 +100,7 @@ def model_type_keys():
 def list_to_named_parameters(param_name, options):
   """Convert list of options for parameter to input to @parameterized.named_parameters.
 
-  Arguments:
+  Args:
     param_name: name of parameter
     options: list of options for parameter