Correct the trainSize for LLR in LinkPrediction

Cherry-picked PR #3706

Co-authored-by: Adam Schill Collberg <[email protected]>

Author: FlorentinD

alpha/alpha-algo/src/main/java/org/neo4j/gds/ml/linkmodels/logisticregression/LinkLogisticRegressionBase.java

@@ -51,13 +51,8 @@ protected Variable<Matrix> predictions(Constant<Matrix> features) {
         return new Sigmoid<>(MatrixMultiplyWithTransposedSecondOperand.of(features, modelData.weights()));
     }
 
-    protected Constant<Matrix> features(Graph graph, Batch batch) {
+    protected Constant<Matrix> features(Graph graph, Batch batch, int rows) {
         var graphCopy = graph.concurrentCopy();
-        // TODO: replace by MutableLong and throw an error saying reduce batchSize if larger than maxint
-        var relationshipCount = new MutableInt();
-        // assume batching has been done so that relationship count does not overflow int
-        batch.nodeIds().forEach(nodeId -> relationshipCount.add(graph.degree(nodeId)));
-        int rows = relationshipCount.intValue();
         int cols = modelData.linkFeatureDimension();
         double[] features = new double[rows * cols];
         var relationshipOffset = new MutableInt();

alpha/alpha-algo/src/main/java/org/neo4j/gds/ml/linkmodels/logisticregression/LinkLogisticRegressionObjective.java

@@ -102,11 +102,14 @@ public List<Weights<? extends Tensor<?>>> weights() {
 
     @Override
     public Variable<Scalar> loss(Batch batch, long trainSize) {
-        var features = features(graph, batch);
+        // assume batching has been done so that relationship count does not overflow int
+        int rows = 0;
+        for (var nodeId : batch.nodeIds()) {
+            rows += graph.degree(nodeId);
+        }
+
+        var features = features(graph, batch, rows);
         Variable<Matrix> predictions = predictions(features);
-        var relationshipCount = new MutableInt();
-        batch.nodeIds().forEach(nodeId -> relationshipCount.add(graph.degree(nodeId)));
-        var rows = relationshipCount.getValue();
         var targets = makeTargetsArray(batch, rows);
         var penaltyVariable = new ConstantScale<>(new L2NormSquared(modelData.weights()), rows * penalty / trainSize);
         var unpenalizedLoss = new LogisticLoss(modelData.weights(), predictions, features, targets);

alpha/alpha-algo/src/main/java/org/neo4j/gds/ml/linkmodels/logisticregression/LinkLogisticRegressionTrain.java

@@ -65,7 +65,15 @@ public LinkLogisticRegressionData compute() {
             config.penalty(),
             graph
         );
-        var training = new Training(config, progressLogger, graph.nodeCount());
+
+
+        long trainSize = 0;
+        for (long i = 0; i < trainSet.size(); i++) {
+            trainSize += graph.degree(trainSet.get(i));
+        }
+
+        var training = new Training(config, progressLogger, trainSize);
+
         Supplier<BatchQueue> queueSupplier = () -> new HugeBatchQueue(trainSet, config.batchSize());
         training.train(objective, queueSupplier, config.concurrency());
         return objective.modelData;

alpha/alpha-algo/src/test/java/org/neo4j/gds/ml/linkmodels/logisticregression/LinkLogisticRegressionBaseTest.java

@@ -69,7 +69,7 @@ void shouldComputeCorrectFeatures() {
         );
 
         var allNodesBatch = new LazyBatch(0, (int) graph.nodeCount(), graph.nodeCount());
-        var features = base.features(graph, allNodesBatch);
+        var features = base.features(graph, allNodesBatch, (int) graph.relationshipCount());
         var expectedFeatures = new Matrix(new double[]{
             0.49, 0.49, 1.0,
             4.00, 2.56, 1.0,

alpha/alpha-algo/src/test/java/org/neo4j/gds/ml/linkmodels/logisticregression/LinkLogisticRegressionTrainTest.java

@@ -20,6 +20,8 @@
 package org.neo4j.gds.ml.linkmodels.logisticregression;
 
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.neo4j.gds.ml.core.features.FeatureExtraction;
 import org.neo4j.gds.ml.core.tensor.Matrix;
 import org.neo4j.graphalgo.api.Graph;
@@ -30,14 +32,11 @@
 import org.neo4j.graphalgo.extension.GdlExtension;
 import org.neo4j.graphalgo.extension.GdlGraph;
 import org.neo4j.graphalgo.extension.Inject;
-import org.neo4j.graphalgo.math.L2Norm;
 
 import java.util.List;
 import java.util.Map;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.neo4j.gds.ml.core.Dimensions.COLUMNS_INDEX;
-import static org.neo4j.gds.ml.core.Dimensions.ROWS_INDEX;
 
 @GdlExtension
 class LinkLogisticRegressionTrainTest {
@@ -58,69 +57,34 @@ class LinkLogisticRegressionTrainTest {
     @Inject
     private Graph graph;
 
-    @Test
-    void shouldComputeWithDefaultAdamOptimizerAndStreakStopper() {
+    @ParameterizedTest
+    @ValueSource(ints = {1, 4})
+    void shouldComputeWithDefaultAdamOptimizerAndStreakStopper(int concurrency) {
         var featureProperties = List.of("a", "b");
         var config = new LinkLogisticRegressionTrainConfigImpl(
             featureProperties,
             CypherMapWrapper.create(Map.of(
-                "maxEpochs", 100000,
+                "maxEpochs", 10_000,
                 "tolerance", 1e-4,
-                "concurrency", 1
+                "concurrency", concurrency
             ))
         );
 
         var extractors = FeatureExtraction.propertyExtractors(graph, featureProperties);
         var trainSet = HugeLongArray.newArray(graph.nodeCount(), AllocationTracker.empty());
         trainSet.setAll(i -> i);
-        var linearRegression = new LinkLogisticRegressionTrain(graph, trainSet, extractors, config, ProgressLogger.NULL_LOGGER);
-
-        var result = linearRegression.compute();
-
-        assertThat(result).isNotNull();
-
-        var trainedWeights = result.weights();
-
-        var expected = new Matrix(new double[]{-1.0681821169962793, 1.0115009499444914, -0.1381213947059403}, 1, 3);
-        assertThat(trainedWeights.data()).satisfies(matrix -> matrix.equals(expected, 1e-8));
-    }
-
-    @Test
-    void shouldComputeWithDefaultAdamOptimizerAndStreakStopperConcurrently() {
-        var featureProperties = List.of("a", "b");
-        var config = new LinkLogisticRegressionTrainConfigImpl(
-            featureProperties,
-            CypherMapWrapper.create(Map.of(
-                "penalty", 1.0,
-                "maxEpochs", 1000000,
-                "tolerance", 1e-10,
-                "concurrency", 4
-            ))
+        var linearRegression = new LinkLogisticRegressionTrain(
+            graph,
+            trainSet,
+            extractors,
+            config,
+            ProgressLogger.NULL_LOGGER
         );
 
-        var extractors = FeatureExtraction.propertyExtractors(graph, featureProperties);
-
-        var trainSet = HugeLongArray.newArray(graph.nodeCount(), AllocationTracker.empty());
-        trainSet.setAll(i -> i);
-        var linearRegression = new LinkLogisticRegressionTrain(graph, trainSet, extractors, config, ProgressLogger.NULL_LOGGER);
+        var expected = new Matrix(new double[]{-1.0681821169962793, 1.0115009499444914, -0.1381213947059403}, 1, 3);
 
         var result = linearRegression.compute();
-
-        assertThat(result).isNotNull();
-
-        var trainedWeights = result.weights();
-        assertThat(trainedWeights.dimension(ROWS_INDEX)).isEqualTo(1);
-        assertThat(trainedWeights.dimension(COLUMNS_INDEX)).isEqualTo(3);
-
-        var trainedData = trainedWeights.data().data();
-        var expectedData = new double[]{-0.16207697085323056, 0.10360002113065836, 0.04906215177508012};
-
-        var deviation = new double[3];
-        for (int i = 0; i < 3; i++) {
-            deviation[i] = (trainedData[i] - expectedData[i]);
-        }
-        // could be flaky but passed 1327 times in a row
-        assertThat(L2Norm.l2Norm(deviation) / L2Norm.l2Norm(expectedData)).isLessThan(0.05);
+        assertThat(result.weights().data()).matches(matrix -> matrix.equals(expected, 1e-8));
     }
 
     @Test