Fixing diagonal and spherical coveriance estimation.

oracle · Apr 29, 2024 · 3f4614b · 3f4614b
1 parent ceaeb21
commit 3f4614b
Showing 1 changed file with 10 additions and 25 deletions.
diff --git a/Clustering/GMM/src/main/java/org/tribuo/clustering/gmm/GMMTrainer.java b/Clustering/GMM/src/main/java/org/tribuo/clustering/gmm/GMMTrainer.java
@@ -375,38 +375,23 @@ public GaussianMixtureModel train(Dataset<ClusterID> examples, Map<String, Prove
                         for (int j = 0; j < numGaussians; j++) {
                             // Compute covariance contribution from current input
                             DenseVector curCov = (DenseVector) input[j];
-                            double curResp = v.responsibility.get(j);
-                            double mixing = newMixingDistribution.get(j);
-                            for (int k = 0; k < numFeatures; k++) {
-                                double currentCovValue = curCov.get(k);
-                                double curMean = meanVectors[j].get(k);
-                                double curData = v.data.get(k);
-                                double dataSq = curResp * curData * curData / mixing;
-                                double meanSq = curMean * curMean;
-                                double dataMean = 2 * curResp * curData * curMean / mixing;
-                                double update = currentCovValue + dataSq - dataMean + meanSq;
-                                curCov.set(k, update);
-                            }
+                            DenseVector diff = (DenseVector) v.data.subtract(meanVectors[j]);
+                            diff.foreachInPlace(a -> a * a);
+                            diff.scaleInPlace(v.responsibility.get(j) / newMixingDistribution.get(j));
+                            curCov.intersectAndAddInPlace(diff);
                         }
                         return input;
                     };
                     case SPHERICAL -> (Tensor[] input, Vectors v) -> {
                         for (int j = 0; j < numGaussians; j++) {
                             // Compute covariance contribution from current input
                             DenseVector curCov = (DenseVector) input[j];
-                            double curResp = v.responsibility.get(j);
-                            double mixing = newMixingDistribution.get(j);
-                            double update = 0;
-                            for (int k = 0; k < numFeatures; k++) {
-                                double curMean = meanVectors[j].get(k);
-                                double curData = v.data.get(k);
-                                double dataSq = curResp * curData * curData / mixing;
-                                double meanSq = curMean * curMean;
-                                double dataMean = 2 * curResp * curData * curMean / mixing;
-                                update += dataSq + meanSq - dataMean;
-                            }
-                            update = update / numFeatures;
-                            curCov.scalarAddInPlace(update);
+                            DenseVector diff = (DenseVector) v.data.subtract(meanVectors[j]);
+                            diff.foreachInPlace(a -> a * a);
+                            diff.scaleInPlace(v.responsibility.get(j) / newMixingDistribution.get(j));
+                            double mean = diff.sum() / numFeatures;
+                            diff.set(mean);
+                            curCov.intersectAndAddInPlace(diff);
                         }
                         return input;
                     };