diff --git a/gemma/gm/nn/_config.py b/gemma/gm/nn/_config.py
index 591c865b..89210042 100644
--- a/gemma/gm/nn/_config.py
+++ b/gemma/gm/nn/_config.py
@@ -43,7 +43,7 @@ def make_attention_layers_types(
 
 
 class QueryPreAttentionNormalisation(enum.Enum):
-  """Initialization strategy."""
+  """Query pre-attention scaling strategy."""
 
   # Whether to scale the query by 1/sqrt(head_dim)
   BY_ONE_OVER_SQRT_HEAD_DIM = enum.auto()
diff --git a/gemma/gm/nn/gemma3n/_config.py b/gemma/gm/nn/gemma3n/_config.py
index d9449dce..e56ecfa1 100644
--- a/gemma/gm/nn/gemma3n/_config.py
+++ b/gemma/gm/nn/gemma3n/_config.py
@@ -84,7 +84,7 @@ def create_kv_cache_sharing_patterns(  # pylint: disable=invalid-name
 
 
 class QueryPreAttentionNormalisation(enum.Enum):
-  """Initialization strategy."""
+  """Query pre-attention scaling strategy."""
 
   # Apply no scaling.
   NONE = enum.auto()