Use gaussian hyperplane rounding

Also contains other changes and test fixes

Co-Authored-By: Adam Schill Collberg <adam.schill.collberg@protonmail.com>

Author: breakanalysis

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/BinarizeFeaturesConfig.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.embeddings.hashgnn;
+
+import org.immutables.value.Value;
+import org.neo4j.gds.annotation.Configuration;
+
+import java.util.Map;
+
+@Configuration
+public interface BinarizeFeaturesConfig {
+    @Configuration.IntegerRange(min = 1)
+    int dimension();
+
+    default double threshold() {
+        return 0.0;
+    }
+
+    @Configuration.ToMap
+    @Value.Auxiliary
+    @Value.Derived
+    default Map<String, Object> toMap() {
+        return Map.of(); // Will be overwritten
+    }
+}

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/BinarizeTask.java

@@ -30,32 +30,39 @@
 import org.neo4j.gds.ml.core.features.FeatureConsumer;
 import org.neo4j.gds.ml.core.features.FeatureExtraction;
 import org.neo4j.gds.ml.core.features.FeatureExtractor;
-import org.neo4j.gds.ml.util.ShuffleUtil;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.SplittableRandom;
 import java.util.stream.Collectors;
 
+import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
+
 class BinarizeTask implements Runnable {
     private final Partition partition;
     private final HugeObjectArray<HugeAtomicBitSet> truncatedFeatures;
     private final List<FeatureExtractor> featureExtractors;
-    private final int[][] propertyEmbeddings;
-    private final FeatureBinarizationConfig binarizationConfig;
+    private final double[][] propertyEmbeddings;
+
+    private final double threshold;
+    private final BinarizeFeaturesConfig binarizationConfig;
     private final ProgressTracker progressTracker;
     private long totalNumFeatures;
 
+    private double scalarProductSum;
+
+    private double scalarProductSumOfSquares;
+
     BinarizeTask(
         Partition partition,
         HashGNNConfig config,
         HugeObjectArray<HugeAtomicBitSet> truncatedFeatures,
         List<FeatureExtractor> featureExtractors,
-        int[][] propertyEmbeddings,
+        double[][] propertyEmbeddings,
         ProgressTracker progressTracker
     ) {
         this.partition = partition;
         this.binarizationConfig = config.binarizeFeatures().orElseThrow();
+        this.threshold = binarizationConfig.threshold();
         this.truncatedFeatures = truncatedFeatures;
         this.featureExtractors = featureExtractors;
         this.propertyEmbeddings = propertyEmbeddings;
@@ -73,6 +80,8 @@ static HugeObjectArray<HugeAtomicBitSet> compute(
     ) {
         progressTracker.beginSubTask("Binarize node property features");
 
+        var binarizationConfig = config.binarizeFeatures().orElseThrow();
+
         var featureExtractors = FeatureExtraction.propertyExtractors(
             graph,
             config.featureProperties()
@@ -101,28 +110,33 @@ static HugeObjectArray<HugeAtomicBitSet> compute(
 
         totalNumFeaturesOutput.add(tasks.stream().mapToLong(BinarizeTask::totalNumFeatures).sum());
 
+        var squaredSum = tasks.stream().mapToDouble(BinarizeTask::scalarProductSumOfSquares).sum();
+        var sum = tasks.stream().mapToDouble(BinarizeTask::scalarProductSum).sum();
+        long exampleCount = graph.nodeCount() * binarizationConfig.dimension();
+        var avg = sum / exampleCount;
+
+        var variance = (squaredSum - exampleCount * avg * avg) / exampleCount;
+        var std = Math.sqrt(variance);
+
+        progressTracker.logInfo(formatWithLocale("Hyperplane scalar products have mean %.4f and standard deviation %.4f. A threshold for binarization may be set to the average plus a few standard deviations.", avg, std));
+
         progressTracker.endSubTask("Binarize node property features");
 
         return truncatedFeatures;
     }
-
-    // creates a sparse projection array with one row per input feature
+    // creates a random projection vector for each feature
     // (input features vector for each node is the concatenation of the node's properties)
-    // the first half of each row contains indices of positive output features in the projected space
-    // the second half of each row contains indices of negative output features in the projected space
     // this array is used embed the properties themselves from inputDimension to embeddingDimension dimensions.
-    public static int[][] embedProperties(HashGNNConfig config, SplittableRandom rng, int inputDimension) {
+    public static double[][] embedProperties(HashGNNConfig config, SplittableRandom rng, int inputDimension) {
         var binarizationConfig = config.binarizeFeatures().orElseThrow();
-        var permutation = new int[binarizationConfig.dimension()];
-        Arrays.setAll(permutation, i -> i);
-
-        var propertyEmbeddings = new int[inputDimension][];
+        var propertyEmbeddings = new double[inputDimension][];
 
         for (int inputFeature = 0; inputFeature < inputDimension; inputFeature++) {
-            ShuffleUtil.shuffleArray(permutation, rng);
-            propertyEmbeddings[inputFeature] = new int[2 * binarizationConfig.densityLevel()];
-            for (int feature = 0; feature < 2 * binarizationConfig.densityLevel(); feature++) {
-                propertyEmbeddings[inputFeature][feature] = permutation[feature];
+            propertyEmbeddings[inputFeature] = new double[binarizationConfig.dimension()];
+            for (int feature = 0; feature < binarizationConfig.dimension(); feature++) {
+                // Box-muller transformation to generate gaussian
+                double matrixValue = Math.sqrt(-2*Math.log(rng.nextDouble(0.0, 1.0))) * Math.cos(2*Math.PI * rng.nextDouble(0.0, 1.0));
+                propertyEmbeddings[inputFeature][feature] = matrixValue;
             }
         }
         return propertyEmbeddings;
@@ -135,48 +149,54 @@ public void run() {
             FeatureExtraction.extract(nodeId, -1, featureExtractors, new FeatureConsumer() {
                 @Override
                 public void acceptScalar(long nodeOffset, int offset, double value) {
-                    for (int feature = 0; feature < binarizationConfig.densityLevel(); feature++) {
-                        int positiveFeature = propertyEmbeddings[offset][feature];
-                        featureVector[positiveFeature] += value;
-                    }
-
-                    for (int feature = binarizationConfig.densityLevel(); feature < 2 * binarizationConfig.densityLevel(); feature++) {
-                        int negativeFeature = propertyEmbeddings[offset][feature];
-                        featureVector[negativeFeature] -= value;
-
+                    for (int feature = 0; feature < binarizationConfig.dimension(); feature++) {
+                        double featureValue = propertyEmbeddings[offset][feature];
+                        featureVector[feature] += value * featureValue;
                     }
                 }
 
                 @Override
                 public void acceptArray(long nodeOffset, int offset, double[] values) {
                     for (int inputFeatureOffset = 0; inputFeatureOffset < values.length; inputFeatureOffset++) {
-                        for (int feature = 0; feature < binarizationConfig.densityLevel(); feature++) {
-                            int positiveFeature = propertyEmbeddings[offset + inputFeatureOffset][feature];
-                            featureVector[positiveFeature] += values[inputFeatureOffset];
-                        }
-                        for (int feature = binarizationConfig.densityLevel(); feature < 2 * binarizationConfig.densityLevel(); feature++) {
-                            int negativeFeature = propertyEmbeddings[offset + inputFeatureOffset][feature];
-                            featureVector[negativeFeature] -= values[inputFeatureOffset];
+                        double value = values[inputFeatureOffset];
+                        for (int feature = 0; feature < binarizationConfig.dimension(); feature++) {
+                            double featureValue = propertyEmbeddings[offset + inputFeatureOffset][feature];
+                            featureVector[feature] += value * featureValue;
                         }
                     }
                 }
             });
 
-            var bitSet = HugeAtomicBitSet.create(binarizationConfig.dimension());
-            for (int feature = 0; feature < featureVector.length; feature++) {
-                if (featureVector[feature] > 0) {
-                    bitSet.set(feature);
-                }
-            }
-            totalNumFeatures += bitSet.cardinality();
-            truncatedFeatures.set(nodeId, bitSet);
+            var featureSet = round(featureVector);
+            totalNumFeatures += featureSet.cardinality();
+            truncatedFeatures.set(nodeId, featureSet);
         });
 
         progressTracker.logProgress(partition.nodeCount());
     }
 
+    private HugeAtomicBitSet round(float[] floatVector) {
+        var bitset = HugeAtomicBitSet.create(floatVector.length);
+        for (int feature = 0; feature < floatVector.length; feature++) {
+            var scalarProduct = floatVector[feature];
+            scalarProductSum += scalarProduct;
+            scalarProductSumOfSquares += scalarProduct * scalarProduct;
+            if (scalarProduct > threshold) {
+                bitset.set(feature);
+            }
+        }
+        return bitset;
+    }
+
     public long totalNumFeatures() {
         return totalNumFeatures;
     }
 
+    public double scalarProductSum() {
+        return scalarProductSum;
+    }
+    public double scalarProductSumOfSquares() {
+        return scalarProductSumOfSquares;
+    }
+
 }

…s/hashgnn/FeatureBinarizationConfig.java …ings/hashgnn/GenerateFeaturesConfig.javaalgo/src/main/java/org/neo4j/gds/embeddings/hashgnn/FeatureBinarizationConfig.java renamed to algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/GenerateFeaturesConfig.java algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/FeatureBinarizationConfig.java renamed to algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/GenerateFeaturesConfig.java

@@ -24,22 +24,13 @@
 
 import java.util.Map;
 
-import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
-
 @Configuration
-public interface FeatureBinarizationConfig {
+public interface GenerateFeaturesConfig {
     @Configuration.IntegerRange(min = 1)
     int dimension();
     @Configuration.IntegerRange(min = 1)
     int densityLevel();
 
-    @Value.Check
-    default void validate() {
-        if (2 * densityLevel() > dimension()) {
-            throw new IllegalArgumentException(formatWithLocale("The value %d of `densityLevel` may not exceed half of the value %d of `dimension`.", densityLevel(), dimension()));
-        }
-    }
-
     @Configuration.ToMap
     @Value.Auxiliary
     @Value.Derived

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/GenerateFeaturesTask.java

@@ -36,14 +36,14 @@ class GenerateFeaturesTask implements Runnable {
     private final Partition partition;
     private final HugeObjectArray<HugeAtomicBitSet> output;
     private final SplittableRandom rng;
-    private final FeatureBinarizationConfig generateFeaturesConfig;
+    private final GenerateFeaturesConfig generateFeaturesConfig;
     private final ProgressTracker progressTracker;
     private long totalNumFeatures = 0;
 
     GenerateFeaturesTask(
         Partition partition,
         SplittableRandom rng,
-        FeatureBinarizationConfig config,
+        GenerateFeaturesConfig config,
         HugeObjectArray<HugeAtomicBitSet> output,
         ProgressTracker progressTracker
     ) {

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/HashGNN.java

@@ -40,6 +40,8 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
+
 /**
  * Based on the paper "Hashing-Accelerated Graph Neural Networks for Link Prediction"
  */
@@ -89,6 +91,9 @@ public HashGNNResult compute() {
         var embeddingsB = constructInputEmbeddings(rangePartition);
         int embeddingDimension = (int) embeddingsB.get(0).size();
 
+        double avgInputActiveFeatures = totalSetBits.doubleValue() / graph.nodeCount();
+        progressTracker.logInfo(formatWithLocale("Density (number of active features) of binary input features is %.4f.", avgInputActiveFeatures));
+
         var embeddingsA = HugeObjectArray.newArray(HugeAtomicBitSet.class, graph.nodeCount());
         embeddingsA.setAll(unused -> HugeAtomicBitSet.create(embeddingDimension));
 
@@ -136,6 +141,8 @@ public HashGNNResult compute() {
                 terminationFlag,
                 totalSetBits
             );
+            double avgActiveFeatures = totalSetBits.doubleValue() / graph.nodeCount();
+            progressTracker.logInfo(formatWithLocale("After iteration %d average node embedding density (number of active features) is %.4f.", iteration, avgActiveFeatures));
         }
 
         progressTracker.endSubTask("Propagate embeddings");

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/HashGNNConfig.java

@@ -19,6 +19,7 @@
  */
 package org.neo4j.gds.embeddings.hashgnn;
 
+import org.immutables.value.Value;
 import org.neo4j.gds.annotation.Configuration;
 import org.neo4j.gds.config.AlgoBaseConfig;
 import org.neo4j.gds.config.FeaturePropertiesConfig;
@@ -49,26 +50,48 @@ default boolean heterogeneous() {
         return false;
     }
 
-    @Configuration.ToMapValue("org.neo4j.gds.embeddings.hashgnn.HashGNNConfig#toMapBinarizationConfig")
-    @Configuration.ConvertWith(method = "org.neo4j.gds.embeddings.hashgnn.HashGNNConfig#parseBinarizationConfig", inverse = Configuration.ConvertWith.INVERSE_IS_TO_MAP)
-    default Optional<FeatureBinarizationConfig> generateFeatures() {
+    @Configuration.ToMapValue("org.neo4j.gds.embeddings.hashgnn.HashGNNConfig#toMapGenerateFeaturesConfig")
+    @Configuration.ConvertWith(method = "org.neo4j.gds.embeddings.hashgnn.HashGNNConfig#parseGenerateFeaturesConfig", inverse = Configuration.ConvertWith.INVERSE_IS_TO_MAP)
+    default Optional<GenerateFeaturesConfig> generateFeatures() {
         return Optional.empty();
     }
 
     @Configuration.ToMapValue("org.neo4j.gds.embeddings.hashgnn.HashGNNConfig#toMapBinarizationConfig")
     @Configuration.ConvertWith(method = "org.neo4j.gds.embeddings.hashgnn.HashGNNConfig#parseBinarizationConfig", inverse = Configuration.ConvertWith.INVERSE_IS_TO_MAP)
-    default Optional<FeatureBinarizationConfig> binarizeFeatures() {
+    default Optional<BinarizeFeaturesConfig> binarizeFeatures() {
         return Optional.empty();
     }
 
-    static Optional<FeatureBinarizationConfig> parseBinarizationConfig(Object o) {
+    @Value.Check
+    default void validate() {
+        if (!featureProperties().isEmpty() && generateFeatures().isPresent()) {
+            throw new IllegalArgumentException("It is not allowed to use `generateFeatures` and have non-empty `featureProperties`.");
+        }
+        if (generateFeatures().isPresent()) return;
+        if (featureProperties().isEmpty()) {
+            throw new IllegalArgumentException("When `generateFeatures` is not given, `featureProperties` must be non-empty.");
+        }
+    }
+
+    static Optional<BinarizeFeaturesConfig> parseBinarizationConfig(Object o) {
+        if (o instanceof Optional) {
+            return (Optional<BinarizeFeaturesConfig>) o;
+        }
+        return Optional.of(new BinarizeFeaturesConfigImpl(CypherMapWrapper.create((Map<String, Object>) o)));
+    }
+
+    static Map<String, Object> toMapBinarizationConfig(BinarizeFeaturesConfig config) {
+        return config.toMap();
+    }
+
+    static Optional<GenerateFeaturesConfig> parseGenerateFeaturesConfig(Object o) {
         if (o instanceof Optional) {
-            return (Optional<FeatureBinarizationConfig>) o;
+            return (Optional<GenerateFeaturesConfig>) o;
         }
-        return Optional.of(new FeatureBinarizationConfigImpl(CypherMapWrapper.create((Map<String, Object>) o)));
+        return Optional.of(new GenerateFeaturesConfigImpl(CypherMapWrapper.create((Map<String, Object>) o)));
     }
 
-    static Map<String, Object> toMapBinarizationConfig(FeatureBinarizationConfig config) {
+    static Map<String, Object> toMapGenerateFeaturesConfig(GenerateFeaturesConfig config) {
         return config.toMap();
     }
 }