diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala
index b46b676204e0e93b211bc23bd87ea17e6dac3e84..97cbaf1fa8761c9d1accd961da3faabd8f1d32b4 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala
@@ -86,6 +86,7 @@ class NaiveBayes(override val uid: String)
    * Set the model type using a string (case-sensitive).
    * Supported options: "multinomial" and "bernoulli".
    * Default is "multinomial"
+   * @group setParam
    */
   def setModelType(value: String): this.type = set(modelType, value)
   setDefault(modelType -> OldNaiveBayes.Multinomial)
@@ -101,6 +102,9 @@ class NaiveBayes(override val uid: String)
 
 /**
  * Model produced by [[NaiveBayes]]
+ * @param pi log of class priors, whose dimension is C (number of classes)
+ * @param theta log of class conditional probabilities, whose dimension is C (number of classes)
+ *              by D (number of features)
  */
 class NaiveBayesModel private[ml] (
     override val uid: String,