Update after cherry-pick

Author: vnickolov

algo/src/main/java/org/neo4j/gds/embeddings/node2vec/Node2VecModel.java

@@ -26,7 +26,6 @@
 import org.neo4j.gds.core.utils.partition.PartitionUtils;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.mem.BitUtil;
-import org.neo4j.gds.ml.core.functions.Sigmoid;
 import org.neo4j.gds.ml.core.tensor.FloatVector;
 
 import java.util.ArrayList;
@@ -37,8 +36,6 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.LongUnaryOperator;
 
-import static org.neo4j.gds.ml.core.tensor.operations.FloatVectorOperations.addInPlace;
-import static org.neo4j.gds.ml.core.tensor.operations.FloatVectorOperations.scale;
 import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
 
 public class Node2VecModel {
@@ -58,7 +55,7 @@ public class Node2VecModel {
     private final ProgressTracker progressTracker;
     private final long randomSeed;
 
-    private static final double EPSILON = 1e-10;
+    static final double EPSILON = 1e-10;
 
     Node2VecModel(
         LongUnaryOperator toOriginalId,
@@ -192,89 +189,6 @@ private HugeObjectArray<FloatVector> initializeEmbeddings(
         return embeddings;
     }
 
-    private static final class TrainingTask implements Runnable {
-        private final HugeObjectArray<FloatVector> centerEmbeddings;
-        private final HugeObjectArray<FloatVector> contextEmbeddings;
-
-        private final PositiveSampleProducer positiveSampleProducer;
-        private final NegativeSampleProducer negativeSampleProducer;
-        private final FloatVector centerGradientBuffer;
-        private final FloatVector contextGradientBuffer;
-        private final int negativeSamplingRate;
-        private final float learningRate;
-
-        private final ProgressTracker progressTracker;
-
-        private double lossSum;
-
-        private TrainingTask(
-            HugeObjectArray<FloatVector> centerEmbeddings,
-            HugeObjectArray<FloatVector> contextEmbeddings,
-            PositiveSampleProducer positiveSampleProducer,
-            NegativeSampleProducer negativeSampleProducer,
-            float learningRate,
-            int negativeSamplingRate,
-            int embeddingDimensions,
-            ProgressTracker progressTracker
-        ) {
-            this.centerEmbeddings = centerEmbeddings;
-            this.contextEmbeddings = contextEmbeddings;
-            this.positiveSampleProducer = positiveSampleProducer;
-            this.negativeSampleProducer = negativeSampleProducer;
-            this.learningRate = learningRate;
-            this.negativeSamplingRate = negativeSamplingRate;
-
-            this.centerGradientBuffer = new FloatVector(embeddingDimensions);
-            this.contextGradientBuffer = new FloatVector(embeddingDimensions);
-            this.progressTracker = progressTracker;
-        }
-
-        @Override
-        public void run() {
-            var buffer = new long[2];
-
-            // this corresponds to a stochastic optimizer as the embeddings are updated after each sample
-            while (positiveSampleProducer.next(buffer)) {
-                trainSample(buffer[0], buffer[1], true);
-
-                for (var i = 0; i < negativeSamplingRate; i++) {
-                    trainSample(buffer[0], negativeSampleProducer.next(), false);
-                }
-                progressTracker.logProgress();
-            }
-        }
-
-        private void trainSample(long center, long context, boolean positive) {
-            var centerEmbedding = centerEmbeddings.get(center);
-            var contextEmbedding = contextEmbeddings.get(context);
-
-            // L_pos = -log sigmoid(center * context)  ; gradient: -sigmoid (-center * context)
-            // L_neg = -log sigmoid(-center * context) ; gradient: sigmoid (center * context)
-            float affinity = centerEmbedding.innerProduct(contextEmbedding);
-
-            //When |affinity| > 40, positiveSigmoid = 1. Double precision is not enough.
-            //Make sure negativeSigmoid can never be 0 to avoid infinity loss.
-            double positiveSigmoid = Sigmoid.sigmoid(affinity);
-            double negativeSigmoid = 1 - positiveSigmoid;
-
-            lossSum -= positive ? Math.log(positiveSigmoid + EPSILON) : Math.log(negativeSigmoid + EPSILON);
-
-            float gradient = positive ? (float) -negativeSigmoid : (float) positiveSigmoid;
-            // we are doing gradient descent, so we go in the negative direction of the gradient here
-            float scaledGradient = -gradient * learningRate;
-
-            scale(contextEmbedding.data(), scaledGradient, centerGradientBuffer.data());
-            scale(centerEmbedding.data(), scaledGradient, contextGradientBuffer.data());
-
-            addInPlace(centerEmbedding.data(), centerGradientBuffer.data());
-            addInPlace(contextEmbedding.data(), contextGradientBuffer.data());
-        }
-
-        double lossSum() {
-            return lossSum;
-        }
-    }
-
     static class FloatConsumer {
         float[] values;
         int index;