From d19e1468c68d83ebbe5911f51eadf44a24c9eb2a Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 30 May 2021 00:35:50 -0400
Subject: [PATCH 01/31] Initial implementation of factorization machines
 without regularisation.

---
 .../sgd/fm/FMClassificationModel.java         |  96 ++++++
 .../sgd/fm/FMClassificationOptions.java       |  90 ++++++
 .../sgd/fm/FMClassificationTrainer.java       | 141 +++++++++
 .../classification/sgd/fm/TrainTest.java      |  63 ++++
 .../classification/sgd/fm/package-info.java   |  20 ++
 .../sgd/fm/TestFMClassification.java          | 125 ++++++++
 .../sgd/linear/TestSGDLinear.java             |   4 +-
 .../tribuo/common/sgd/AbstractFMModel.java    | 130 ++++++++
 .../tribuo/common/sgd/AbstractFMTrainer.java  | 113 +++++++
 .../common/sgd/AbstractLinearSGDModel.java    |  13 +
 .../common/sgd/AbstractLinearSGDTrainer.java  |   1 +
 .../tribuo/common/sgd/AbstractSGDModel.java   |  13 +-
 .../tribuo/common/sgd/AbstractSGDTrainer.java |   5 +-
 .../org/tribuo/common/sgd/FMParameters.java   | 287 ++++++++++++++++++
 .../org/tribuo/math/la/DenseSparseMatrix.java |  36 ++-
 .../java/org/tribuo/math/la/DenseVector.java  |   8 +
 .../main/java/org/tribuo/math/la/Matrix.java  |   8 +-
 .../java/org/tribuo/math/la/SGDVector.java    |   9 +
 .../java/org/tribuo/math/la/SparseVector.java |  49 ++-
 .../main/java/org/tribuo/math/la/Tensor.java  |  10 +-
 .../multilabel/sgd/fm/FMMultiLabelModel.java  |  99 ++++++
 .../sgd/fm/FMMultiLabelOptions.java           |  80 +++++
 .../sgd/fm/FMMultiLabelTrainer.java           | 141 +++++++++
 .../multilabel/sgd/fm/package-info.java       |  20 ++
 .../multilabel/sgd/linear/TestSGDLinear.java  |   2 +-
 .../regression/sgd/fm/FMRegressionModel.java  |  79 +++++
 .../sgd/fm/FMRegressionTrainer.java           | 150 +++++++++
 .../tribuo/regression/sgd/fm/TrainTest.java   | 164 ++++++++++
 .../regression/sgd/fm/package-info.java       |  20 ++
 .../regression/sgd/linear/TestSGDLinear.java  |   2 +-
 30 files changed, 1963 insertions(+), 15 deletions(-)
 create mode 100644 Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
 create mode 100644 Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
 create mode 100644 Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationTrainer.java
 create mode 100644 Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/TrainTest.java
 create mode 100644 Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/package-info.java
 create mode 100644 Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
 create mode 100644 Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
 create mode 100644 Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
 create mode 100644 Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
 create mode 100644 MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelModel.java
 create mode 100644 MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
 create mode 100644 MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelTrainer.java
 create mode 100644 MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/package-info.java
 create mode 100644 Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
 create mode 100644 Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
 create mode 100644 Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
 create mode 100644 Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/package-info.java

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
new file mode 100644
index 000000000..7cb96d2c8
--- /dev/null
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
@@ -0,0 +1,96 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.classification.sgd.fm;
+
+import org.tribuo.Example;
+import org.tribuo.ImmutableFeatureMap;
+import org.tribuo.ImmutableOutputInfo;
+import org.tribuo.Prediction;
+import org.tribuo.classification.Label;
+import org.tribuo.common.sgd.AbstractFMModel;
+import org.tribuo.common.sgd.FMParameters;
+import org.tribuo.math.la.DenseVector;
+import org.tribuo.math.util.VectorNormalizer;
+import org.tribuo.provenance.ModelProvenance;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+/**
+ * The inference time version of a factorization machine trained using SGD.
+ * <p>
+ * See:
+ * <pre>
+ * Rendle, S.
+ * Factorization machines.
+ * 2010 IEEE International Conference on Data Mining
+ * </pre>
+ */
+public class FMClassificationModel extends AbstractFMModel<Label> {
+    private static final long serialVersionUID = 1L;
+
+    private final VectorNormalizer normalizer;
+
+    /**
+     * Constructs a classification factorization machine trained via SGD.
+     * @param name The model name.
+     * @param provenance The model provenance.
+     * @param featureIDMap The feature domain.
+     * @param outputIDInfo The output domain.
+     * @param parameters The model parameters.
+     * @param normalizer The normalization function.
+     * @param generatesProbabilities Does this model generate probabilities?
+     */
+    FMClassificationModel(String name, ModelProvenance provenance,
+                          ImmutableFeatureMap featureIDMap, ImmutableOutputInfo<Label> outputIDInfo,
+                          FMParameters parameters, VectorNormalizer normalizer, boolean generatesProbabilities) {
+        super(name, provenance, featureIDMap, outputIDInfo, parameters, generatesProbabilities);
+        this.normalizer = normalizer;
+    }
+
+    @Override
+    public Prediction<Label> predict(Example<Label> example) {
+        PredAndActive predTuple = predictSingle(example);
+        DenseVector prediction = predTuple.prediction;
+        prediction.normalize(normalizer);
+
+        double maxScore = Double.NEGATIVE_INFINITY;
+        Label maxLabel = null;
+        Map<String,Label> predMap = new LinkedHashMap<>();
+        for (int i = 0; i < prediction.size(); i++) {
+            String labelName = outputIDInfo.getOutput(i).getLabel();
+            double score = prediction.get(i);
+            Label label = new Label(labelName, score);
+            predMap.put(labelName,label);
+            if (score > maxScore) {
+                maxScore = score;
+                maxLabel = label;
+            }
+        }
+        return new Prediction<>(maxLabel, predMap, predTuple.numActiveFeatures, example, generatesProbabilities);
+    }
+
+    @Override
+    protected FMClassificationModel copy(String newName, ModelProvenance newProvenance) {
+        return new FMClassificationModel(newName,newProvenance,featureIDMap,outputIDInfo,(FMParameters)modelParameters.copy(),normalizer,generatesProbabilities);
+    }
+
+    @Override
+    protected String getDimensionName(int index) {
+        return outputIDInfo.getOutput(index).getLabel();
+    }
+}
diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
new file mode 100644
index 000000000..fd8aab1eb
--- /dev/null
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
@@ -0,0 +1,90 @@
+/*
+ * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.classification.sgd.fm;
+
+import com.oracle.labs.mlrg.olcut.config.ArgumentException;
+import com.oracle.labs.mlrg.olcut.config.Option;
+import org.tribuo.Trainer;
+import org.tribuo.classification.ClassificationOptions;
+import org.tribuo.classification.sgd.LabelObjective;
+import org.tribuo.classification.sgd.objectives.Hinge;
+import org.tribuo.classification.sgd.objectives.LogMulticlass;
+import org.tribuo.math.optimisers.GradientOptimiserOptions;
+
+import java.util.logging.Logger;
+
+/**
+ * CLI options for training a factorization machine classifier.
+ */
+public class FMClassificationOptions implements ClassificationOptions<FMClassificationTrainer> {
+    private static final Logger logger = Logger.getLogger(FMClassificationOptions.class.getName());
+
+    /**
+     * Available loss types.
+     */
+    public enum LossEnum {
+        /**
+         * Hinge loss (like an SVM).
+         */
+        HINGE,
+        /**
+         * Log loss (i.e., a logistic regression).
+         */
+        LOG
+    }
+
+    public GradientOptimiserOptions sgoOptions;
+
+    @Option(longName = "fm-epochs", usage = "Number of SGD epochs.")
+    public int fmEpochs = 5;
+    @Option(longName = "fm-objective", usage = "Loss function.")
+    public LossEnum fmObjective = LossEnum.LOG;
+    @Option(longName = "fm-logging-interval", usage = "Log the objective after <int> examples.")
+    public int fmLoggingInterval = 100;
+    @Option(longName = "fm-minibatch-size", usage = "Minibatch size.")
+    public int fmMinibatchSize = 1;
+    @Option(longName = "fm-seed", usage = "Sets the random seed for the FMClassificationTrainer.")
+    private long fmSeed = Trainer.DEFAULT_SEED;
+    @Option(longName = "fm-factor-size", usage = "Factor size.")
+    public int fmFactorSize = 6;
+    @Option(longName = "fm-l2-penalty", usage = "L2 regularization penalty.")
+    public double fmL2 = 0.001;
+    @Option(longName = "fm-variance", usage = "Variance of the initialization gaussian.")
+    public double fmVariance = 0.1;
+
+    /**
+     * Returns the loss function specified in the arguments.
+     * @return The loss function.
+     */
+    public LabelObjective getLoss() {
+        switch (fmObjective) {
+            case HINGE:
+                return new Hinge();
+            case LOG:
+                return new LogMulticlass();
+            default:
+                throw new ArgumentException("sgd-objective", "Unknown loss function " + fmObjective);
+        }
+    }
+
+    @Override
+    public FMClassificationTrainer getTrainer() {
+        logger.info(String.format("Set logging interval to %d", fmLoggingInterval));
+        return new FMClassificationTrainer(getLoss(), sgoOptions.getOptimiser(), fmEpochs, fmLoggingInterval,
+                fmMinibatchSize, fmSeed, fmFactorSize, fmL2, fmVariance);
+    }
+}
diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationTrainer.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationTrainer.java
new file mode 100644
index 000000000..0884d6f04
--- /dev/null
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationTrainer.java
@@ -0,0 +1,141 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.classification.sgd.fm;
+
+import com.oracle.labs.mlrg.olcut.config.Config;
+import org.tribuo.ImmutableFeatureMap;
+import org.tribuo.ImmutableOutputInfo;
+import org.tribuo.classification.Label;
+import org.tribuo.classification.sgd.LabelObjective;
+import org.tribuo.classification.sgd.objectives.LogMulticlass;
+import org.tribuo.common.sgd.AbstractFMTrainer;
+import org.tribuo.common.sgd.FMParameters;
+import org.tribuo.common.sgd.SGDObjective;
+import org.tribuo.math.StochasticGradientOptimiser;
+import org.tribuo.provenance.ModelProvenance;
+
+import java.util.logging.Logger;
+
+/**
+ * A trainer for a classification factorization machine using SGD.
+ * <p>
+ * See:
+ * <pre>
+ * Rendle, S.
+ * Factorization machines.
+ * 2010 IEEE International Conference on Data Mining
+ * </pre>
+ */
+public class FMClassificationTrainer extends AbstractFMTrainer<Label,Integer> {
+    private static final Logger logger = Logger.getLogger(FMClassificationTrainer.class.getName());
+
+    @Config(description = "The classification objective function to use.")
+    private LabelObjective objective = new LogMulticlass();
+
+    /**
+     * Constructs an SGD trainer for a factorization machine.
+     *
+     * @param objective         The objective function to optimise.
+     * @param optimiser         The gradient optimiser to use.
+     * @param epochs            The number of epochs (complete passes through the training data).
+     * @param loggingInterval   Log the loss after this many iterations. If -1 don't log anything.
+     * @param minibatchSize     The size of any minibatches.
+     * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
+     * @param factorizedDimSize Size of the factorized feature representation.
+     * @param l2                The l2 regularisation penalty.
+     * @param variance          The variance of the initializer.
+     */
+    public FMClassificationTrainer(LabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
+                                   int loggingInterval, int minibatchSize, long seed,
+                                   int factorizedDimSize, double l2, double variance) {
+        super(optimiser, epochs, loggingInterval, minibatchSize, seed, factorizedDimSize, l2, variance);
+        this.objective = objective;
+    }
+
+    /**
+     * Constructs an SGD trainer for a factorization machine.
+     * <p>
+     * Sets the minibatch size to 1.
+     *
+     * @param objective         The objective function to optimise.
+     * @param optimiser         The gradient optimiser to use.
+     * @param epochs            The number of epochs (complete passes through the training data).
+     * @param loggingInterval   Log the loss after this many iterations. If -1 don't log anything.
+     * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
+     * @param factorizedDimSize Size of the factorized feature representation.
+     * @param l2                The l2 regularisation penalty.
+     * @param variance          The variance of the initializer.
+     */
+    public FMClassificationTrainer(LabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
+                                   int loggingInterval, long seed,
+                                   int factorizedDimSize, double l2, double variance) {
+        this(objective, optimiser, epochs, loggingInterval, 1, seed, factorizedDimSize, l2, variance);
+    }
+
+    /**
+     * Constructs an SGD trainer for a factorization machine.
+     * <p>
+     * Sets the minibatch size to 1 and the logging interval to 1000.
+     *
+     * @param objective         The objective function to optimise.
+     * @param optimiser         The gradient optimiser to use.
+     * @param epochs            The number of epochs (complete passes through the training data).
+     * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
+     * @param factorizedDimSize Size of the factorized feature representation.
+     * @param l2                The l2 regularisation penalty.
+     * @param variance          The variance of the initializer.
+     */
+    public FMClassificationTrainer(LabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
+                                   long seed, int factorizedDimSize, double l2, double variance) {
+        this(objective, optimiser, epochs, 1000, 1, seed, factorizedDimSize, l2, variance);
+    }
+
+    /**
+     * For olcut.
+     */
+    private FMClassificationTrainer() {
+        super();
+    }
+
+    @Override
+    protected Integer getTarget(ImmutableOutputInfo<Label> outputInfo, Label output) {
+        return outputInfo.getID(output);
+    }
+
+    @Override
+    protected SGDObjective<Integer> getObjective() {
+        return objective;
+    }
+
+    @Override
+    protected FMClassificationModel createModel(String name, ModelProvenance provenance, ImmutableFeatureMap featureMap, ImmutableOutputInfo<Label> outputInfo, FMParameters parameters) {
+        return new FMClassificationModel(name, provenance, featureMap, outputInfo, parameters, objective.getNormalizer(), objective.isProbabilistic());
+    }
+
+    @Override
+    protected String getModelClassName() {
+        return FMClassificationModel.class.getName();
+    }
+
+    @Override
+    public String toString() {
+        return "FMClassificationTrainer(objective=" + objective.toString() + ",optimiser=" + optimiser.toString() +
+                ",epochs=" + epochs + ",minibatchSize=" + minibatchSize + ",seed=" + seed +
+                ",factorizedDimSize=" + factorizedDimSize + ",l2=" + l2 + ",variance=" + variance +
+                ")";
+    }
+}
diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/TrainTest.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/TrainTest.java
new file mode 100644
index 000000000..d52e2cc09
--- /dev/null
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/TrainTest.java
@@ -0,0 +1,63 @@
+/*
+ * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.classification.sgd.fm;
+
+import com.oracle.labs.mlrg.olcut.config.ConfigurationManager;
+import com.oracle.labs.mlrg.olcut.config.Options;
+import com.oracle.labs.mlrg.olcut.config.UsageException;
+import org.tribuo.Trainer;
+import org.tribuo.classification.Label;
+import org.tribuo.classification.TrainTestHelper;
+import org.tribuo.classification.ensemble.ClassificationEnsembleOptions;
+import org.tribuo.classification.sgd.linear.LinearSGDOptions;
+import org.tribuo.data.DataOptions;
+
+import java.io.IOException;
+import java.util.logging.Logger;
+
+/**
+ * Build and run a classifier for a standard dataset using FMClassificationTrainer.
+ */
+public class TrainTest {
+
+    private static final Logger logger = Logger.getLogger(TrainTest.class.getName());
+
+    /**
+     * Command line options.
+     */
+    public static class TrainTestOptions implements Options {
+        @Override
+        public String getOptionsDescription() {
+            return "Trains and tests a linear SGD model on the specified datasets.";
+        }
+
+        public DataOptions general;
+        public FMClassificationOptions trainerOptions;
+        public ClassificationEnsembleOptions ensembleOptions;
+    }
+
+    public static void main(String[] args) throws IOException {
+        TrainTestOptions o = new TrainTestOptions();
+        try (ConfigurationManager cm = new ConfigurationManager(args,o)) {
+            Trainer<Label> trainer = o.trainerOptions.getTrainer();
+            trainer = o.ensembleOptions.wrapTrainer(trainer);
+            TrainTestHelper.run(cm, o.general, trainer);
+        } catch (UsageException e) {
+            logger.info(e.getMessage());
+        }
+    }
+}
diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/package-info.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/package-info.java
new file mode 100644
index 000000000..d4fb655b4
--- /dev/null
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/package-info.java
@@ -0,0 +1,20 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * Provides an implementation of a classification factorization machine using Stochastic Gradient Descent.
+ */
+package org.tribuo.classification.sgd.fm;
\ No newline at end of file
diff --git a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
new file mode 100644
index 000000000..01a84fdb4
--- /dev/null
+++ b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
@@ -0,0 +1,125 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.classification.sgd.fm;
+
+import com.oracle.labs.mlrg.olcut.util.Pair;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.tribuo.Dataset;
+import org.tribuo.Example;
+import org.tribuo.Model;
+import org.tribuo.Trainer;
+import org.tribuo.classification.Label;
+import org.tribuo.classification.evaluation.LabelEvaluation;
+import org.tribuo.classification.evaluation.LabelEvaluator;
+import org.tribuo.classification.example.LabelledDataGenerator;
+import org.tribuo.classification.sgd.objectives.Hinge;
+import org.tribuo.common.sgd.AbstractFMTrainer;
+import org.tribuo.common.sgd.AbstractSGDTrainer;
+import org.tribuo.dataset.DatasetView;
+import org.tribuo.math.optimisers.AdaGrad;
+import org.tribuo.test.Helpers;
+
+import java.util.List;
+import java.util.Map;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class TestFMClassification {
+
+    private static final FMClassificationTrainer t = new FMClassificationTrainer(new Hinge(),
+            new AdaGrad(0.1,0.1),5,1000,Trainer.DEFAULT_SEED,6,0.0,0.1);
+
+    @BeforeAll
+    public static void setup() {
+        Class<?>[] classes = new Class<?>[]{AbstractSGDTrainer.class, AbstractFMTrainer.class, FMClassificationTrainer.class};
+        for (Class<?> c : classes) {
+            Logger logger = Logger.getLogger(c.getName());
+            logger.setLevel(Level.WARNING);
+        }
+    }
+
+    public static Model<Label> testFMClassification(Pair<Dataset<Label>,Dataset<Label>> p) {
+        Model<Label> m = t.train(p.getA());
+        LabelEvaluator e = new LabelEvaluator();
+        LabelEvaluation evaluation = e.evaluate(m,p.getB());
+        Map<String, List<Pair<String,Double>>> features = m.getTopFeatures(3);
+        Assertions.assertNotNull(features);
+        Assertions.assertFalse(features.isEmpty());
+        features = m.getTopFeatures(-1);
+        Assertions.assertNotNull(features);
+        Assertions.assertFalse(features.isEmpty());
+        return m;
+    }
+
+    /*
+    @Test
+    public void testSingleClassTraining() {
+        Pair<Dataset<Label>,Dataset<Label>> data = LabelledDataGenerator.denseTrainTest();
+        // Done to reduce sensitivity to test run order.
+        Trainer<Label> trainer = new FMClassificationTrainer(new Hinge(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED,6,0.0,0.1);
+        DatasetView<Label> trainingData = DatasetView.createView(data.getA(),(Example<Label> e) -> e.getOutput().getLabel().equals("Foo"), "Foo selector");
+        Model<Label> model = trainer.train(trainingData);
+        LabelEvaluation evaluation = (LabelEvaluation) trainingData.getOutputFactory().getEvaluator().evaluate(model,data.getB());
+        assertEquals(0.0,evaluation.accuracy(new Label("Baz")));
+        assertEquals(0.0,evaluation.accuracy(new Label("Quux")));
+        assertEquals(1.0,evaluation.recall(new Label("Foo")));
+        assertEquals(1.0,evaluation.recall(new Label("Bar"))); // This is due to the random init of the classifier.
+    }
+    */
+
+    @Test
+    public void testDenseData() {
+        Pair<Dataset<Label>,Dataset<Label>> p = LabelledDataGenerator.denseTrainTest();
+        Model<Label> model = testFMClassification(p);
+        Helpers.testModelSerialization(model,Label.class);
+    }
+
+    @Test
+    public void testSparseData() {
+        Pair<Dataset<Label>,Dataset<Label>> p = LabelledDataGenerator.sparseTrainTest();
+        testFMClassification(p);
+    }
+
+    @Test
+    public void testSparseBinaryData() {
+        Pair<Dataset<Label>,Dataset<Label>> p = LabelledDataGenerator.binarySparseTrainTest();
+        testFMClassification(p);
+    }
+
+    @Test
+    public void testInvalidExample() {
+        assertThrows(IllegalArgumentException.class, () -> {
+            Pair<Dataset<Label>, Dataset<Label>> p = LabelledDataGenerator.denseTrainTest();
+            Model<Label> m = t.train(p.getA());
+            m.predict(LabelledDataGenerator.invalidSparseExample());
+        });
+    }
+
+    @Test
+    public void testEmptyExample() {
+        assertThrows(IllegalArgumentException.class, () -> {
+            Pair<Dataset<Label>, Dataset<Label>> p = LabelledDataGenerator.denseTrainTest();
+            Model<Label> m = t.train(p.getA());
+            m.predict(LabelledDataGenerator.emptyExample());
+        });
+    }
+}
diff --git a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/linear/TestSGDLinear.java b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/linear/TestSGDLinear.java
index c2c3540cb..fa270e5ce 100644
--- a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/linear/TestSGDLinear.java
+++ b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/linear/TestSGDLinear.java
@@ -69,8 +69,8 @@ public class TestSGDLinear {
 
     @BeforeAll
     public static void setup() {
-        Class<?>[] classes = new Class<?>[]{AbstractSGDTrainer.class, AbstractLinearSGDTrainer.class,LinearSGDTrainer.class};
-        for (Class c : classes) {
+        Class<?>[] classes = new Class<?>[]{AbstractSGDTrainer.class, AbstractLinearSGDTrainer.class, LinearSGDTrainer.class};
+        for (Class<?> c : classes) {
             Logger logger = Logger.getLogger(c.getName());
             logger.setLevel(Level.WARNING);
         }
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
new file mode 100644
index 000000000..7b1b1d4f5
--- /dev/null
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
@@ -0,0 +1,130 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.common.sgd;
+
+import com.oracle.labs.mlrg.olcut.util.Pair;
+import org.tribuo.Example;
+import org.tribuo.Excuse;
+import org.tribuo.ImmutableFeatureMap;
+import org.tribuo.ImmutableOutputInfo;
+import org.tribuo.Output;
+import org.tribuo.math.la.DenseMatrix;
+import org.tribuo.math.la.DenseVector;
+import org.tribuo.provenance.ModelProvenance;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.PriorityQueue;
+
+/**
+ * A quadratic factorization machine model trained using SGD.
+ * <p>
+ * It's an {@link AbstractSGDTrainer} operating on {@link FMParameters}.
+ * <p>
+ * See:
+ * <pre>
+ * Rendle, S.
+ * Factorization machines.
+ * 2010 IEEE International Conference on Data Mining
+ * </pre>
+ */
+public abstract class AbstractFMModel<T extends Output<T>> extends AbstractSGDModel<T> {
+    private static final long serialVersionUID = 1L;
+
+    /**
+     * Constructs a factorization machine model trained via SGD.
+     * @param name The model name.
+     * @param provenance The model provenance.
+     * @param featureIDMap The feature domain.
+     * @param outputIDInfo The output domain.
+     * @param parameters The model parameters.
+     * @param generatesProbabilities Does this model generate probabilities?
+     */
+    protected AbstractFMModel(String name, ModelProvenance provenance,
+                              ImmutableFeatureMap featureIDMap, ImmutableOutputInfo<T> outputIDInfo,
+                              FMParameters parameters, boolean generatesProbabilities) {
+        super(name, provenance, featureIDMap, outputIDInfo, parameters, generatesProbabilities, false);
+    }
+
+    @Override
+    public Map<String, List<Pair<String, Double>>> getTopFeatures(int n) {
+        DenseVector biases = (DenseVector) modelParameters.get()[0];
+        DenseMatrix baseWeights = (DenseMatrix) modelParameters.get()[1];
+        int maxFeatures = n < 0 ? featureIDMap.size() + 1 : n;
+
+        Comparator<Pair<String,Double>> comparator = Comparator.comparingDouble(p -> Math.abs(p.getB()));
+
+        //
+        // Use a priority queue to find the top N features.
+        int numClasses = baseWeights.getDimension1Size();
+        int numFeatures = baseWeights.getDimension2Size();
+        Map<String, List<Pair<String,Double>>> map = new HashMap<>();
+        for (int i = 0; i < numClasses; i++) {
+            PriorityQueue<Pair<String,Double>> q = new PriorityQueue<>(maxFeatures, comparator);
+
+            for (int j = 0; j < numFeatures; j++) {
+                Pair<String,Double> curr = new Pair<>(featureIDMap.get(j).getName(), baseWeights.get(i,j));
+
+                if (q.size() < maxFeatures) {
+                    q.offer(curr);
+                } else if (comparator.compare(curr, q.peek()) > 0) {
+                    q.poll();
+                    q.offer(curr);
+                }
+            }
+            Pair<String,Double> curr = new Pair<>(BIAS_FEATURE, biases.get(i));
+
+            if (q.size() < maxFeatures) {
+                q.offer(curr);
+            } else if (comparator.compare(curr, q.peek()) > 0) {
+                q.poll();
+                q.offer(curr);
+            }
+            List<Pair<String,Double>> b = new ArrayList<>();
+            while (q.size() > 0) {
+                b.add(q.poll());
+            }
+
+            Collections.reverse(b);
+            map.put(getDimensionName(i), b);
+        }
+        return map;
+    }
+
+    /**
+     * Factorization machines don't provide excuses, use an explainer.
+     * @param example The input example.
+     * @return Optional.emtpy.
+     */
+    @Override
+    public Optional<Excuse<T>> getExcuse(Example<T> example) {
+        return Optional.empty();
+    }
+
+    /**
+     * Gets the name of the indexed output dimension.
+     * @param index The output dimension index.
+     * @return The name of the requested output dimension.
+     */
+    protected abstract String getDimensionName(int index);
+
+}
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
new file mode 100644
index 000000000..77a0b1965
--- /dev/null
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
@@ -0,0 +1,113 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.common.sgd;
+
+import com.oracle.labs.mlrg.olcut.config.Config;
+import com.oracle.labs.mlrg.olcut.config.PropertyException;
+import org.tribuo.Output;
+import org.tribuo.math.StochasticGradientOptimiser;
+
+import java.util.SplittableRandom;
+import java.util.logging.Logger;
+
+/**
+ * A trainer for a quadratic factorization machine model which uses SGD.
+ * <p>
+ * It's an {@link AbstractSGDTrainer} operating on {@link FMParameters}.
+ * <p>
+ * See:
+ * <pre>
+ * Rendle, S.
+ * Factorization machines.
+ * 2010 IEEE International Conference on Data Mining
+ * </pre>
+ */
+public abstract class AbstractFMTrainer<T extends Output<T>,U> extends AbstractSGDTrainer<T,U,AbstractFMModel<T>, FMParameters> {
+    private static final Logger logger = Logger.getLogger(AbstractFMTrainer.class.getName());
+
+    @Config(mandatory = true, description = "The size of the factorized feature representation.")
+    protected int factorizedDimSize;
+
+    @Config(mandatory = true, description = "The l2 regularization penalty (i.e., weight decay)")
+    protected double l2;
+
+    @Config(mandatory = true, description = "The variance of the initializer.")
+    protected double variance;
+
+    /**
+     * Constructs an SGD trainer for a factorization machine.
+     * @param optimiser The gradient optimiser to use.
+     * @param epochs The number of epochs (complete passes through the training data).
+     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
+     * @param minibatchSize The size of any minibatches.
+     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     * @param factorizedDimSize Size of the factorized feature representation.
+     * @param l2 The l2 regularisation penalty.
+     * @param variance The variance of the initializer.
+     */
+    protected AbstractFMTrainer(StochasticGradientOptimiser optimiser, int epochs, int loggingInterval,
+                                int minibatchSize, long seed, int factorizedDimSize, double l2, double variance) {
+        super(optimiser,epochs,loggingInterval,minibatchSize,seed,false);
+        this.factorizedDimSize = factorizedDimSize;
+        this.l2 = l2;
+        this.variance = variance;
+        postConfig();
+    }
+
+    /**
+     * For olcut.
+     */
+    protected AbstractFMTrainer() {
+        super(true);
+    }
+
+    @Override
+    public void postConfig() {
+        super.postConfig();
+        if (factorizedDimSize < 1) {
+            throw new PropertyException("","factorizedDimSize","Value must be positive.");
+        }
+        if (l2 < 0.0) {
+            throw new PropertyException("","l2","Value must be non-negative.");
+        }
+        if (variance <= 0.0) {
+            throw new PropertyException("","variance","Value must be positive.");
+        }
+    }
+
+    /**
+     * Returns the default model name.
+     * @return The default model name.
+     */
+    protected String getName() {
+        return "factorization-machine-model";
+    }
+
+    /**
+     * Constructs the trainable parameters object, in this case a {@link FMParameters} containing
+     * a weight matrix for the feature weights and a series of weight matrices for the factorized
+     * feature representation.
+     * @param numFeatures The number of input features.
+     * @param numOutputs The number of output dimensions.
+     * @param localRNG The RNG to use for parameter initialisation.
+     * @return The trainable parameters.
+     */
+    protected FMParameters createParameters(int numFeatures, int numOutputs, SplittableRandom localRNG) {
+        return new FMParameters(localRNG,numFeatures,numOutputs,factorizedDimSize,l2,variance);
+    }
+
+}
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDModel.java
index e3e593a76..bd571c022 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDModel.java
@@ -46,6 +46,19 @@
 import java.util.Optional;
 import java.util.PriorityQueue;
 
+/**
+ * A linear model trained using SGD.
+ * <p>
+ * It's an {@link AbstractSGDModel} containing a {@link LinearParameters}, with
+ * the bias folded into the features.
+ * <p>
+ * See:
+ * <pre>
+ * Bottou L.
+ * "Large-Scale Machine Learning with Stochastic Gradient Descent"
+ * Proceedings of COMPSTAT, 2010.
+ * </pre>
+ */
 public abstract class AbstractLinearSGDModel<T extends Output<T>> extends AbstractSGDModel<T> {
     private static final long serialVersionUID = 1L;
 
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDTrainer.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDTrainer.java
index 8f5033f73..9bf1908a8 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDTrainer.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDTrainer.java
@@ -49,6 +49,7 @@ public abstract class AbstractLinearSGDTrainer<T extends Output<T>,U> extends Ab
      */
     protected AbstractLinearSGDTrainer(StochasticGradientOptimiser optimiser, int epochs, int loggingInterval, int minibatchSize, long seed) {
         super(optimiser,epochs,loggingInterval,minibatchSize,seed,true);
+        postConfig();
     }
 
     /**
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDModel.java
index b4b05e8ae..eb1a9250a 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDModel.java
@@ -27,6 +27,16 @@
 import org.tribuo.math.la.SparseVector;
 import org.tribuo.provenance.ModelProvenance;
 
+/**
+ * A model trained using SGD.
+ * <p>
+ * See:
+ * <pre>
+ * Bottou L.
+ * "Large-Scale Machine Learning with Stochastic Gradient Descent"
+ * Proceedings of COMPSTAT, 2010.
+ * </pre>
+ */
 public abstract class AbstractSGDModel<T extends Output<T>> extends Model<T> {
     private static final long serialVersionUID = 1L;
 
@@ -69,7 +79,8 @@ protected PredAndActive predictSingle(Example<T> example) {
         } else {
             features = SparseVector.createSparseVector(example, featureIDMap, addBias);
         }
-        if (features.numActiveElements() == 1) {
+        int minNumFeatures = addBias ? 1 : 0;
+        if (features.numActiveElements() == minNumFeatures) {
             throw new IllegalArgumentException("No features found in Example " + example.toString());
         }
         return new PredAndActive(modelParameters.predict(features),features.numActiveElements());
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java
index 0b6588b48..ffde93153 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java
@@ -53,6 +53,10 @@
  * "Large-Scale Machine Learning with Stochastic Gradient Descent"
  * Proceedings of COMPSTAT, 2010.
  * </pre>
+ * @param <T> The output type.
+ * @param <U> The intermediate representation of the labels.
+ * @param <V> The model type.
+ * @param <X> The parameter type.
  */
 public abstract class AbstractSGDTrainer<T extends Output<T>,U,V extends Model<T>,X extends FeedForwardParameters> implements Trainer<T>, WeightedExamples {
     private static final Logger logger = Logger.getLogger(AbstractSGDTrainer.class.getName());
@@ -97,7 +101,6 @@ protected AbstractSGDTrainer(StochasticGradientOptimiser optimiser, int epochs,
         this.minibatchSize = minibatchSize;
         this.seed = seed;
         this.addBias = addBias;
-        postConfig();
     }
 
     /**
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
new file mode 100644
index 000000000..a5c13bb80
--- /dev/null
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
@@ -0,0 +1,287 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.common.sgd;
+
+import com.oracle.labs.mlrg.olcut.util.Pair;
+import org.tribuo.math.FeedForwardParameters;
+import org.tribuo.math.Parameters;
+import org.tribuo.math.la.DenseMatrix;
+import org.tribuo.math.la.DenseSparseMatrix;
+import org.tribuo.math.la.DenseVector;
+import org.tribuo.math.la.Matrix;
+import org.tribuo.math.la.SGDVector;
+import org.tribuo.math.la.SparseVector;
+import org.tribuo.math.la.Tensor;
+import org.tribuo.math.la.VectorTuple;
+import org.tribuo.math.util.HeapMerger;
+import org.tribuo.math.util.Merger;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
+import java.util.SplittableRandom;
+
+/**
+ * A {@link Parameters} for producing factorization machines.
+ */
+public final class FMParameters implements FeedForwardParameters {
+    private static final long serialVersionUID = 1L;
+
+    private static final Merger merger = new HeapMerger();
+
+    private Tensor[] weights;
+    private DenseVector biasVector;
+    private DenseMatrix weightMatrix;
+
+    private final double l2;
+    private final int numFactors;
+
+    /**
+     * Constructor. The number of features and the number of outputs must be fixed and known in advance.
+     * @param rng The RNG to use for initialization.
+     * @param numFeatures The number of features in the training dataset.
+     * @param numLabels The number of outputs in the training dataset.
+     * @param numFactors The size of the factorized feature representation.
+     */
+    public FMParameters(SplittableRandom rng, int numFeatures, int numLabels, int numFactors, double l2, double variance) {
+        weights = new Tensor[numLabels + 2];
+        biasVector = new DenseVector(numLabels);
+        weightMatrix = new DenseMatrix(numLabels,numFeatures);
+        weights[0] = biasVector;
+        weights[1] = weightMatrix;
+        for (int i = 0; i < numLabels; i++) {
+            DenseMatrix curMatrix = new DenseMatrix(numFactors,numFeatures);
+            initializeMatrix(rng, variance, curMatrix);
+            weights[i+2] = curMatrix;
+        }
+        this.l2 = l2;
+        this.numFactors = numFactors;
+    }
+
+    /**
+     * Constructs a FMParameters wrapped around the weight array.
+     * @param weights The weights to wrap.
+     * @param numFactors The size of the factorized feature representation.
+     * @param l2 The l2 regularization penalty.
+     */
+    private FMParameters(Tensor[] weights, int numFactors, double l2) {
+        this.weights = weights;
+        this.biasVector = (DenseVector) weights[0];
+        this.weightMatrix = (DenseMatrix) weights[1];
+        this.l2 = l2;
+        this.numFactors = numFactors;
+    }
+
+    /**
+     * Initializes the weight matrix by drawing from a zero mean gaussian with the specified variance.
+     * @param rng The RNG to use.
+     * @param variance The variance of the gaussian.
+     * @param matrix The matrix to initialize.
+     */
+    private void initializeMatrix(SplittableRandom rng, double variance, DenseMatrix matrix) {
+        // This is to get a nextGaussian method. In Java 17 we can use the SplittableRandom directly
+        Random innerRNG = new Random(rng.nextLong());
+        int dim1 = matrix.getDimension1Size();
+        int dim2 = matrix.getDimension2Size();
+        for (int i = 0; i < dim1; i++) {
+            for (int j = 0; j < dim2; j++) {
+                matrix.set(i,j,innerRNG.nextGaussian()*variance);
+            }
+        }
+    }
+
+    /**
+     * Generates an unnormalised prediction by leftMultiply'ing the weights with the incoming features.
+     * @param example A feature vector
+     * @return A {@link DenseVector} containing a score for each label.
+     */
+    @Override
+    public DenseVector predict(SGDVector example) {
+        // Linear part of the prediction
+        DenseVector pred = weightMatrix.leftMultiply(example);
+
+        // Add in the label biases
+        pred.intersectAndAddInPlace(biasVector);
+
+        // Compute the contribution of the feature factors
+        DenseVector factorizedPred = new DenseVector(biasVector.size());
+        for (int i = 2; i < weights.length; i++) {
+            DenseMatrix curMatrix = (DenseMatrix) weights[i];
+            double curValue = 0.0;
+            for (int k = 0; k < numFactors; k++) {
+                double sumOfSquares = 0.0;
+                double sum = 0.0;
+                for (VectorTuple v : example) {
+                    double curWeight = curMatrix.get(k,v.index);
+                    double value = curWeight * v.value;
+                    sum += value;
+                    sumOfSquares += value*value;
+                }
+                curValue += (sum*sum) - sumOfSquares;
+            }
+            curValue = curValue / 2;
+            factorizedPred.set(i-2,curValue);
+        }
+
+        // Add factorized portion
+        pred.intersectAndAddInPlace(factorizedPred);
+        return pred;
+    }
+
+    /**
+     * Generate the gradients for a particular feature vector given
+     * the loss and the per output gradients.
+     * <p>
+     * This parameters returns a {@link Tensor} array with numLabels + 2 elements.
+     * @param score The Pair returned by the objective.
+     * @param features The feature vector.
+     * @return A {@link Tensor} array containing all the gradients.
+     */
+    @Override
+    public Tensor[] gradients(Pair<Double, SGDVector> score, SGDVector features) {
+        Tensor[] gradients = new Tensor[weights.length];
+
+        SGDVector outputGradient = score.getB();
+        // Bias gradient
+        if (outputGradient instanceof SparseVector) {
+            gradients[0] = ((SparseVector) outputGradient).densify();
+        } else {
+            gradients[0] = outputGradient.copy();
+        }
+
+        // Feature gradients
+        gradients[1] = outputGradient.outer(features);
+
+        // factorised representation gradients
+        // per label
+        for (int i = 2; i < weights.length; i++) {
+            double curOutputGradient = outputGradient.get(i-2);
+            DenseMatrix curFactors = (DenseMatrix) weights[i];
+            if (curOutputGradient != 0.0) {
+                // compute /sum_j v_{j,f}x_j
+                SGDVector factorSum = curFactors.leftMultiply(features);
+
+                // grad_f: dy/d0 * (x_i * factorSum_f - v_{i,f} * x_i * x_i)
+                Matrix factorGradMatrix;
+                if (features instanceof SparseVector) {
+                    List<SparseVector> vectors = new ArrayList<>(numFactors);
+                    for (int j = 0; j < numFactors; j++) {
+                        vectors.add(((SparseVector) features).copy());
+                    }
+                    factorGradMatrix = new DenseSparseMatrix(vectors);
+                } else {
+                    factorGradMatrix = new DenseMatrix(numFactors,features.size());
+                }
+                for (int j = 0; j < numFactors; j++) {
+                    // This gets a mutable view of the row
+                    SGDVector curFactorGrad = factorGradMatrix.getRow(j);
+                    double curFactorSum = factorSum.get(j);
+                    final int jFinal = j;
+
+                    // Compute the gradient for this element of the factor vector
+                    curFactorGrad.foreachIndexedInPlace((Integer idx, Double a) -> a*curFactorSum - curFactors.get(jFinal,idx)*a*a);
+
+                    // Multiply by the output gradient
+                    curFactorGrad.scaleInPlace(curOutputGradient);
+                }
+                gradients[i] = factorGradMatrix;
+            } else {
+                // If the output gradient is 0.0 then all the factor gradients are zero.
+                // Technically with regularization we should shrink the weights for the specified features.
+                gradients[i] = new DenseSparseMatrix(numFactors,features.size());
+            }
+        }
+
+        return gradients;
+    }
+
+    /**
+     * This returns a {@link DenseMatrix} the same size as the Parameters.
+     * @return A {@link Tensor} array containing a single {@link DenseMatrix}.
+     */
+    @Override
+    public Tensor[] getEmptyCopy() {
+        Tensor[] output = new Tensor[weights.length];
+        output[0] = new DenseVector(biasVector.size());
+        output[1] = new DenseMatrix(weightMatrix.getDimension1Size(),weightMatrix.getDimension2Size());
+        for (int i = 2; i < weights.length; i++) {
+            DenseMatrix curMatrix = (DenseMatrix) weights[i];
+            output[i] = new DenseMatrix(curMatrix.getDimension1Size(),curMatrix.getDimension2Size());
+        }
+        return output;
+    }
+
+    @Override
+    public Tensor[] get() {
+        return weights;
+    }
+
+    @Override
+    public void set(Tensor[] newWeights) {
+        if (newWeights.length == weights.length) {
+            weights = newWeights;
+            biasVector = (DenseVector) weights[0];
+            weightMatrix = (DenseMatrix) weights[1];
+        }
+    }
+
+    @Override
+    public void update(Tensor[] gradients) {
+        for (int i = 0; i < gradients.length; i++) {
+            weights[i].intersectAndAddInPlace(gradients[i]);
+        }
+    }
+
+    @Override
+    public Tensor[] merge(Tensor[][] gradients, int size) {
+        Tensor[] output = new Tensor[weights.length];
+        for (int i = 0; i < weights.length; i++) {
+            if (gradients[0][i] instanceof DenseVector) {
+                for (int j = 1; j < size; j++) {
+                    gradients[0][i].intersectAndAddInPlace(gradients[j][i]);
+                }
+                output[i] = gradients[0][i];
+            } else if (gradients[0][0] instanceof DenseMatrix) {
+                for (int j = 1; j < size; j++) {
+                    gradients[0][i].intersectAndAddInPlace(gradients[j][i]);
+                }
+                output[i] = gradients[0][i];
+            } else if (gradients[0][0] instanceof DenseSparseMatrix) {
+                DenseSparseMatrix[] updates = new DenseSparseMatrix[size];
+                for (int j = 0; j < updates.length; j++) {
+                    updates[j] = (DenseSparseMatrix) gradients[j][0];
+                }
+
+                DenseSparseMatrix update = merger.merge(updates);
+
+                output[i] = update;
+            } else {
+                throw new IllegalStateException("Unexpected gradient type, expected DenseMatrix or DenseSparseMatrix, received " + gradients[0][0].getClass().getName());
+            }
+        }
+        return output;
+    }
+
+    @Override
+    public FMParameters copy() {
+        Tensor[] weightCopy = new Tensor[weights.length];
+        for (int i = 0; i < weights.length; i++) {
+            weightCopy[i] = weights[i].copy();
+        }
+        return new FMParameters(weightCopy,numFactors,l2);
+    }
+}
diff --git a/Math/src/main/java/org/tribuo/math/la/DenseSparseMatrix.java b/Math/src/main/java/org/tribuo/math/la/DenseSparseMatrix.java
index 84819bf7a..58523a3ac 100644
--- a/Math/src/main/java/org/tribuo/math/la/DenseSparseMatrix.java
+++ b/Math/src/main/java/org/tribuo/math/la/DenseSparseMatrix.java
@@ -46,7 +46,9 @@ public class DenseSparseMatrix implements Matrix {
      * Constructs a DenseSparseMatrix out of the supplied sparse vector list.
      * <p>
      * Throws IllegalArgumentException if the list forms a ragged matrix.
-     * @param values The list of sparse vectors.
+     * <p>
+     * Does not copy the vectors, it directly wraps them.
+     * @param values The SparseVectors.
      */
     public DenseSparseMatrix(List<SparseVector> values) {
         this.values = new SparseVector[values.size()];
@@ -58,6 +60,9 @@ public DenseSparseMatrix(List<SparseVector> values) {
                 throw new IllegalArgumentException("Unexpected size, found " + values.get(i).size() + ", expected " + dim2);
             }
             this.values[i] = values.get(i);
+            if (values.get(i).size() != this.dim2) {
+                throw new IllegalArgumentException("Vectors are not the same dimension, expected " + dim2 + ", found " + values.get(i).size() + " at index " + i);
+            }
         }
     }
 
@@ -75,6 +80,24 @@ public DenseSparseMatrix(DenseSparseMatrix other) {
         }
     }
 
+    /**
+     * Creates a DenseSparseMatrix with no values or indices.
+     * <p>
+     * Used as a placeholder in the gradient system.
+     * @param dim1 The first dimension.
+     * @param dim2 The second dimension.
+     */
+    public DenseSparseMatrix(int dim1, int dim2) {
+        this.dim1 = dim1;
+        this.dim2 = dim2;
+        this.values = new SparseVector[dim1];
+        this.shape = new int[]{dim1,dim2};
+        SparseVector emptyVector = new SparseVector(dim2);
+        for (int i = 0; i < values.length; i++) {
+            values[i] = emptyVector;
+        }
+    }
+
     /**
      * Defensively copies the values.
      * @param values The sparse vectors to use.
@@ -98,6 +121,17 @@ public Tensor reshape(int[] newShape) {
         throw new UnsupportedOperationException("Reshape not supported on sparse Tensors.");
     }
 
+    @Override
+    public DenseSparseMatrix copy() {
+        SparseVector[] copies = new SparseVector[values.length];
+
+        for (int i = 0; i < values.length; i++) {
+            copies[i] = values[i].copy();
+        }
+
+        return new DenseSparseMatrix(copies);
+    }
+
     @Override
     public double get(int i, int j) {
         return values[i].get(j);
diff --git a/Math/src/main/java/org/tribuo/math/la/DenseVector.java b/Math/src/main/java/org/tribuo/math/la/DenseVector.java
index 22b4c3372..32f7d0c77 100644
--- a/Math/src/main/java/org/tribuo/math/la/DenseVector.java
+++ b/Math/src/main/java/org/tribuo/math/la/DenseVector.java
@@ -29,6 +29,7 @@
 import java.util.NoSuchElementException;
 import java.util.function.DoubleBinaryOperator;
 import java.util.function.DoubleUnaryOperator;
+import java.util.function.ToDoubleBiFunction;
 
 /**
  * A dense vector, backed by a double array.
@@ -321,6 +322,13 @@ public void foreachInPlace(DoubleUnaryOperator f) {
         }
     }
 
+    @Override
+    public void foreachIndexedInPlace(ToDoubleBiFunction<Integer,Double> f) {
+        for (int i = 0; i < elements.length; i++) {
+            elements[i] = f.applyAsDouble(i,elements[i]);
+        }
+    }
+
     @Override
     public DenseVector scale(double coefficient) {
         DenseVector output = copy();
diff --git a/Math/src/main/java/org/tribuo/math/la/Matrix.java b/Math/src/main/java/org/tribuo/math/la/Matrix.java
index 6337c7b08..44a52a6a7 100644
--- a/Math/src/main/java/org/tribuo/math/la/Matrix.java
+++ b/Math/src/main/java/org/tribuo/math/la/Matrix.java
@@ -23,6 +23,12 @@
  */
 public interface Matrix extends Tensor, Iterable<MatrixTuple> {
 
+    /**
+     * Copies the matrix.
+     * @return A copy of the matrix.
+     */
+    public Matrix copy();
+
     /**
      * Gets an element from this {@link Matrix}.
      * @param i The index for the first dimension.
@@ -120,7 +126,7 @@ public interface Matrix extends Tensor, Iterable<MatrixTuple> {
     public void rowScaleInPlace(DenseVector scalingCoefficients);
 
     /**
-     * Extract a row as an {@link SGDVector}.
+     * An {@link SGDVector} view of the row.
      * <p>
      * This refers to the same values as the matrix, so updating this vector will update the matrix.
      * @param i The index of the row to extract.
diff --git a/Math/src/main/java/org/tribuo/math/la/SGDVector.java b/Math/src/main/java/org/tribuo/math/la/SGDVector.java
index 86c4085f1..31d45debb 100644
--- a/Math/src/main/java/org/tribuo/math/la/SGDVector.java
+++ b/Math/src/main/java/org/tribuo/math/la/SGDVector.java
@@ -20,6 +20,7 @@
 
 import java.util.function.DoubleBinaryOperator;
 import java.util.function.DoubleUnaryOperator;
+import java.util.function.ToDoubleBiFunction;
 
 /**
  * Interface for 1 dimensional {@link Tensor}s.
@@ -28,6 +29,14 @@
  */
 public interface SGDVector extends Tensor, Iterable<VectorTuple> {
 
+    /**
+     * Applies a {@link ToDoubleBiFunction} elementwise to this {@link SGDVector}.
+     * <p>
+     * The first argument to the function is the index, the second argument is the current value.
+     * @param f The function to apply.
+     */
+    public void foreachIndexedInPlace(ToDoubleBiFunction<Integer,Double> f);
+
     /**
      * Returns a deep copy of this vector.
      * @return A copy of this vector.
diff --git a/Math/src/main/java/org/tribuo/math/la/SparseVector.java b/Math/src/main/java/org/tribuo/math/la/SparseVector.java
index 480ac403e..acbf408a0 100644
--- a/Math/src/main/java/org/tribuo/math/la/SparseVector.java
+++ b/Math/src/main/java/org/tribuo/math/la/SparseVector.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2015-2021, Oracle and/or its affiliates. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -35,6 +35,7 @@
 import java.util.Objects;
 import java.util.function.DoubleBinaryOperator;
 import java.util.function.DoubleUnaryOperator;
+import java.util.function.ToDoubleBiFunction;
 import java.util.stream.Collectors;
 
 /**
@@ -43,7 +44,7 @@
  * Uses binary search to look up a specific index, so it's usually faster to
  * use the iterator to iterate the values.
  * <p>
- * This vector has immutable indices. It cannot get new indices after construction,
+ * This vector has immutable indices. It cannot get new indices after construction,
  * and will throw {@link IllegalArgumentException} if such an operation is tried.
  */
 public class SparseVector implements SGDVector {
@@ -54,6 +55,17 @@ public class SparseVector implements SGDVector {
     protected final double[] values;
     private final int size;
 
+    /**
+     * Creates an empty sparse vector.
+     * @param size The dimension.
+     */
+    SparseVector(int size) {
+        this.indices = new int[0];
+        this.values = new double[0];
+        this.size = size;
+        this.shape = new int[]{size};
+    }
+
     /**
      * Used internally for performance.
      * Does not defensively copy the input, nor check it's sorted.
@@ -93,8 +105,8 @@ private SparseVector(SparseVector other) {
     /**
      * Creates a sparse vector of the specified size, with the supplied value at each of the indices.
      * @param size The vector size.
-     * @param indices The indices of the sparse vector.
-     * @param value The initial value.
+     * @param indices The active indices.
+     * @param value The value for those indices.
      */
     public SparseVector(int size, int[] indices, double value) {
         this.indices = Arrays.copyOf(indices,indices.length);
@@ -178,7 +190,7 @@ public static <T extends Output<T>> SparseVector createSparseVector(Example<T> e
     /**
      * Defensively copies the input, and checks that the indices are sorted. If not,
      * it sorts them.
-     *  <p>
+     * <p>
      * Throws {@link IllegalArgumentException} if the arrays are not the same length, or if size is less than
      * the max index.
      * @param dimension The dimension of this vector.
@@ -459,6 +471,16 @@ public void hadamardProductInPlace(Tensor other, DoubleUnaryOperator f) {
         }
     }
 
+    /**
+     * Applies a {@link ToDoubleBiFunction} elementwise to this {@link SGDVector}.
+     * <p>
+     * The first argument to the function is the index, the second argument is the current value.
+     * <p>
+     * Only applies the function to the elements which are present.
+     * <p>
+     * If you need to operate over the whole vector then densify it first.
+     * @param f The function to apply.
+     */
     @Override
     public void foreachInPlace(DoubleUnaryOperator f) {
         for (int i = 0; i < values.length; i++) {
@@ -466,6 +488,23 @@ public void foreachInPlace(DoubleUnaryOperator f) {
         }
     }
 
+    /**
+     * Applies a {@link ToDoubleBiFunction} elementwise to this {@link SGDVector}.
+     * <p>
+     * The first argument to the function is the index, the second argument is the current value.
+     * <p>
+     * Only applies the function to the elements which are present.
+     * <p>
+     * If you need to operate over the whole vector then densify it first.
+     * @param f The function to apply.
+     */
+    @Override
+    public void foreachIndexedInPlace(ToDoubleBiFunction<Integer,Double> f) {
+        for (int i = 0; i < values.length; i++) {
+            values[i] = f.applyAsDouble(indices[i],values[i]);
+        }
+    }
+
     @Override
     public SparseVector scale(double coefficient) {
         double[] newValues = Arrays.copyOf(values, values.length);
diff --git a/Math/src/main/java/org/tribuo/math/la/Tensor.java b/Math/src/main/java/org/tribuo/math/la/Tensor.java
index b0cd22efa..7daa08459 100644
--- a/Math/src/main/java/org/tribuo/math/la/Tensor.java
+++ b/Math/src/main/java/org/tribuo/math/la/Tensor.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2015-2021, Oracle and/or its affiliates. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,7 +28,7 @@ public interface Tensor extends Serializable {
     /**
      * The number of elements in this shape, i.e., the product of the shape array.
      * @param shape The tensor shape.
-     * @return The number of elements in the shape.
+     * @return The total number of elements.
      */
     public static int shapeSum(int[] shape) {
         int sum = 1;
@@ -68,6 +68,12 @@ public static boolean shapeCheck(Tensor first, Tensor second) {
      */
     public Tensor reshape(int[] shape);
 
+    /**
+     * Returns a copy of this Tensor.
+     * @return A copy of the Tensor.
+     */
+    public Tensor copy();
+
     /**
      * Updates this {@link Tensor} by adding all the values from the intersection with {@code other}.
      * <p>
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelModel.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelModel.java
new file mode 100644
index 000000000..a88f48d8f
--- /dev/null
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelModel.java
@@ -0,0 +1,99 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.multilabel.sgd.fm;
+
+import org.tribuo.Example;
+import org.tribuo.ImmutableFeatureMap;
+import org.tribuo.ImmutableOutputInfo;
+import org.tribuo.Prediction;
+import org.tribuo.classification.Label;
+import org.tribuo.common.sgd.AbstractFMModel;
+import org.tribuo.common.sgd.FMParameters;
+import org.tribuo.math.la.DenseVector;
+import org.tribuo.math.util.VectorNormalizer;
+import org.tribuo.multilabel.MultiLabel;
+import org.tribuo.provenance.ModelProvenance;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * The inference time version of a multi-label factorization machine trained using SGD.
+ * <p>
+ * See:
+ * <pre>
+ * Rendle, S.
+ * Factorization machines.
+ * 2010 IEEE International Conference on Data Mining
+ * </pre>
+ */
+public class FMMultiLabelModel extends AbstractFMModel<MultiLabel> {
+    private static final long serialVersionUID = 2L;
+
+    private final VectorNormalizer normalizer;
+    private final double threshold;
+
+    /**
+     * Constructs a linear regression model trained via SGD.
+     * @param name The model name.
+     * @param provenance The model provenance.
+     * @param featureIDMap The feature domain.
+     * @param outputIDInfo The output domain.
+     * @param parameters The model parameters (i.e., the weight matrix).
+     * @param normalizer The output normalizer (usually sigmoid or no-op).
+     * @param generatesProbabilities Does this model produce probabilistic outputs.
+     * @param threshold The threshold for emitting a label.
+     */
+    FMMultiLabelModel(String name, ModelProvenance provenance,
+                      ImmutableFeatureMap featureIDMap, ImmutableOutputInfo<MultiLabel> outputIDInfo,
+                      FMParameters parameters, VectorNormalizer normalizer, boolean generatesProbabilities, double threshold) {
+        super(name, provenance, featureIDMap, outputIDInfo, parameters, generatesProbabilities);
+        this.normalizer = normalizer;
+        this.threshold = threshold;
+    }
+
+    @Override
+    public Prediction<MultiLabel> predict(Example<MultiLabel> example) {
+        PredAndActive predTuple = predictSingle(example);
+        DenseVector outputs = predTuple.prediction;
+        outputs.normalize(normalizer);
+        Map<String,MultiLabel> fullLabels = new HashMap<>();
+        Set<Label> predictedLabels = new HashSet<>();
+        for (int i = 0; i < outputs.size(); i++) {
+            String labelName = outputIDInfo.getOutput(i).getLabelString();
+            double labelScore = outputs.get(i);
+            Label score = new Label(outputIDInfo.getOutput(i).getLabelString(),labelScore);
+            if (labelScore > threshold) {
+                predictedLabels.add(score);
+            }
+            fullLabels.put(labelName,new MultiLabel(score));
+        }
+        return new Prediction<>(new MultiLabel(predictedLabels), fullLabels, predTuple.numActiveFeatures - 1, example, generatesProbabilities);
+    }
+
+    @Override
+    protected String getDimensionName(int index) {
+        return outputIDInfo.getOutput(index).getLabelString();
+    }
+
+    @Override
+    protected FMMultiLabelModel copy(String newName, ModelProvenance newProvenance) {
+        return new FMMultiLabelModel(newName,newProvenance,featureIDMap,outputIDInfo,(FMParameters)modelParameters.copy(),normalizer,generatesProbabilities,threshold);
+    }
+}
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
new file mode 100644
index 000000000..d7a750588
--- /dev/null
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.multilabel.sgd.fm;
+
+import com.oracle.labs.mlrg.olcut.config.ArgumentException;
+import com.oracle.labs.mlrg.olcut.config.Option;
+import com.oracle.labs.mlrg.olcut.config.Options;
+import org.tribuo.Trainer;
+import org.tribuo.math.optimisers.GradientOptimiserOptions;
+import org.tribuo.multilabel.sgd.MultiLabelObjective;
+import org.tribuo.multilabel.sgd.objectives.BinaryCrossEntropy;
+import org.tribuo.multilabel.sgd.objectives.Hinge;
+
+import java.util.logging.Logger;
+
+/**
+ * CLI options for training a linear classifier.
+ */
+public class FMMultiLabelOptions implements Options {
+    private static final Logger logger = Logger.getLogger(FMMultiLabelOptions.class.getName());
+
+    /**
+     * Available loss types.
+     */
+    public enum LossEnum {HINGE, SIGMOID}
+
+    public GradientOptimiserOptions sgoOptions;
+
+    @Option(longName = "fm-epochs", usage = "Number of SGD epochs.")
+    public int fmEpochs = 5;
+    @Option(longName = "fm-objective", usage = "Loss function.")
+    public LossEnum fmObjective = LossEnum.SIGMOID;
+    @Option(longName = "fm-logging-interval", usage = "Log the objective after <int> examples.")
+    public int fmLoggingInterval = 100;
+    @Option(longName = "fm-minibatch-size", usage = "Minibatch size.")
+    public int fmMinibatchSize = 1;
+    @Option(longName = "fm-seed", usage = "Sets the random seed for the FMMultiLabelTrainer.")
+    private long fmSeed = Trainer.DEFAULT_SEED;
+    @Option(longName = "fm-factor-size", usage = "Factor size.")
+    public int fmFactorSize = 6;
+    @Option(longName = "fm-l2-penalty", usage = "L2 regularization penalty.")
+    public double fmL2 = 0.001;
+    @Option(longName = "fm-variance", usage = "Variance of the initialization gaussian.")
+    public double fmVariance = 0.1;
+
+    /**
+     * Returns the loss function specified in the arguments.
+     * @return The loss function.
+     */
+    public MultiLabelObjective getLoss() {
+        switch (fmObjective) {
+            case HINGE:
+                return new Hinge();
+            case SIGMOID:
+                return new BinaryCrossEntropy();
+            default:
+                throw new ArgumentException("sgd-objective", "Unknown loss function " + fmObjective);
+        }
+    }
+
+    public FMMultiLabelTrainer getTrainer() {
+        logger.info(String.format("Set logging interval to %d", fmLoggingInterval));
+        return new FMMultiLabelTrainer(getLoss(), sgoOptions.getOptimiser(), fmEpochs, fmLoggingInterval,
+                fmMinibatchSize, fmSeed, fmFactorSize, fmL2, fmVariance);
+    }
+}
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelTrainer.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelTrainer.java
new file mode 100644
index 000000000..eecc10bbf
--- /dev/null
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelTrainer.java
@@ -0,0 +1,141 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.multilabel.sgd.fm;
+
+import com.oracle.labs.mlrg.olcut.config.Config;
+import org.tribuo.ImmutableFeatureMap;
+import org.tribuo.ImmutableOutputInfo;
+import org.tribuo.common.sgd.AbstractFMTrainer;
+import org.tribuo.common.sgd.FMParameters;
+import org.tribuo.common.sgd.SGDObjective;
+import org.tribuo.math.StochasticGradientOptimiser;
+import org.tribuo.math.la.SGDVector;
+import org.tribuo.math.la.SparseVector;
+import org.tribuo.multilabel.MultiLabel;
+import org.tribuo.multilabel.sgd.MultiLabelObjective;
+import org.tribuo.multilabel.sgd.objectives.BinaryCrossEntropy;
+import org.tribuo.provenance.ModelProvenance;
+
+import java.util.logging.Logger;
+
+/**
+ * A trainer for a multi-label classification factorization machine using SGD.
+ * <p>
+ * See:
+ * <pre>
+ * Rendle, S.
+ * Factorization machines.
+ * 2010 IEEE International Conference on Data Mining
+ * </pre>
+ */
+public class FMMultiLabelTrainer extends AbstractFMTrainer<MultiLabel,SGDVector> {
+    private static final Logger logger = Logger.getLogger(FMMultiLabelTrainer.class.getName());
+
+    @Config(description="The classification objective function to use.")
+    private MultiLabelObjective objective = new BinaryCrossEntropy();
+
+    /**
+     * Constructs an SGD trainer for a multi-label factorization machine.
+     * @param objective The objective function to optimise.
+     * @param optimiser The gradient optimiser to use.
+     * @param epochs The number of epochs (complete passes through the training data).
+     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
+     * @param minibatchSize The size of any minibatches.
+     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     * @param factorizedDimSize Size of the factorized feature representation.
+     * @param l2                The l2 regularisation penalty.
+     * @param variance          The variance of the initializer.
+     */
+    public FMMultiLabelTrainer(MultiLabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
+                               int loggingInterval, int minibatchSize, long seed,
+                               int factorizedDimSize, double l2, double variance) {
+        super(optimiser,epochs,loggingInterval,minibatchSize,seed, factorizedDimSize, l2, variance);
+        this.objective = objective;
+    }
+
+    /**
+     * Constructs an SGD trainer for a multi-label factorization machine.
+     * <p>
+     * Sets the minibatch size to 1.
+     * @param objective The objective function to optimise.
+     * @param optimiser The gradient optimiser to use.
+     * @param epochs The number of epochs (complete passes through the training data).
+     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
+     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     * @param factorizedDimSize Size of the factorized feature representation.
+     * @param l2                The l2 regularisation penalty.
+     * @param variance          The variance of the initializer.
+     */
+    public FMMultiLabelTrainer(MultiLabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
+                               int loggingInterval, long seed,
+                               int factorizedDimSize, double l2, double variance) {
+        this(objective,optimiser,epochs,loggingInterval,1,seed, factorizedDimSize, l2, variance);
+    }
+
+    /**
+     * Constructs an SGD trainer for a multi-label factorization machine.
+     * <p>
+     * Sets the minibatch size to 1 and the logging interval to 1000.
+     * @param objective The objective function to optimise.
+     * @param optimiser The gradient optimiser to use.
+     * @param epochs The number of epochs (complete passes through the training data).
+     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     * @param factorizedDimSize Size of the factorized feature representation.
+     * @param l2                The l2 regularisation penalty.
+     * @param variance          The variance of the initializer.
+     */
+    public FMMultiLabelTrainer(MultiLabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
+                               long seed, int factorizedDimSize, double l2, double variance) {
+        this(objective,optimiser,epochs,1000,1,seed, factorizedDimSize, l2, variance);
+    }
+
+    /**
+     * For olcut.
+     */
+    private FMMultiLabelTrainer() {
+        super();
+    }
+
+    @Override
+    protected SparseVector getTarget(ImmutableOutputInfo<MultiLabel> outputInfo, MultiLabel output) {
+        return output.convertToSparseVector(outputInfo);
+    }
+
+    @Override
+    protected SGDObjective<SGDVector> getObjective() {
+        return objective;
+    }
+
+    @Override
+    protected FMMultiLabelModel createModel(String name, ModelProvenance provenance, ImmutableFeatureMap featureMap, ImmutableOutputInfo<MultiLabel> outputInfo, FMParameters parameters) {
+        return new FMMultiLabelModel(name,provenance,featureMap,outputInfo,parameters,objective.getNormalizer(),objective.isProbabilistic(),objective.threshold());
+    }
+
+    @Override
+    protected String getModelClassName() {
+        return FMMultiLabelModel.class.getName();
+    }
+
+    @Override
+    public String toString() {
+        return "FMMultiLabelTrainer(objective=" + objective.toString() + ",optimiser=" + optimiser.toString() +
+                ",epochs=" + epochs + ",minibatchSize=" + minibatchSize + ",seed=" + seed +
+                ",factorizedDimSize=" + factorizedDimSize + ",l2=" + l2 + ",variance=" + variance +
+                ")";
+    }
+
+}
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/package-info.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/package-info.java
new file mode 100644
index 000000000..be7cb7702
--- /dev/null
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/package-info.java
@@ -0,0 +1,20 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * Provides an implementation of a multi-label classification factorization machine model using Stochastic Gradient Descent.
+ */
+package org.tribuo.multilabel.sgd.fm;
\ No newline at end of file
diff --git a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java
index fce49f683..94a1ab92a 100644
--- a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java
+++ b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java
@@ -64,7 +64,7 @@ public class TestSGDLinear {
     @BeforeAll
     public static void setup() {
         Class<?>[] classes = new Class<?>[]{AbstractSGDTrainer.class, AbstractLinearSGDTrainer.class,LinearSGDTrainer.class};
-        for (Class c : classes) {
+        for (Class<?> c : classes) {
             Logger logger = Logger.getLogger(c.getName());
             logger.setLevel(Level.WARNING);
         }
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
new file mode 100644
index 000000000..f6d8fbcaf
--- /dev/null
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
@@ -0,0 +1,79 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.regression.sgd.fm;
+
+import org.tribuo.Example;
+import org.tribuo.ImmutableFeatureMap;
+import org.tribuo.ImmutableOutputInfo;
+import org.tribuo.Prediction;
+import org.tribuo.common.sgd.AbstractFMModel;
+import org.tribuo.common.sgd.FMParameters;
+import org.tribuo.provenance.ModelProvenance;
+import org.tribuo.regression.Regressor;
+
+import java.util.Arrays;
+
+/**
+ * The inference time model of a regression factorization machine trained using SGD.
+ * Independently predicts each output dimension, unless they are tied together in the
+ * optimiser.
+ * <p>
+ * See:
+ * <pre>
+ * Rendle, S.
+ * Factorization machines.
+ * 2010 IEEE International Conference on Data Mining
+ * </pre>
+ */
+public class FMRegressionModel extends AbstractFMModel<Regressor> {
+    private static final long serialVersionUID = 3L;
+
+    private final String[] dimensionNames;
+
+    /**
+     * Constructs a linear regression model trained via SGD.
+     * @param name The model name.
+     * @param dimensionNames The regression dimension names.
+     * @param provenance The model provenance.
+     * @param featureIDMap The feature domain.
+     * @param outputIDInfo The output domain.
+     * @param parameters The model parameters.
+     */
+    FMRegressionModel(String name, String[] dimensionNames, ModelProvenance provenance,
+                      ImmutableFeatureMap featureIDMap, ImmutableOutputInfo<Regressor> outputIDInfo,
+                      FMParameters parameters) {
+        super(name, provenance, featureIDMap, outputIDInfo, parameters, false);
+        this.dimensionNames = dimensionNames;
+    }
+
+    @Override
+    public Prediction<Regressor> predict(Example<Regressor> example) {
+        PredAndActive predTuple = predictSingle(example);
+        return new Prediction<>(new Regressor(dimensionNames,predTuple.prediction.toArray()), predTuple.numActiveFeatures, example);
+    }
+
+    @Override
+    protected FMRegressionModel copy(String newName, ModelProvenance newProvenance) {
+        return new FMRegressionModel(newName,Arrays.copyOf(dimensionNames,dimensionNames.length),newProvenance,featureIDMap,outputIDInfo,(FMParameters)modelParameters.copy());
+    }
+
+    @Override
+    protected String getDimensionName(int index) {
+        return dimensionNames[index];
+    }
+
+}
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
new file mode 100644
index 000000000..a445bab47
--- /dev/null
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
@@ -0,0 +1,150 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.regression.sgd.fm;
+
+import com.oracle.labs.mlrg.olcut.config.Config;
+import org.tribuo.ImmutableFeatureMap;
+import org.tribuo.ImmutableOutputInfo;
+import org.tribuo.common.sgd.AbstractFMTrainer;
+import org.tribuo.common.sgd.FMParameters;
+import org.tribuo.common.sgd.SGDObjective;
+import org.tribuo.math.StochasticGradientOptimiser;
+import org.tribuo.math.la.DenseVector;
+import org.tribuo.provenance.ModelProvenance;
+import org.tribuo.regression.Regressor;
+import org.tribuo.regression.sgd.RegressionObjective;
+
+import java.util.logging.Logger;
+
+/**
+ * A trainer for a regression factorization machine using SGD.
+ * Independently trains each output dimension, unless they are tied together in the
+ * optimiser.
+ * <p>
+ * See:
+ * <pre>
+ * Rendle, S.
+ * Factorization machines.
+ * 2010 IEEE International Conference on Data Mining
+ * </pre>
+ */
+public class FMRegressionTrainer extends AbstractFMTrainer<Regressor,DenseVector> {
+    private static final Logger logger = Logger.getLogger(FMRegressionTrainer.class.getName());
+
+    @Config(mandatory = true,description="The regression objective to use.")
+    private RegressionObjective objective;
+
+    /**
+     * Constructs an SGD trainer for a factorization machine.
+     * @param objective The objective function to optimise.
+     * @param optimiser The gradient optimiser to use.
+     * @param epochs The number of epochs (complete passes through the training data).
+     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
+     * @param minibatchSize The size of any minibatches.
+     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     * @param factorizedDimSize Size of the factorized feature representation.
+     * @param l2                The l2 regularisation penalty.
+     * @param variance          The variance of the initializer.
+     */
+    public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOptimiser optimiser, int epochs,
+                               int loggingInterval, int minibatchSize, long seed,
+                               int factorizedDimSize, double l2, double variance) {
+        super(optimiser, epochs, loggingInterval, minibatchSize, seed, factorizedDimSize, l2, variance);
+        this.objective = objective;
+    }
+
+    /**
+     * Constructs an SGD trainer for a factorization machine.
+     * <p>
+     * Sets the minibatch size to 1.
+     * @param objective The objective function to optimise.
+     * @param optimiser The gradient optimiser to use.
+     * @param epochs The number of epochs (complete passes through the training data).
+     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
+     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     * @param factorizedDimSize Size of the factorized feature representation.
+     * @param l2                The l2 regularisation penalty.
+     * @param variance          The variance of the initializer.
+     */
+    public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOptimiser optimiser, int epochs,
+                               int loggingInterval, long seed,
+                               int factorizedDimSize, double l2, double variance) {
+        this(objective,optimiser,epochs,loggingInterval,1,seed,factorizedDimSize,l2,variance);
+    }
+
+    /**
+     * Constructs an SGD trainer for a factorization machine.
+     * <p>
+     * Sets the minibatch size to 1 and the logging interval to 1000.
+     * @param objective The objective function to optimise.
+     * @param optimiser The gradient optimiser to use.
+     * @param epochs The number of epochs (complete passes through the training data).
+     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     * @param factorizedDimSize Size of the factorized feature representation.
+     * @param l2                The l2 regularisation penalty.
+     * @param variance          The variance of the initializer.
+     */
+    public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOptimiser optimiser, int epochs,
+                               long seed, int factorizedDimSize, double l2, double variance) {
+        this(objective,optimiser,epochs,1000,1,seed,factorizedDimSize,l2,variance);
+    }
+
+    /**
+     * For olcut.
+     */
+    private FMRegressionTrainer() {
+        super();
+    }
+
+    @Override
+    protected DenseVector getTarget(ImmutableOutputInfo<Regressor> outputInfo, Regressor output) {
+        double[] regressorsBuffer = new double[outputInfo.size()];
+        for (Regressor.DimensionTuple r : output) {
+            int id = outputInfo.getID(r);
+            regressorsBuffer[id] = r.getValue();
+        }
+        return DenseVector.createDenseVector(regressorsBuffer);
+    }
+
+    @Override
+    protected SGDObjective<DenseVector> getObjective() {
+        return objective;
+    }
+
+    @Override
+    protected FMRegressionModel createModel(String name, ModelProvenance provenance, ImmutableFeatureMap featureMap, ImmutableOutputInfo<Regressor> outputInfo, FMParameters parameters) {
+        String[] dimensionNames = new String[outputInfo.size()];
+        for (Regressor r : outputInfo.getDomain()) {
+            int id = outputInfo.getID(r);
+            dimensionNames[id] = r.getNames()[0];
+        }
+        return new FMRegressionModel("linear-sgd-model",dimensionNames,provenance,featureMap,outputInfo,parameters);
+    }
+
+    @Override
+    protected String getModelClassName() {
+        return FMRegressionModel.class.getName();
+    }
+
+    @Override
+    public String toString() {
+        return "FMRegressionTrainer(objective=" + objective.toString() + ",optimiser=" + optimiser.toString() +
+                ",epochs=" + epochs + ",minibatchSize=" + minibatchSize + ",seed=" + seed +
+                ",factorizedDimSize=" + factorizedDimSize + ",l2=" + l2 + ",variance=" + variance +
+                ")";
+    }
+}
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
new file mode 100644
index 000000000..a289d4bfa
--- /dev/null
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
@@ -0,0 +1,164 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.regression.sgd.fm;
+
+import com.oracle.labs.mlrg.olcut.config.ConfigurationManager;
+import com.oracle.labs.mlrg.olcut.config.Option;
+import com.oracle.labs.mlrg.olcut.config.Options;
+import com.oracle.labs.mlrg.olcut.config.UsageException;
+import com.oracle.labs.mlrg.olcut.util.LabsLogFormatter;
+import com.oracle.labs.mlrg.olcut.util.Pair;
+import org.tribuo.Dataset;
+import org.tribuo.Model;
+import org.tribuo.Trainer;
+import org.tribuo.data.DataOptions;
+import org.tribuo.math.StochasticGradientOptimiser;
+import org.tribuo.math.optimisers.GradientOptimiserOptions;
+import org.tribuo.regression.RegressionFactory;
+import org.tribuo.regression.Regressor;
+import org.tribuo.regression.evaluation.RegressionEvaluation;
+import org.tribuo.regression.sgd.RegressionObjective;
+import org.tribuo.regression.sgd.linear.LinearSGDTrainer;
+import org.tribuo.regression.sgd.objectives.AbsoluteLoss;
+import org.tribuo.regression.sgd.objectives.Huber;
+import org.tribuo.regression.sgd.objectives.SquaredLoss;
+import org.tribuo.util.Util;
+
+import java.io.IOException;
+import java.util.logging.Logger;
+
+/**
+ * Build and run a regression factorization machine for a standard dataset.
+ */
+public class TrainTest {
+
+    private static final Logger logger = Logger.getLogger(TrainTest.class.getName());
+
+    /**
+     * Loss function.
+     */
+    public enum LossEnum {
+        /**
+         * Uses {@link AbsoluteLoss}.
+         */
+        ABSOLUTE,
+        /**
+         * Uses {@link SquaredLoss}.
+         */
+        SQUARED,
+        /**
+         * Uses {@link Huber} loss.
+         */
+        HUBER
+    }
+
+    /**
+     * Command line options.
+     */
+    public static class SGDOptions implements Options {
+        @Override
+        public String getOptionsDescription() {
+            return "Trains and tests a linear SGD regression model on the specified datasets.";
+        }
+        public DataOptions general;
+        public GradientOptimiserOptions gradientOptions;
+
+        @Option(charName='i',longName="epochs",usage="Number of SGD epochs. Defaults to 5.")
+        public int epochs = 5;
+        @Option(charName='o',longName="objective",usage="Loss function. Defaults to SQUARED.")
+        public LossEnum loss = LossEnum.SQUARED;
+        @Option(charName='p',longName="logging-interval",usage="Log the objective after <int> examples. Defaults to 100.")
+        public int loggingInterval = 100;
+        @Option(charName='z',longName="minibatch-size",usage="Minibatch size. Defaults to 1.")
+        public int minibatchSize = 1;
+        @Option(charName = 'd', longName = "factor-size", usage = "Factor size.")
+        public int factorSize = 6;
+        @Option(charName = 'k', longName = "l2-penalty", usage = "L2 regularization penalty.")
+        public double l2 = 0.001;
+        @Option(longName = "variance", usage = "Variance of the initialization gaussian.")
+        public double variance = 0.1;
+    }
+
+    /**
+     * @param args the command line arguments
+     * @throws IOException if there is any error reading the examples.
+     */
+    public static void main(String[] args) throws IOException {
+
+        //
+        // Use the labs format logging.
+        LabsLogFormatter.setAllLogFormatters();
+
+        SGDOptions o = new SGDOptions();
+        ConfigurationManager cm;
+        try {
+            cm = new ConfigurationManager(args,o);
+        } catch (UsageException e) {
+            logger.info(e.getMessage());
+            return;
+        }
+
+        if (o.general.trainingPath == null || o.general.testingPath == null) {
+            logger.info(cm.usage());
+            return;
+        }
+
+        logger.info("Configuring gradient optimiser");
+        RegressionObjective obj = null;
+        switch (o.loss) {
+            case ABSOLUTE:
+                obj = new AbsoluteLoss();
+                break;
+            case SQUARED:
+                obj = new SquaredLoss();
+                break;
+            case HUBER:
+                obj = new Huber();
+                break;
+            default:
+                logger.warning("Unknown objective function " + o.loss);
+                logger.info(cm.usage());
+                return;
+        }
+        StochasticGradientOptimiser grad = o.gradientOptions.getOptimiser();
+
+        logger.info(String.format("Set logging interval to %d",o.loggingInterval));
+        RegressionFactory factory = new RegressionFactory();
+
+        Pair<Dataset<Regressor>,Dataset<Regressor>> data = o.general.load(factory);
+        Dataset<Regressor> train = data.getA();
+        Dataset<Regressor> test = data.getB();
+
+        Trainer<Regressor> trainer = new FMRegressionTrainer(obj,grad,o.epochs,o.loggingInterval,o.minibatchSize,o.general.seed,o.factorSize,o.l2,o.variance);
+        logger.info("Training using " + trainer.toString());
+        final long trainStart = System.currentTimeMillis();
+        Model<Regressor> model = trainer.train(train);
+        final long trainStop = System.currentTimeMillis();
+
+        logger.info("Finished training regressor " + Util.formatDuration(trainStart,trainStop));
+
+        final long testStart = System.currentTimeMillis();
+        RegressionEvaluation evaluation = factory.getEvaluator().evaluate(model,test);
+        final long testStop = System.currentTimeMillis();
+        logger.info("Finished evaluating model " + Util.formatDuration(testStart,testStop));
+        System.out.println(evaluation.toString());
+
+        if (o.general.outputPath != null) {
+            o.general.saveModel(model);
+        }
+    }
+}
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/package-info.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/package-info.java
new file mode 100644
index 000000000..7f192e3e3
--- /dev/null
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/package-info.java
@@ -0,0 +1,20 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * Provides an implementation of factorization machines for regression using Stochastic Gradient Descent.
+ */
+package org.tribuo.regression.sgd.fm;
\ No newline at end of file
diff --git a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java
index fbe968ffc..cbce5b3ba 100644
--- a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java
+++ b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java
@@ -75,7 +75,7 @@ public class TestSGDLinear {
     @BeforeAll
     public static void setup() {
         Class<?>[] classes = new Class<?>[]{AbstractSGDTrainer.class, AbstractLinearSGDTrainer.class,LinearSGDTrainer.class};
-        for (Class c : classes) {
+        for (Class<?> c : classes) {
             Logger logger = Logger.getLogger(c.getName());
             logger.setLevel(Level.WARNING);
         }

From 9aec04407a9517ec2803153bdd4c1df7c4fb0579 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 30 May 2021 00:43:02 -0400
Subject: [PATCH 02/31] Tidying up copyrights and some comments.

---
 .../sgd/fm/FMClassificationOptions.java          |  2 +-
 .../tribuo/classification/sgd/fm/TrainTest.java  |  2 +-
 .../sgd/fm/TestFMClassification.java             | 16 ----------------
 .../java/org/tribuo/common/sgd/FMParameters.java |  2 +-
 .../multilabel/sgd/fm/FMMultiLabelOptions.java   |  2 +-
 5 files changed, 4 insertions(+), 20 deletions(-)

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
index fd8aab1eb..0825686f7 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/TrainTest.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/TrainTest.java
index d52e2cc09..61e5e58a7 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/TrainTest.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/TrainTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
index 01a84fdb4..9839b124a 100644
--- a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
+++ b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
@@ -70,22 +70,6 @@ public static Model<Label> testFMClassification(Pair<Dataset<Label>,Dataset<Labe
         return m;
     }
 
-    /*
-    @Test
-    public void testSingleClassTraining() {
-        Pair<Dataset<Label>,Dataset<Label>> data = LabelledDataGenerator.denseTrainTest();
-        // Done to reduce sensitivity to test run order.
-        Trainer<Label> trainer = new FMClassificationTrainer(new Hinge(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED,6,0.0,0.1);
-        DatasetView<Label> trainingData = DatasetView.createView(data.getA(),(Example<Label> e) -> e.getOutput().getLabel().equals("Foo"), "Foo selector");
-        Model<Label> model = trainer.train(trainingData);
-        LabelEvaluation evaluation = (LabelEvaluation) trainingData.getOutputFactory().getEvaluator().evaluate(model,data.getB());
-        assertEquals(0.0,evaluation.accuracy(new Label("Baz")));
-        assertEquals(0.0,evaluation.accuracy(new Label("Quux")));
-        assertEquals(1.0,evaluation.recall(new Label("Foo")));
-        assertEquals(1.0,evaluation.recall(new Label("Bar"))); // This is due to the random init of the classifier.
-    }
-    */
-
     @Test
     public void testDenseData() {
         Pair<Dataset<Label>,Dataset<Label>> p = LabelledDataGenerator.denseTrainTest();
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
index a5c13bb80..422e63621 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
@@ -36,7 +36,7 @@
 import java.util.SplittableRandom;
 
 /**
- * A {@link Parameters} for producing factorization machines.
+ * A {@link Parameters} for factorization machines.
  */
 public final class FMParameters implements FeedForwardParameters {
     private static final long serialVersionUID = 1L;
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
index d7a750588..bf7de9a68 100644
--- a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

From ebfc347129ed8366521d38808afbf13d6f7125e0 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 30 May 2021 00:49:38 -0400
Subject: [PATCH 03/31] Fixing a bug in the model naming for FMRegressionModel
 and LinearSGDModel.

---
 .../java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java  | 2 +-
 .../java/org/tribuo/regression/sgd/linear/LinearSGDTrainer.java | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
index a445bab47..cb5408324 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
@@ -132,7 +132,7 @@ protected FMRegressionModel createModel(String name, ModelProvenance provenance,
             int id = outputInfo.getID(r);
             dimensionNames[id] = r.getNames()[0];
         }
-        return new FMRegressionModel("linear-sgd-model",dimensionNames,provenance,featureMap,outputInfo,parameters);
+        return new FMRegressionModel(name,dimensionNames,provenance,featureMap,outputInfo,parameters);
     }
 
     @Override
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDTrainer.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDTrainer.java
index 111cfb185..fe8399a63 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDTrainer.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDTrainer.java
@@ -118,7 +118,7 @@ protected LinearSGDModel createModel(String name, ModelProvenance provenance, Im
             int id = outputInfo.getID(r);
             dimensionNames[id] = r.getNames()[0];
         }
-        return new LinearSGDModel("linear-sgd-model",dimensionNames,provenance,featureMap,outputInfo,parameters);
+        return new LinearSGDModel(name,dimensionNames,provenance,featureMap,outputInfo,parameters);
     }
 
     @Override

From 9d18400918ffc1544106bf9299e48a7b333f252c Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 30 May 2021 21:31:21 -0400
Subject: [PATCH 04/31] Fixing a gradient initialisation bug for dense FMs.

---
 .../src/main/java/org/tribuo/common/sgd/FMParameters.java    | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
index 422e63621..0da8a5d93 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
@@ -185,6 +185,11 @@ public Tensor[] gradients(Pair<Double, SGDVector> score, SGDVector features) {
                     factorGradMatrix = new DenseSparseMatrix(vectors);
                 } else {
                     factorGradMatrix = new DenseMatrix(numFactors,features.size());
+                    for (int j = 0; j < numFactors; j++) {
+                        for (int k = 0; k < features.size(); k++) {
+                            factorGradMatrix.set(j,k,features.get(k));
+                        }
+                    }
                 }
                 for (int j = 0; j < numFactors; j++) {
                     // This gets a mutable view of the row

From cf14b490d87ab20b113b963dcff2a05875b73d66 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 30 May 2021 21:32:03 -0400
Subject: [PATCH 05/31] Adding int id lookup methods to
 ImmutableRegressionInfo's dimension stats.

---
 .../regression/ImmutableRegressionInfo.java   | 70 +++++++++++++++++++
 .../regression/MutableRegressionInfo.java     |  2 +-
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/Regression/Core/src/main/java/org/tribuo/regression/ImmutableRegressionInfo.java b/Regression/Core/src/main/java/org/tribuo/regression/ImmutableRegressionInfo.java
index 927cc2f12..dd2e6d9c4 100644
--- a/Regression/Core/src/main/java/org/tribuo/regression/ImmutableRegressionInfo.java
+++ b/Regression/Core/src/main/java/org/tribuo/regression/ImmutableRegressionInfo.java
@@ -21,6 +21,7 @@
 import com.oracle.labs.mlrg.olcut.util.Pair;
 import org.tribuo.ImmutableOutputInfo;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
@@ -50,6 +51,11 @@ public class ImmutableRegressionInfo extends RegressionInfo implements Immutable
 
     private final Set<Regressor> domain;
 
+    private transient double[] minArray;
+    private transient double[] maxArray;
+    private transient double[] meanArray;
+    private transient double[] varianceArray;
+
     /**
      * Copies an ImmutableRegressionInfo including the ids.
      * @param info The info to copy.
@@ -61,6 +67,7 @@ private ImmutableRegressionInfo(ImmutableRegressionInfo info) {
         labelIDMap = new LinkedHashMap<>();
         labelIDMap.putAll(info.labelIDMap);
         domain = calculateDomain(minMap);
+        computeStatisticArrays();
     }
 
     /**
@@ -82,6 +89,7 @@ private ImmutableRegressionInfo(ImmutableRegressionInfo info) {
             counter++;
         }
         domain = calculateDomain(minMap);
+        computeStatisticArrays();
     }
 
     /**
@@ -125,6 +133,7 @@ else if (curNames.length == 1) {
             throw new IllegalArgumentException("Mapping must contain an entry for each element in the info, found " + labelIDMap.keySet() + " and " + countMap.keySet());
         }
         domain = calculateDomain(minMap);
+        computeStatisticArrays();
     }
 
     /**
@@ -167,11 +176,65 @@ public Regressor getOutput(int id) {
         }
     }
 
+    /**
+     * Gets the minimum value this RegressionInfo has seen for that id.
+     * @param id The dimension to check.
+     * @return The minimum value for that dimension.
+     */
+    public double getMin(int id) {
+        return minArray[id];
+    }
+
+    /**
+     * Gets the maximum value this RegressionInfo has seen for that id.
+     * @param id The dimension to check.
+     * @return The maximum value for that dimension.
+     */
+    public double getMax(int id) {
+        return maxArray[id];
+    }
+
+    /**
+     * Gets the mean value this RegressionInfo has seen for that id.
+     * @param id The dimension to check.
+     * @return The mean value for that dimension.
+     */
+    public double getMean(int id) {
+        return meanArray[id];
+    }
+
+    /**
+     * Gets the variance this RegressionInfo has seen for that id.
+     * @param id The dimension to check.
+     * @return The variance for that dimension.
+     */
+    public double getVariance(int id) {
+        return varianceArray[id];
+    }
+
     @Override
     public long getTotalObservations() {
         return overallCount;
     }
 
+    /**
+     * Computes the statistic arrays on construction and after deserialization.
+     */
+    private void computeStatisticArrays() {
+        int size = labelIDMap.size();
+        minArray = new double[size];
+        maxArray = new double[size];
+        meanArray = new double[size];
+        varianceArray = new double[size];
+        for (int i = 0; i < size; i++) {
+            String name = idLabelMap.get(i);
+            minArray[i] = minMap.get(name).doubleValue();
+            maxArray[i] = maxMap.get(name).doubleValue();
+            meanArray[i] = meanMap.get(name).doubleValue();
+            varianceArray[i] = sumSquaresMap.get(name).doubleValue() / (countMap.get(name).longValue()-1);
+        }
+    }
+
     @Override
     public ImmutableRegressionInfo copy() {
         return new ImmutableRegressionInfo(this);
@@ -290,4 +353,11 @@ public Pair<Integer, Regressor> next() {
             return new Pair<>(e.getKey(),new Regressor.DimensionTuple(e.getValue(),1.0));
         }
     }
+
+    private void readObject(java.io.ObjectInputStream in) throws IOException, ClassNotFoundException {
+        in.defaultReadObject();
+
+        // Rebuild statistic arrays
+        computeStatisticArrays();
+    }
 }
diff --git a/Regression/Core/src/main/java/org/tribuo/regression/MutableRegressionInfo.java b/Regression/Core/src/main/java/org/tribuo/regression/MutableRegressionInfo.java
index f6f09b875..1b7e4bb0f 100644
--- a/Regression/Core/src/main/java/org/tribuo/regression/MutableRegressionInfo.java
+++ b/Regression/Core/src/main/java/org/tribuo/regression/MutableRegressionInfo.java
@@ -101,7 +101,7 @@ public MutableRegressionInfo copy() {
     @Override
     public String toString() {
         StringBuilder builder = new StringBuilder();
-        builder.append("MultipleRegressionOutput(");
+        builder.append("RegressionInfo(");
         for (Map.Entry<String,MutableLong> e : countMap.entrySet()) {
             String name = e.getKey();
             long count = e.getValue().longValue();

From bdffcde9e36f23801a669ee0f3b245f18067e3b4 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 30 May 2021 21:32:22 -0400
Subject: [PATCH 06/31] Adding standardisation to factorisation machine
 regressors.

---
 .../regression/sgd/fm/FMRegressionModel.java  | 30 +++++++++++++++++--
 .../sgd/fm/FMRegressionTrainer.java           | 29 +++++++++++++-----
 .../tribuo/regression/sgd/fm/TrainTest.java   |  6 +++-
 3 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
index f6d8fbcaf..f0fa55ba1 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
@@ -23,6 +23,8 @@
 import org.tribuo.common.sgd.AbstractFMModel;
 import org.tribuo.common.sgd.FMParameters;
 import org.tribuo.provenance.ModelProvenance;
+import org.tribuo.regression.ImmutableRegressionInfo;
+import org.tribuo.regression.RegressionInfo;
 import org.tribuo.regression.Regressor;
 
 import java.util.Arrays;
@@ -44,6 +46,8 @@ public class FMRegressionModel extends AbstractFMModel<Regressor> {
 
     private final String[] dimensionNames;
 
+    private final boolean standardise;
+
     /**
      * Constructs a linear regression model trained via SGD.
      * @param name The model name.
@@ -52,23 +56,43 @@ public class FMRegressionModel extends AbstractFMModel<Regressor> {
      * @param featureIDMap The feature domain.
      * @param outputIDInfo The output domain.
      * @param parameters The model parameters.
+     * @param standardise Is the model fitted on standardised regressors?
      */
     FMRegressionModel(String name, String[] dimensionNames, ModelProvenance provenance,
                       ImmutableFeatureMap featureIDMap, ImmutableOutputInfo<Regressor> outputIDInfo,
-                      FMParameters parameters) {
+                      FMParameters parameters, boolean standardise) {
         super(name, provenance, featureIDMap, outputIDInfo, parameters, false);
         this.dimensionNames = dimensionNames;
+        this.standardise = standardise;
     }
 
     @Override
     public Prediction<Regressor> predict(Example<Regressor> example) {
         PredAndActive predTuple = predictSingle(example);
-        return new Prediction<>(new Regressor(dimensionNames,predTuple.prediction.toArray()), predTuple.numActiveFeatures, example);
+        double[] predictions = predTuple.prediction.toArray();
+        if (standardise) {
+            predictions = unstandardisePredictions(predictions);
+        }
+        return new Prediction<>(new Regressor(dimensionNames,predictions), predTuple.numActiveFeatures, example);
+    }
+
+    /**
+     * Converts zero mean unit variance predictions into the true range.
+     * @param predictions The predictions to convert.
+     */
+    private double[] unstandardisePredictions(double[] predictions) {
+        ImmutableRegressionInfo info = (ImmutableRegressionInfo) outputIDInfo;
+        for (int i = 0; i < predictions.length; i++) {
+            double mean = info.getMean(i);
+            double variance = info.getVariance(i);
+            predictions[i] = (predictions[i] * variance) + mean;
+        }
+        return predictions;
     }
 
     @Override
     protected FMRegressionModel copy(String newName, ModelProvenance newProvenance) {
-        return new FMRegressionModel(newName,Arrays.copyOf(dimensionNames,dimensionNames.length),newProvenance,featureIDMap,outputIDInfo,(FMParameters)modelParameters.copy());
+        return new FMRegressionModel(newName,Arrays.copyOf(dimensionNames,dimensionNames.length),newProvenance,featureIDMap,outputIDInfo,(FMParameters)modelParameters.copy(),standardise);
     }
 
     @Override
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
index cb5408324..8d12e6755 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
@@ -25,6 +25,7 @@
 import org.tribuo.math.StochasticGradientOptimiser;
 import org.tribuo.math.la.DenseVector;
 import org.tribuo.provenance.ModelProvenance;
+import org.tribuo.regression.ImmutableRegressionInfo;
 import org.tribuo.regression.Regressor;
 import org.tribuo.regression.sgd.RegressionObjective;
 
@@ -48,6 +49,9 @@ public class FMRegressionTrainer extends AbstractFMTrainer<Regressor,DenseVector
     @Config(mandatory = true,description="The regression objective to use.")
     private RegressionObjective objective;
 
+    @Config(mandatory = true,description="Standardise the output variables before fitting the model.")
+    private boolean standardise;
+
     /**
      * Constructs an SGD trainer for a factorization machine.
      * @param objective The objective function to optimise.
@@ -59,12 +63,14 @@ public class FMRegressionTrainer extends AbstractFMTrainer<Regressor,DenseVector
      * @param factorizedDimSize Size of the factorized feature representation.
      * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
+     * @param standardise       Standardise the output regressors before fitting the model.
      */
     public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOptimiser optimiser, int epochs,
                                int loggingInterval, int minibatchSize, long seed,
-                               int factorizedDimSize, double l2, double variance) {
+                               int factorizedDimSize, double l2, double variance, boolean standardise) {
         super(optimiser, epochs, loggingInterval, minibatchSize, seed, factorizedDimSize, l2, variance);
         this.objective = objective;
+        this.standardise = standardise;
     }
 
     /**
@@ -79,11 +85,12 @@ public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOpti
      * @param factorizedDimSize Size of the factorized feature representation.
      * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
+     * @param standardise       Standardise the output regressors before fitting the model.
      */
     public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOptimiser optimiser, int epochs,
                                int loggingInterval, long seed,
-                               int factorizedDimSize, double l2, double variance) {
-        this(objective,optimiser,epochs,loggingInterval,1,seed,factorizedDimSize,l2,variance);
+                               int factorizedDimSize, double l2, double variance, boolean standardise) {
+        this(objective,optimiser,epochs,loggingInterval,1,seed,factorizedDimSize,l2,variance,standardise);
     }
 
     /**
@@ -97,10 +104,11 @@ public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOpti
      * @param factorizedDimSize Size of the factorized feature representation.
      * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
+     * @param standardise       Standardise the output regressors before fitting the model.
      */
     public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOptimiser optimiser, int epochs,
-                               long seed, int factorizedDimSize, double l2, double variance) {
-        this(objective,optimiser,epochs,1000,1,seed,factorizedDimSize,l2,variance);
+                               long seed, int factorizedDimSize, double l2, double variance, boolean standardise) {
+        this(objective,optimiser,epochs,1000,1,seed,factorizedDimSize,l2,variance,standardise);
     }
 
     /**
@@ -112,10 +120,15 @@ private FMRegressionTrainer() {
 
     @Override
     protected DenseVector getTarget(ImmutableOutputInfo<Regressor> outputInfo, Regressor output) {
+        ImmutableRegressionInfo regressionInfo = (ImmutableRegressionInfo) outputInfo;
         double[] regressorsBuffer = new double[outputInfo.size()];
         for (Regressor.DimensionTuple r : output) {
             int id = outputInfo.getID(r);
-            regressorsBuffer[id] = r.getValue();
+            double curValue = r.getValue();
+            if (standardise) {
+                curValue = (curValue - regressionInfo.getMean(id)) / regressionInfo.getVariance(id);
+            }
+            regressorsBuffer[id] = curValue;
         }
         return DenseVector.createDenseVector(regressorsBuffer);
     }
@@ -132,7 +145,7 @@ protected FMRegressionModel createModel(String name, ModelProvenance provenance,
             int id = outputInfo.getID(r);
             dimensionNames[id] = r.getNames()[0];
         }
-        return new FMRegressionModel(name,dimensionNames,provenance,featureMap,outputInfo,parameters);
+        return new FMRegressionModel(name,dimensionNames,provenance,featureMap,outputInfo,parameters,standardise);
     }
 
     @Override
@@ -145,6 +158,6 @@ public String toString() {
         return "FMRegressionTrainer(objective=" + objective.toString() + ",optimiser=" + optimiser.toString() +
                 ",epochs=" + epochs + ",minibatchSize=" + minibatchSize + ",seed=" + seed +
                 ",factorizedDimSize=" + factorizedDimSize + ",l2=" + l2 + ",variance=" + variance +
-                ")";
+                ",standardise=" + standardise + ")";
     }
 }
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
index a289d4bfa..3b558d226 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
@@ -91,6 +91,8 @@ public String getOptionsDescription() {
         public double l2 = 0.001;
         @Option(longName = "variance", usage = "Variance of the initialization gaussian.")
         public double variance = 0.1;
+        @Option(longName = "standardise", usage = "Standardise the output regressors before model fitting.")
+        public boolean standardise = false;
     }
 
     /**
@@ -143,7 +145,9 @@ public static void main(String[] args) throws IOException {
         Dataset<Regressor> train = data.getA();
         Dataset<Regressor> test = data.getB();
 
-        Trainer<Regressor> trainer = new FMRegressionTrainer(obj,grad,o.epochs,o.loggingInterval,o.minibatchSize,o.general.seed,o.factorSize,o.l2,o.variance);
+        logger.info("Feature domain - " + train.getFeatureIDMap());
+
+        Trainer<Regressor> trainer = new FMRegressionTrainer(obj,grad,o.epochs,o.loggingInterval,o.minibatchSize,o.general.seed,o.factorSize,o.l2,o.variance,o.standardise);
         logger.info("Training using " + trainer.toString());
         final long trainStart = System.currentTimeMillis();
         Model<Regressor> model = trainer.train(train);

From 6c2b17f84313251d5e520b045973961f00ac0ef1 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 30 May 2021 22:51:23 -0400
Subject: [PATCH 07/31] Adding feature scaling to DataOptions for easier
 testing of factorization machines.

---
 .../java/org/tribuo/data/DataOptions.java     | 26 +++++++++++++++----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/Data/src/main/java/org/tribuo/data/DataOptions.java b/Data/src/main/java/org/tribuo/data/DataOptions.java
index a897b54f0..88507ad94 100644
--- a/Data/src/main/java/org/tribuo/data/DataOptions.java
+++ b/Data/src/main/java/org/tribuo/data/DataOptions.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2015-2021, Oracle and/or its affiliates. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -36,6 +36,9 @@
 import org.tribuo.data.text.impl.TokenPipeline;
 import org.tribuo.dataset.MinimumCardinalityDataset;
 import org.tribuo.datasource.LibSVMDataSource;
+import org.tribuo.transform.TransformationMap;
+import org.tribuo.transform.TransformerMap;
+import org.tribuo.transform.transformations.LinearScalingTransformation;
 import org.tribuo.util.tokens.impl.BreakIteratorTokenizer;
 
 import java.io.BufferedInputStream;
@@ -45,6 +48,7 @@
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.nio.file.Path;
+import java.util.Collections;
 import java.util.Locale;
 import java.util.logging.Logger;
 
@@ -63,11 +67,11 @@ public enum InputFormat {
          */
         SERIALIZED,
         /**
-         * LibSVM format data.
+         * LibSVM/svm-light format data.
          */
         LIBSVM,
         /**
-         * Text data in Tribuo's default text format "output ## text".
+         * Text data in Tribuo's standard format (i.e., each line is "output ## text data").
          */
         TEXT,
         /**
@@ -75,7 +79,7 @@ public enum InputFormat {
          */
         CSV,
         /**
-         * A CSV file using a {@link RowProcessor}.
+         * A CSV file parsed using a configured {@link RowProcessor}.
          */
         COLUMNAR
     }
@@ -177,6 +181,8 @@ public String getOptionsDescription() {
      */
     @Option(charName = 'v', longName = "testing-file", usage = "Path to the testing file.")
     public Path testingPath;
+    @Option(longName="scale-features",usage="Scales the features to the range 0-1 independently.")
+    public boolean scaleFeatures;
 
     /**
      * Loads the datasets specified in this options.
@@ -285,7 +291,16 @@ public <T extends Output<T>> Pair<Dataset<T>, Dataset<T>> load(OutputFactory<T>
                 throw new IllegalArgumentException("Unsupported input format " + inputFormat);
         }
         logger.info(String.format("Loaded %d testing examples", test.size()));
-        return new Pair<>(train, test);
+        if (scaleFeatures) {
+            logger.info("Fitting feature scaling");
+            TransformationMap map = new TransformationMap(Collections.singletonList(new LinearScalingTransformation()));
+            TransformerMap transformers = train.createTransformers(map);
+            logger.info("Applying scaling to training dataset");
+            train = transformers.transformDataset(train);
+            logger.info("Applying scaling to testing dataset");
+            test = transformers.transformDataset(test);
+        }
+        return new Pair<>(train,test);
     }
 
     /**
@@ -301,3 +316,4 @@ public <T extends Output<T>> void saveModel(Model<T> model) throws IOException {
         }
     }
 }
+

From aa563f15781beaaab04b5c48c1d27917a77ff65b Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 1 Jun 2021 21:40:07 -0400
Subject: [PATCH 08/31] Adding additional fine level logging to
 AbstractSGDTrainer for debugging.

---
 .../main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java
index ffde93153..9f45905ba 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java
@@ -161,17 +161,23 @@ public V train(Dataset<T> examples, Map<String, Provenance> runProvenance) {
         U[] sgdTargets = (U[]) new Object[examples.size()];
         double[] weights = new double[examples.size()];
         int n = 0;
+        long featureSize = 0;
+        long denseCount = 0;
         for (Example<T> example : examples) {
             weights[n] = example.getWeight();
             if (example.size() == featureSpaceSize) {
                 sgdFeatures[n] = DenseVector.createDenseVector(example, featureIDMap, addBias);
+                denseCount++;
             } else {
                 sgdFeatures[n] = SparseVector.createSparseVector(example, featureIDMap, addBias);
             }
             sgdTargets[n] = getTarget(outputIDInfo,example.getOutput());
+            featureSize += sgdFeatures[n].numActiveElements();
             n++;
         }
         logger.info(String.format("Training SGD model with %d examples", n));
+        logger.fine("Average feature size = " + featureSize / (double)n);
+        logger.fine("Dense count = " + denseCount);
         logger.info("Outputs - " + outputIDInfo.toReadableString());
 
         X parameters = createParameters(featureIDMap.size(), outputIDInfo.size(), localRNG);

From 2cb5d4c354d8befc4f400926eeab1c5f823dbf1d Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 1 Jun 2021 21:40:44 -0400
Subject: [PATCH 09/31] Fixing a bug in ArrayExample.transform which reduced
 the effectiveness of a performance optimization.

---
 Core/src/main/java/org/tribuo/impl/ArrayExample.java | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Core/src/main/java/org/tribuo/impl/ArrayExample.java b/Core/src/main/java/org/tribuo/impl/ArrayExample.java
index 05dcb0181..015355f59 100644
--- a/Core/src/main/java/org/tribuo/impl/ArrayExample.java
+++ b/Core/src/main/java/org/tribuo/impl/ArrayExample.java
@@ -428,7 +428,7 @@ public void set(Feature feature) {
 
     @Override
     public void transform(TransformerMap transformerMap) {
-        if(transformerMap.size() < featureNames.length) {
+        if(transformerMap.size() < size) {
             //
             // We have fewer transformers than feature names, so let's
             // iterate through the map and find the features.

From 23d28451f10e597145187c2d2d577f69052e7438 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 1 Jun 2021 21:54:52 -0400
Subject: [PATCH 10/31] Adding another rescaling option to DataOptions.

---
 .../java/org/tribuo/data/DataOptions.java     | 24 +++++++++++++++----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/Data/src/main/java/org/tribuo/data/DataOptions.java b/Data/src/main/java/org/tribuo/data/DataOptions.java
index 88507ad94..ad494e3e4 100644
--- a/Data/src/main/java/org/tribuo/data/DataOptions.java
+++ b/Data/src/main/java/org/tribuo/data/DataOptions.java
@@ -183,8 +183,11 @@ public String getOptionsDescription() {
     public Path testingPath;
     @Option(longName="scale-features",usage="Scales the features to the range 0-1 independently.")
     public boolean scaleFeatures;
+    @Option(longName="scale-including-zeros",usage="Includes implicit zeros in the scale range calculation")
+    public boolean scaleIncZeros;
 
     /**
+<<<<<<< HEAD
      * Loads the datasets specified in this options.
      * @param outputFactory The output factory to use.
      * @param <T> The type of the dataset.
@@ -192,6 +195,17 @@ public String getOptionsDescription() {
      * @throws IOException If the datasets failed to load.
      */
     public <T extends Output<T>> Pair<Dataset<T>, Dataset<T>> load(OutputFactory<T> outputFactory) throws IOException {
+=======
+     * Loads the training and testing data from {@link #trainingPath} and {@link #testingPath}
+     * according to the other parameters specified in this class.
+     * @param outputFactory The output factory to use to process the inputs.
+     * @param <T> The dataset output type.
+     * @return A pair containing the training and testing datasets. The training dataset is element 'A' and the
+     * testing dataset is element 'B'.
+     * @throws IOException If the paths could not be loaded.
+     */
+    public <T extends Output<T>> Pair<Dataset<T>,Dataset<T>> load(OutputFactory<T> outputFactory) throws IOException {
+>>>>>>> c459390 (Adding another rescaling option to DataOptions.)
         logger.info(String.format("Loading data from %s", trainingPath));
         Dataset<T> train;
         Dataset<T> test;
@@ -294,7 +308,7 @@ public <T extends Output<T>> Pair<Dataset<T>, Dataset<T>> load(OutputFactory<T>
         if (scaleFeatures) {
             logger.info("Fitting feature scaling");
             TransformationMap map = new TransformationMap(Collections.singletonList(new LinearScalingTransformation()));
-            TransformerMap transformers = train.createTransformers(map);
+            TransformerMap transformers = train.createTransformers(map,scaleIncZeros);
             logger.info("Applying scaling to training dataset");
             train = transformers.transformDataset(train);
             logger.info("Applying scaling to testing dataset");
@@ -304,14 +318,14 @@ public <T extends Output<T>> Pair<Dataset<T>, Dataset<T>> load(OutputFactory<T>
     }
 
     /**
-     * Saves the model out to the path specified in this options.
+     * Saves the model out to the path in {@link #outputPath}.
      * @param model The model to save.
-     * @param <T> The type of the model output.
+     * @param <T> The model's output type.
      * @throws IOException If the model could not be saved.
      */
     public <T extends Output<T>> void saveModel(Model<T> model) throws IOException {
-        try (ObjectOutputStream objOut = new ObjectOutputStream(new FileOutputStream(outputPath.toFile()))) {
-            objOut.writeObject(model);
+        try (ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(outputPath.toFile()))) {
+            oos.writeObject(model);
             logger.info("Serialized model to file: " + outputPath);
         }
     }

From d9e0083ae4adf6de17970ab892f60ff45c16d8d6 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Fri, 4 Jun 2021 12:06:04 -0400
Subject: [PATCH 11/31] Adding copy accessors for the FM model parameters.

---
 .../tribuo/common/sgd/AbstractFMModel.java    | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
index 7b1b1d4f5..5af86c41b 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
@@ -24,6 +24,7 @@
 import org.tribuo.Output;
 import org.tribuo.math.la.DenseMatrix;
 import org.tribuo.math.la.DenseVector;
+import org.tribuo.math.la.Tensor;
 import org.tribuo.provenance.ModelProvenance;
 
 import java.util.ArrayList;
@@ -110,6 +111,35 @@ public Map<String, List<Pair<String, Double>>> getTopFeatures(int n) {
         return map;
     }
 
+    /**
+     * Returns a copy of the linear weights.
+     * @return The linear weights.
+     */
+    public DenseMatrix getLinearWeightsCopy() {
+        return ((DenseMatrix)modelParameters.get()[1]).copy();
+    }
+
+    /**
+     * Returns a copy of the output dimension biases.
+     * @return The biases.
+     */
+    public DenseVector getBiasesCopy() {
+        return ((DenseVector)modelParameters.get()[0]).copy();
+    }
+
+    /**
+     * Returns a copy of the factors. There is one factor matrix per output dimension.
+     * @return The factors.
+     */
+    public Tensor[] getFactorsCopy() {
+        Tensor[] params = modelParameters.get();
+        Tensor[] paramCopy = new Tensor[params.length-2];
+        for (int i = 0; i < paramCopy.length; i++) {
+            paramCopy[i] = params[i+2].copy();
+        }
+        return paramCopy;
+    }
+
     /**
      * Factorization machines don't provide excuses, use an explainer.
      * @param example The input example.

From fd16538c39375e6da82bcf604c2da6f734ddf30e Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Fri, 4 Jun 2021 12:07:27 -0400
Subject: [PATCH 12/31] Adding smoke tests for FM multi-label and regression.

---
 .../tribuo/common/sgd/AbstractFMModel.java    |   5 +-
 .../multilabel/sgd/fm/TestFMMultiLabel.java   |  85 ++++++++++++
 .../regression/sgd/fm/TestFMRegression.java   | 122 ++++++++++++++++++
 3 files changed, 211 insertions(+), 1 deletion(-)
 create mode 100644 MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
 create mode 100644 Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java

diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
index 5af86c41b..bec2675f6 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
@@ -128,7 +128,10 @@ public DenseVector getBiasesCopy() {
     }
 
     /**
-     * Returns a copy of the factors. There is one factor matrix per output dimension.
+     * Returns a copy of the factors.
+     * There is one factor matrix per output dimension.
+     * The first factor matrix dimension is the factor dimension,
+     * the second is the number of features.
      * @return The factors.
      */
     public Tensor[] getFactorsCopy() {
diff --git a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
new file mode 100644
index 000000000..dbed244d3
--- /dev/null
+++ b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.multilabel.sgd.fm;
+
+import com.oracle.labs.mlrg.olcut.util.Pair;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.tribuo.Dataset;
+import org.tribuo.Model;
+import org.tribuo.Prediction;
+import org.tribuo.Trainer;
+import org.tribuo.common.sgd.AbstractFMTrainer;
+import org.tribuo.common.sgd.AbstractSGDTrainer;
+import org.tribuo.math.optimisers.AdaGrad;
+import org.tribuo.multilabel.MultiLabel;
+import org.tribuo.multilabel.evaluation.MultiLabelEvaluation;
+import org.tribuo.multilabel.example.MultiLabelDataGenerator;
+import org.tribuo.multilabel.sgd.objectives.BinaryCrossEntropy;
+import org.tribuo.multilabel.sgd.objectives.Hinge;
+import org.tribuo.test.Helpers;
+
+import java.util.List;
+import java.util.Map;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class TestFMMultiLabel {
+
+    private static final FMMultiLabelTrainer hinge = new FMMultiLabelTrainer(new Hinge(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED, 5,0.0,0.1);
+    private static final FMMultiLabelTrainer sigmoid = new FMMultiLabelTrainer(new BinaryCrossEntropy(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED,5,0.0,0.1);
+
+    @BeforeAll
+    public static void setup() {
+        Class<?>[] classes = new Class<?>[]{AbstractSGDTrainer.class, AbstractFMTrainer.class,FMMultiLabelTrainer.class};
+        for (Class<?> c : classes) {
+            Logger logger = Logger.getLogger(c.getName());
+            logger.setLevel(Level.WARNING);
+        }
+    }
+
+    @Test
+    public void testPredictions() {
+        Dataset<MultiLabel> train = MultiLabelDataGenerator.generateTrainData();
+        Dataset<MultiLabel> test = MultiLabelDataGenerator.generateTestData();
+
+        testTrainer(train,test,hinge);
+        testTrainer(train,test,sigmoid);
+    }
+
+    private static void testTrainer(Dataset<MultiLabel> train, Dataset<MultiLabel> test, Trainer<MultiLabel> trainer) {
+        Model<MultiLabel> model = trainer.train(train);
+
+        List<Prediction<MultiLabel>> predictions = model.predict(test);
+        Prediction<MultiLabel> first = predictions.get(0);
+        MultiLabel trueLabel = train.getOutputFactory().generateOutput("MONKEY,PUZZLE,TREE");
+        assertEquals(trueLabel, first.getOutput(), "Predicted labels not equal");
+        Map<String, List<Pair<String, Double>>> features = model.getTopFeatures(2);
+        Assertions.assertNotNull(features);
+        Assertions.assertFalse(features.isEmpty());
+
+        MultiLabelEvaluation evaluation = (MultiLabelEvaluation) train.getOutputFactory().getEvaluator().evaluate(model,test);
+
+        Assertions.assertEquals(1.0,evaluation.microAveragedRecall());
+
+        Helpers.testModelSerialization(model, MultiLabel.class);
+    }
+
+}
diff --git a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
new file mode 100644
index 000000000..9a5810eec
--- /dev/null
+++ b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
@@ -0,0 +1,122 @@
+/*
+ * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tribuo.regression.sgd.fm;
+
+import com.oracle.labs.mlrg.olcut.util.Pair;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.tribuo.Dataset;
+import org.tribuo.Model;
+import org.tribuo.Trainer;
+import org.tribuo.common.sgd.AbstractFMModel;
+import org.tribuo.common.sgd.AbstractFMTrainer;
+import org.tribuo.common.sgd.AbstractSGDTrainer;
+import org.tribuo.math.la.DenseMatrix;
+import org.tribuo.math.la.DenseVector;
+import org.tribuo.math.optimisers.AdaGrad;
+import org.tribuo.regression.Regressor;
+import org.tribuo.regression.evaluation.RegressionEvaluation;
+import org.tribuo.regression.evaluation.RegressionEvaluator;
+import org.tribuo.regression.example.RegressionDataGenerator;
+import org.tribuo.regression.sgd.objectives.SquaredLoss;
+import org.tribuo.test.Helpers;
+
+import java.util.List;
+import java.util.Map;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class TestFMRegression {
+
+    private static final FMRegressionTrainer t = new FMRegressionTrainer(new SquaredLoss(), new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED,1,0.0,0.1,true);
+
+    private static final RegressionEvaluator e = new RegressionEvaluator();
+
+    @BeforeAll
+    public static void setup() {
+        Class<?>[] classes = new Class<?>[]{AbstractSGDTrainer.class, AbstractFMTrainer.class,FMRegressionTrainer.class};
+        for (Class<?> c : classes) {
+            Logger logger = Logger.getLogger(c.getName());
+            logger.setLevel(Level.WARNING);
+        }
+    }
+
+    public static Model<Regressor> testFMRegression(Pair<Dataset<Regressor>,Dataset<Regressor>> p) {
+        Model<Regressor> m = t.train(p.getA());
+        RegressionEvaluation evaluation = e.evaluate(m,p.getB());
+        Map<String, List<Pair<String,Double>>> features = m.getTopFeatures(3);
+        Assertions.assertNotNull(features);
+        Assertions.assertFalse(features.isEmpty());
+        features = m.getTopFeatures(-1);
+        Assertions.assertNotNull(features);
+        Assertions.assertFalse(features.isEmpty());
+        return m;
+    }
+
+    @Test
+    public void testDenseData() {
+        Pair<Dataset<Regressor>,Dataset<Regressor>> p = RegressionDataGenerator.denseTrainTest();
+        Model<Regressor> model = testFMRegression(p);
+        Helpers.testModelSerialization(model,Regressor.class);
+    }
+
+    @Test
+    public void testSparseData() {
+        Pair<Dataset<Regressor>,Dataset<Regressor>> p = RegressionDataGenerator.sparseTrainTest();
+        testFMRegression(p);
+    }
+
+    @Test
+    public void testInvalidExample() {
+        assertThrows(IllegalArgumentException.class, () -> {
+            Pair<Dataset<Regressor>, Dataset<Regressor>> p = RegressionDataGenerator.denseTrainTest();
+            Model<Regressor> m = t.train(p.getA());
+            m.predict(RegressionDataGenerator.invalidSparseExample());
+        });
+    }
+
+    @Test
+    public void testEmptyExample() {
+        assertThrows(IllegalArgumentException.class, () -> {
+            Pair<Dataset<Regressor>, Dataset<Regressor>> p = RegressionDataGenerator.denseTrainTest();
+            Model<Regressor> m = t.train(p.getA());
+            m.predict(RegressionDataGenerator.emptyExample());
+        });
+    }
+
+    @Test
+    public void testMultiInvalidExample() {
+        assertThrows(IllegalArgumentException.class, () -> {
+            Pair<Dataset<Regressor>, Dataset<Regressor>> p = RegressionDataGenerator.multiDimDenseTrainTest();
+            Model<Regressor> m = t.train(p.getA());
+            m.predict(RegressionDataGenerator.invalidMultiDimSparseExample());
+        });
+    }
+
+    @Test
+    public void testMultiEmptyExample() {
+        assertThrows(IllegalArgumentException.class, () -> {
+            Pair<Dataset<Regressor>, Dataset<Regressor>> p = RegressionDataGenerator.multiDimDenseTrainTest();
+            Model<Regressor> m = t.train(p.getA());
+            m.predict(RegressionDataGenerator.emptyMultiDimExample());
+        });
+    }
+}

From 49d6dbcd226b20a3a4ab1e78f5667301e1b19b61 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Wed, 7 Jul 2021 15:41:58 -0400
Subject: [PATCH 13/31] Adding missing override annotations to various methods
 in SGD and Math.

---
 .../src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java  | 2 ++
 .../java/org/tribuo/common/sgd/AbstractLinearSGDTrainer.java    | 2 ++
 Math/src/main/java/org/tribuo/math/la/DenseMatrix.java          | 1 +
 Math/src/main/java/org/tribuo/math/la/DenseVector.java          | 1 +
 Math/src/main/java/org/tribuo/math/la/Matrix.java               | 1 +
 Math/src/main/java/org/tribuo/math/la/SGDVector.java            | 1 +
 6 files changed, 8 insertions(+)

diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
index 77a0b1965..3067cb78f 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
@@ -93,6 +93,7 @@ public void postConfig() {
      * Returns the default model name.
      * @return The default model name.
      */
+    @Override
     protected String getName() {
         return "factorization-machine-model";
     }
@@ -106,6 +107,7 @@ protected String getName() {
      * @param localRNG The RNG to use for parameter initialisation.
      * @return The trainable parameters.
      */
+    @Override
     protected FMParameters createParameters(int numFeatures, int numOutputs, SplittableRandom localRNG) {
         return new FMParameters(localRNG,numFeatures,numOutputs,factorizedDimSize,l2,variance);
     }
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDTrainer.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDTrainer.java
index 9bf1908a8..7d0ed6288 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDTrainer.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDTrainer.java
@@ -63,6 +63,7 @@ protected AbstractLinearSGDTrainer() {
      * Returns the default model name.
      * @return The default model name.
      */
+    @Override
     protected String getName() {
         return "linear-sgd-model";
     }
@@ -75,6 +76,7 @@ protected String getName() {
      * @param localRNG The RNG to use for parameter initialisation.
      * @return The trainable parameters.
      */
+    @Override
     protected LinearParameters createParameters(int numFeatures, int numOutputs, SplittableRandom localRNG) {
         return new LinearParameters(numFeatures+1,numOutputs);
     }
diff --git a/Math/src/main/java/org/tribuo/math/la/DenseMatrix.java b/Math/src/main/java/org/tribuo/math/la/DenseMatrix.java
index 95c9b4bf6..9e4b306ad 100644
--- a/Math/src/main/java/org/tribuo/math/la/DenseMatrix.java
+++ b/Math/src/main/java/org/tribuo/math/la/DenseMatrix.java
@@ -162,6 +162,7 @@ public Tensor reshape(int[] newShape) {
      * Copies the matrix.
      * @return A deep copy of the matrix.
      */
+    @Override
     public DenseMatrix copy() {
         return new DenseMatrix(this);
     }
diff --git a/Math/src/main/java/org/tribuo/math/la/DenseVector.java b/Math/src/main/java/org/tribuo/math/la/DenseVector.java
index 32f7d0c77..29b4128b2 100644
--- a/Math/src/main/java/org/tribuo/math/la/DenseVector.java
+++ b/Math/src/main/java/org/tribuo/math/la/DenseVector.java
@@ -131,6 +131,7 @@ public static <T extends Output<T>> DenseVector createDenseVector(Example<T> exa
      * get function has been modified.
      * @return A copy of the values in this DenseVector.
      */
+    @Override
     public double[] toArray() {
         return Arrays.copyOf(elements, elements.length);
     }
diff --git a/Math/src/main/java/org/tribuo/math/la/Matrix.java b/Math/src/main/java/org/tribuo/math/la/Matrix.java
index 44a52a6a7..34b4fbd79 100644
--- a/Math/src/main/java/org/tribuo/math/la/Matrix.java
+++ b/Math/src/main/java/org/tribuo/math/la/Matrix.java
@@ -27,6 +27,7 @@ public interface Matrix extends Tensor, Iterable<MatrixTuple> {
      * Copies the matrix.
      * @return A copy of the matrix.
      */
+    @Override
     public Matrix copy();
 
     /**
diff --git a/Math/src/main/java/org/tribuo/math/la/SGDVector.java b/Math/src/main/java/org/tribuo/math/la/SGDVector.java
index 31d45debb..308d50828 100644
--- a/Math/src/main/java/org/tribuo/math/la/SGDVector.java
+++ b/Math/src/main/java/org/tribuo/math/la/SGDVector.java
@@ -41,6 +41,7 @@ public interface SGDVector extends Tensor, Iterable<VectorTuple> {
      * Returns a deep copy of this vector.
      * @return A copy of this vector.
      */
+    @Override
     public SGDVector copy();
 
     /**

From 34e4c30753242607ba34fe8f3f9e7cdee19ae652 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Wed, 7 Jul 2021 16:02:52 -0400
Subject: [PATCH 14/31] Adding missing override annotations to other classes.

---
 .../tribuo/regression/rtree/IndependentRegressionTreeModel.java  | 1 +
 .../java/org/tribuo/util/tokens/impl/SplitFunctionTokenizer.java | 1 +
 2 files changed, 2 insertions(+)

diff --git a/Regression/RegressionTree/src/main/java/org/tribuo/regression/rtree/IndependentRegressionTreeModel.java b/Regression/RegressionTree/src/main/java/org/tribuo/regression/rtree/IndependentRegressionTreeModel.java
index 963a9087a..8c3bacc1f 100644
--- a/Regression/RegressionTree/src/main/java/org/tribuo/regression/rtree/IndependentRegressionTreeModel.java
+++ b/Regression/RegressionTree/src/main/java/org/tribuo/regression/rtree/IndependentRegressionTreeModel.java
@@ -90,6 +90,7 @@ private static Map<String,List<String>> gatherActiveFeatures(ImmutableFeatureMap
      * Probes the trees to find the depth.
      * @return The maximum depth across the trees.
      */
+    @Override
     public int getDepth() {
         int maxDepth = 0;
         for (Node<Regressor> curRoot : roots.values()) {
diff --git a/Util/Tokenization/src/main/java/org/tribuo/util/tokens/impl/SplitFunctionTokenizer.java b/Util/Tokenization/src/main/java/org/tribuo/util/tokens/impl/SplitFunctionTokenizer.java
index b87d2f5b5..68c9d8ef8 100644
--- a/Util/Tokenization/src/main/java/org/tribuo/util/tokens/impl/SplitFunctionTokenizer.java
+++ b/Util/Tokenization/src/main/java/org/tribuo/util/tokens/impl/SplitFunctionTokenizer.java
@@ -386,6 +386,7 @@ public TokenType getType() {
         return currentToken.type;
     }
 
+    @Override
     public Tokenizer clone() throws CloneNotSupportedException {
         throw new UnsupportedOperationException(
                 "abstract class SplitFunctionTokenizer does not implement clone method.  Subclasses must implement this method.");

From e0467b6ac4b508ba7bf7ba35f8593a57a9e2b114 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 7 Sep 2021 17:59:14 -0400
Subject: [PATCH 15/31] Removing the unimplemented l2 regularisation hooks.
 It's been pushed to the next release.

---
 .../sgd/fm/FMClassificationOptions.java       |  4 +-
 .../sgd/fm/FMClassificationTrainer.java       | 19 +++---
 .../sgd/fm/TestFMClassification.java          |  5 +-
 .../tribuo/common/sgd/AbstractFMTrainer.java  | 39 ++++++------
 .../org/tribuo/common/sgd/FMParameters.java   | 61 ++++++++++---------
 .../sgd/fm/FMMultiLabelOptions.java           |  4 +-
 .../sgd/fm/FMMultiLabelTrainer.java           | 56 ++++++++---------
 .../multilabel/sgd/fm/TestFMMultiLabel.java   |  6 +-
 .../sgd/fm/FMRegressionTrainer.java           | 58 +++++++++---------
 .../tribuo/regression/sgd/fm/TrainTest.java   |  2 +-
 .../regression/sgd/fm/TestFMRegression.java   |  2 +-
 11 files changed, 122 insertions(+), 134 deletions(-)

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
index 0825686f7..f31739094 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
@@ -61,8 +61,6 @@ public enum LossEnum {
     private long fmSeed = Trainer.DEFAULT_SEED;
     @Option(longName = "fm-factor-size", usage = "Factor size.")
     public int fmFactorSize = 6;
-    @Option(longName = "fm-l2-penalty", usage = "L2 regularization penalty.")
-    public double fmL2 = 0.001;
     @Option(longName = "fm-variance", usage = "Variance of the initialization gaussian.")
     public double fmVariance = 0.1;
 
@@ -85,6 +83,6 @@ public LabelObjective getLoss() {
     public FMClassificationTrainer getTrainer() {
         logger.info(String.format("Set logging interval to %d", fmLoggingInterval));
         return new FMClassificationTrainer(getLoss(), sgoOptions.getOptimiser(), fmEpochs, fmLoggingInterval,
-                fmMinibatchSize, fmSeed, fmFactorSize, fmL2, fmVariance);
+                fmMinibatchSize, fmSeed, fmFactorSize, fmVariance);
     }
 }
diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationTrainer.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationTrainer.java
index 0884d6f04..52372e8d5 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationTrainer.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationTrainer.java
@@ -40,7 +40,7 @@
  * 2010 IEEE International Conference on Data Mining
  * </pre>
  */
-public class FMClassificationTrainer extends AbstractFMTrainer<Label,Integer> {
+public class FMClassificationTrainer extends AbstractFMTrainer<Label, Integer> {
     private static final Logger logger = Logger.getLogger(FMClassificationTrainer.class.getName());
 
     @Config(description = "The classification objective function to use.")
@@ -56,13 +56,12 @@ public class FMClassificationTrainer extends AbstractFMTrainer<Label,Integer> {
      * @param minibatchSize     The size of any minibatches.
      * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
      * @param factorizedDimSize Size of the factorized feature representation.
-     * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
      */
     public FMClassificationTrainer(LabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
                                    int loggingInterval, int minibatchSize, long seed,
-                                   int factorizedDimSize, double l2, double variance) {
-        super(optimiser, epochs, loggingInterval, minibatchSize, seed, factorizedDimSize, l2, variance);
+                                   int factorizedDimSize, double variance) {
+        super(optimiser, epochs, loggingInterval, minibatchSize, seed, factorizedDimSize, variance);
         this.objective = objective;
     }
 
@@ -77,13 +76,12 @@ public FMClassificationTrainer(LabelObjective objective, StochasticGradientOptim
      * @param loggingInterval   Log the loss after this many iterations. If -1 don't log anything.
      * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
      * @param factorizedDimSize Size of the factorized feature representation.
-     * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
      */
     public FMClassificationTrainer(LabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
                                    int loggingInterval, long seed,
-                                   int factorizedDimSize, double l2, double variance) {
-        this(objective, optimiser, epochs, loggingInterval, 1, seed, factorizedDimSize, l2, variance);
+                                   int factorizedDimSize, double variance) {
+        this(objective, optimiser, epochs, loggingInterval, 1, seed, factorizedDimSize, variance);
     }
 
     /**
@@ -96,12 +94,11 @@ public FMClassificationTrainer(LabelObjective objective, StochasticGradientOptim
      * @param epochs            The number of epochs (complete passes through the training data).
      * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
      * @param factorizedDimSize Size of the factorized feature representation.
-     * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
      */
     public FMClassificationTrainer(LabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
-                                   long seed, int factorizedDimSize, double l2, double variance) {
-        this(objective, optimiser, epochs, 1000, 1, seed, factorizedDimSize, l2, variance);
+                                   long seed, int factorizedDimSize, double variance) {
+        this(objective, optimiser, epochs, 1000, 1, seed, factorizedDimSize, variance);
     }
 
     /**
@@ -135,7 +132,7 @@ protected String getModelClassName() {
     public String toString() {
         return "FMClassificationTrainer(objective=" + objective.toString() + ",optimiser=" + optimiser.toString() +
                 ",epochs=" + epochs + ",minibatchSize=" + minibatchSize + ",seed=" + seed +
-                ",factorizedDimSize=" + factorizedDimSize + ",l2=" + l2 + ",variance=" + variance +
+                ",factorizedDimSize=" + factorizedDimSize + ",variance=" + variance +
                 ")";
     }
 }
diff --git a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
index 9839b124a..3db6682b1 100644
--- a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
+++ b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
@@ -21,7 +21,6 @@
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.tribuo.Dataset;
-import org.tribuo.Example;
 import org.tribuo.Model;
 import org.tribuo.Trainer;
 import org.tribuo.classification.Label;
@@ -31,7 +30,6 @@
 import org.tribuo.classification.sgd.objectives.Hinge;
 import org.tribuo.common.sgd.AbstractFMTrainer;
 import org.tribuo.common.sgd.AbstractSGDTrainer;
-import org.tribuo.dataset.DatasetView;
 import org.tribuo.math.optimisers.AdaGrad;
 import org.tribuo.test.Helpers;
 
@@ -40,13 +38,12 @@
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class TestFMClassification {
 
     private static final FMClassificationTrainer t = new FMClassificationTrainer(new Hinge(),
-            new AdaGrad(0.1,0.1),5,1000,Trainer.DEFAULT_SEED,6,0.0,0.1);
+            new AdaGrad(0.1,0.1),5,1000,Trainer.DEFAULT_SEED,6,0.1);
 
     @BeforeAll
     public static void setup() {
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
index 3067cb78f..07f66afce 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
@@ -36,34 +36,30 @@
  * 2010 IEEE International Conference on Data Mining
  * </pre>
  */
-public abstract class AbstractFMTrainer<T extends Output<T>,U> extends AbstractSGDTrainer<T,U,AbstractFMModel<T>, FMParameters> {
+public abstract class AbstractFMTrainer<T extends Output<T>, U> extends AbstractSGDTrainer<T, U, AbstractFMModel<T>, FMParameters> {
     private static final Logger logger = Logger.getLogger(AbstractFMTrainer.class.getName());
 
     @Config(mandatory = true, description = "The size of the factorized feature representation.")
     protected int factorizedDimSize;
 
-    @Config(mandatory = true, description = "The l2 regularization penalty (i.e., weight decay)")
-    protected double l2;
-
     @Config(mandatory = true, description = "The variance of the initializer.")
     protected double variance;
 
     /**
      * Constructs an SGD trainer for a factorization machine.
-     * @param optimiser The gradient optimiser to use.
-     * @param epochs The number of epochs (complete passes through the training data).
-     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
-     * @param minibatchSize The size of any minibatches.
-     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     *
+     * @param optimiser         The gradient optimiser to use.
+     * @param epochs            The number of epochs (complete passes through the training data).
+     * @param loggingInterval   Log the loss after this many iterations. If -1 don't log anything.
+     * @param minibatchSize     The size of any minibatches.
+     * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
      * @param factorizedDimSize Size of the factorized feature representation.
-     * @param l2 The l2 regularisation penalty.
-     * @param variance The variance of the initializer.
+     * @param variance          The variance of the initializer.
      */
     protected AbstractFMTrainer(StochasticGradientOptimiser optimiser, int epochs, int loggingInterval,
-                                int minibatchSize, long seed, int factorizedDimSize, double l2, double variance) {
-        super(optimiser,epochs,loggingInterval,minibatchSize,seed,false);
+                                int minibatchSize, long seed, int factorizedDimSize, double variance) {
+        super(optimiser, epochs, loggingInterval, minibatchSize, seed, false);
         this.factorizedDimSize = factorizedDimSize;
-        this.l2 = l2;
         this.variance = variance;
         postConfig();
     }
@@ -79,18 +75,16 @@ protected AbstractFMTrainer() {
     public void postConfig() {
         super.postConfig();
         if (factorizedDimSize < 1) {
-            throw new PropertyException("","factorizedDimSize","Value must be positive.");
-        }
-        if (l2 < 0.0) {
-            throw new PropertyException("","l2","Value must be non-negative.");
+            throw new PropertyException("", "factorizedDimSize", "Value must be positive.");
         }
         if (variance <= 0.0) {
-            throw new PropertyException("","variance","Value must be positive.");
+            throw new PropertyException("", "variance", "Value must be positive.");
         }
     }
 
     /**
      * Returns the default model name.
+     *
      * @return The default model name.
      */
     @Override
@@ -102,14 +96,15 @@ protected String getName() {
      * Constructs the trainable parameters object, in this case a {@link FMParameters} containing
      * a weight matrix for the feature weights and a series of weight matrices for the factorized
      * feature representation.
+     *
      * @param numFeatures The number of input features.
-     * @param numOutputs The number of output dimensions.
-     * @param localRNG The RNG to use for parameter initialisation.
+     * @param numOutputs  The number of output dimensions.
+     * @param localRNG    The RNG to use for parameter initialisation.
      * @return The trainable parameters.
      */
     @Override
     protected FMParameters createParameters(int numFeatures, int numOutputs, SplittableRandom localRNG) {
-        return new FMParameters(localRNG,numFeatures,numOutputs,factorizedDimSize,l2,variance);
+        return new FMParameters(localRNG, numFeatures, numOutputs, factorizedDimSize, variance);
     }
 
 }
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
index 0da8a5d93..32ce4c831 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
@@ -47,50 +47,50 @@ public final class FMParameters implements FeedForwardParameters {
     private DenseVector biasVector;
     private DenseMatrix weightMatrix;
 
-    private final double l2;
     private final int numFactors;
 
     /**
      * Constructor. The number of features and the number of outputs must be fixed and known in advance.
-     * @param rng The RNG to use for initialization.
+     *
+     * @param rng         The RNG to use for initialization.
      * @param numFeatures The number of features in the training dataset.
-     * @param numLabels The number of outputs in the training dataset.
-     * @param numFactors The size of the factorized feature representation.
+     * @param numLabels   The number of outputs in the training dataset.
+     * @param numFactors  The size of the factorized feature representation.
+     * @param variance    The variance of the factor initializer.
      */
-    public FMParameters(SplittableRandom rng, int numFeatures, int numLabels, int numFactors, double l2, double variance) {
+    public FMParameters(SplittableRandom rng, int numFeatures, int numLabels, int numFactors, double variance) {
         weights = new Tensor[numLabels + 2];
         biasVector = new DenseVector(numLabels);
-        weightMatrix = new DenseMatrix(numLabels,numFeatures);
+        weightMatrix = new DenseMatrix(numLabels, numFeatures);
         weights[0] = biasVector;
         weights[1] = weightMatrix;
         for (int i = 0; i < numLabels; i++) {
-            DenseMatrix curMatrix = new DenseMatrix(numFactors,numFeatures);
+            DenseMatrix curMatrix = new DenseMatrix(numFactors, numFeatures);
             initializeMatrix(rng, variance, curMatrix);
-            weights[i+2] = curMatrix;
+            weights[i + 2] = curMatrix;
         }
-        this.l2 = l2;
         this.numFactors = numFactors;
     }
 
     /**
      * Constructs a FMParameters wrapped around the weight array.
-     * @param weights The weights to wrap.
+     *
+     * @param weights    The weights to wrap.
      * @param numFactors The size of the factorized feature representation.
-     * @param l2 The l2 regularization penalty.
      */
-    private FMParameters(Tensor[] weights, int numFactors, double l2) {
+    private FMParameters(Tensor[] weights, int numFactors) {
         this.weights = weights;
         this.biasVector = (DenseVector) weights[0];
         this.weightMatrix = (DenseMatrix) weights[1];
-        this.l2 = l2;
         this.numFactors = numFactors;
     }
 
     /**
      * Initializes the weight matrix by drawing from a zero mean gaussian with the specified variance.
-     * @param rng The RNG to use.
+     *
+     * @param rng      The RNG to use.
      * @param variance The variance of the gaussian.
-     * @param matrix The matrix to initialize.
+     * @param matrix   The matrix to initialize.
      */
     private void initializeMatrix(SplittableRandom rng, double variance, DenseMatrix matrix) {
         // This is to get a nextGaussian method. In Java 17 we can use the SplittableRandom directly
@@ -99,13 +99,14 @@ private void initializeMatrix(SplittableRandom rng, double variance, DenseMatrix
         int dim2 = matrix.getDimension2Size();
         for (int i = 0; i < dim1; i++) {
             for (int j = 0; j < dim2; j++) {
-                matrix.set(i,j,innerRNG.nextGaussian()*variance);
+                matrix.set(i, j, innerRNG.nextGaussian() * variance);
             }
         }
     }
 
     /**
      * Generates an unnormalised prediction by leftMultiply'ing the weights with the incoming features.
+     *
      * @param example A feature vector
      * @return A {@link DenseVector} containing a score for each label.
      */
@@ -126,15 +127,15 @@ public DenseVector predict(SGDVector example) {
                 double sumOfSquares = 0.0;
                 double sum = 0.0;
                 for (VectorTuple v : example) {
-                    double curWeight = curMatrix.get(k,v.index);
+                    double curWeight = curMatrix.get(k, v.index);
                     double value = curWeight * v.value;
                     sum += value;
-                    sumOfSquares += value*value;
+                    sumOfSquares += value * value;
                 }
-                curValue += (sum*sum) - sumOfSquares;
+                curValue += (sum * sum) - sumOfSquares;
             }
             curValue = curValue / 2;
-            factorizedPred.set(i-2,curValue);
+            factorizedPred.set(i - 2, curValue);
         }
 
         // Add factorized portion
@@ -147,7 +148,8 @@ public DenseVector predict(SGDVector example) {
      * the loss and the per output gradients.
      * <p>
      * This parameters returns a {@link Tensor} array with numLabels + 2 elements.
-     * @param score The Pair returned by the objective.
+     *
+     * @param score    The Pair returned by the objective.
      * @param features The feature vector.
      * @return A {@link Tensor} array containing all the gradients.
      */
@@ -169,7 +171,7 @@ public Tensor[] gradients(Pair<Double, SGDVector> score, SGDVector features) {
         // factorised representation gradients
         // per label
         for (int i = 2; i < weights.length; i++) {
-            double curOutputGradient = outputGradient.get(i-2);
+            double curOutputGradient = outputGradient.get(i - 2);
             DenseMatrix curFactors = (DenseMatrix) weights[i];
             if (curOutputGradient != 0.0) {
                 // compute /sum_j v_{j,f}x_j
@@ -184,10 +186,10 @@ public Tensor[] gradients(Pair<Double, SGDVector> score, SGDVector features) {
                     }
                     factorGradMatrix = new DenseSparseMatrix(vectors);
                 } else {
-                    factorGradMatrix = new DenseMatrix(numFactors,features.size());
+                    factorGradMatrix = new DenseMatrix(numFactors, features.size());
                     for (int j = 0; j < numFactors; j++) {
                         for (int k = 0; k < features.size(); k++) {
-                            factorGradMatrix.set(j,k,features.get(k));
+                            factorGradMatrix.set(j, k, features.get(k));
                         }
                     }
                 }
@@ -198,7 +200,7 @@ public Tensor[] gradients(Pair<Double, SGDVector> score, SGDVector features) {
                     final int jFinal = j;
 
                     // Compute the gradient for this element of the factor vector
-                    curFactorGrad.foreachIndexedInPlace((Integer idx, Double a) -> a*curFactorSum - curFactors.get(jFinal,idx)*a*a);
+                    curFactorGrad.foreachIndexedInPlace((Integer idx, Double a) -> a * curFactorSum - curFactors.get(jFinal, idx) * a * a);
 
                     // Multiply by the output gradient
                     curFactorGrad.scaleInPlace(curOutputGradient);
@@ -207,7 +209,7 @@ public Tensor[] gradients(Pair<Double, SGDVector> score, SGDVector features) {
             } else {
                 // If the output gradient is 0.0 then all the factor gradients are zero.
                 // Technically with regularization we should shrink the weights for the specified features.
-                gradients[i] = new DenseSparseMatrix(numFactors,features.size());
+                gradients[i] = new DenseSparseMatrix(numFactors, features.size());
             }
         }
 
@@ -216,16 +218,17 @@ public Tensor[] gradients(Pair<Double, SGDVector> score, SGDVector features) {
 
     /**
      * This returns a {@link DenseMatrix} the same size as the Parameters.
+     *
      * @return A {@link Tensor} array containing a single {@link DenseMatrix}.
      */
     @Override
     public Tensor[] getEmptyCopy() {
         Tensor[] output = new Tensor[weights.length];
         output[0] = new DenseVector(biasVector.size());
-        output[1] = new DenseMatrix(weightMatrix.getDimension1Size(),weightMatrix.getDimension2Size());
+        output[1] = new DenseMatrix(weightMatrix.getDimension1Size(), weightMatrix.getDimension2Size());
         for (int i = 2; i < weights.length; i++) {
             DenseMatrix curMatrix = (DenseMatrix) weights[i];
-            output[i] = new DenseMatrix(curMatrix.getDimension1Size(),curMatrix.getDimension2Size());
+            output[i] = new DenseMatrix(curMatrix.getDimension1Size(), curMatrix.getDimension2Size());
         }
         return output;
     }
@@ -287,6 +290,6 @@ public FMParameters copy() {
         for (int i = 0; i < weights.length; i++) {
             weightCopy[i] = weights[i].copy();
         }
-        return new FMParameters(weightCopy,numFactors,l2);
+        return new FMParameters(weightCopy, numFactors);
     }
 }
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
index bf7de9a68..e3366a99f 100644
--- a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
@@ -52,8 +52,6 @@ public enum LossEnum {HINGE, SIGMOID}
     private long fmSeed = Trainer.DEFAULT_SEED;
     @Option(longName = "fm-factor-size", usage = "Factor size.")
     public int fmFactorSize = 6;
-    @Option(longName = "fm-l2-penalty", usage = "L2 regularization penalty.")
-    public double fmL2 = 0.001;
     @Option(longName = "fm-variance", usage = "Variance of the initialization gaussian.")
     public double fmVariance = 0.1;
 
@@ -75,6 +73,6 @@ public MultiLabelObjective getLoss() {
     public FMMultiLabelTrainer getTrainer() {
         logger.info(String.format("Set logging interval to %d", fmLoggingInterval));
         return new FMMultiLabelTrainer(getLoss(), sgoOptions.getOptimiser(), fmEpochs, fmLoggingInterval,
-                fmMinibatchSize, fmSeed, fmFactorSize, fmL2, fmVariance);
+                fmMinibatchSize, fmSeed, fmFactorSize, fmVariance);
     }
 }
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelTrainer.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelTrainer.java
index eecc10bbf..6e9fb903e 100644
--- a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelTrainer.java
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelTrainer.java
@@ -42,28 +42,28 @@
  * 2010 IEEE International Conference on Data Mining
  * </pre>
  */
-public class FMMultiLabelTrainer extends AbstractFMTrainer<MultiLabel,SGDVector> {
+public class FMMultiLabelTrainer extends AbstractFMTrainer<MultiLabel, SGDVector> {
     private static final Logger logger = Logger.getLogger(FMMultiLabelTrainer.class.getName());
 
-    @Config(description="The classification objective function to use.")
+    @Config(description = "The classification objective function to use.")
     private MultiLabelObjective objective = new BinaryCrossEntropy();
 
     /**
      * Constructs an SGD trainer for a multi-label factorization machine.
-     * @param objective The objective function to optimise.
-     * @param optimiser The gradient optimiser to use.
-     * @param epochs The number of epochs (complete passes through the training data).
-     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
-     * @param minibatchSize The size of any minibatches.
-     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     *
+     * @param objective         The objective function to optimise.
+     * @param optimiser         The gradient optimiser to use.
+     * @param epochs            The number of epochs (complete passes through the training data).
+     * @param loggingInterval   Log the loss after this many iterations. If -1 don't log anything.
+     * @param minibatchSize     The size of any minibatches.
+     * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
      * @param factorizedDimSize Size of the factorized feature representation.
-     * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
      */
     public FMMultiLabelTrainer(MultiLabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
                                int loggingInterval, int minibatchSize, long seed,
-                               int factorizedDimSize, double l2, double variance) {
-        super(optimiser,epochs,loggingInterval,minibatchSize,seed, factorizedDimSize, l2, variance);
+                               int factorizedDimSize, double variance) {
+        super(optimiser, epochs, loggingInterval, minibatchSize, seed, factorizedDimSize, variance);
         this.objective = objective;
     }
 
@@ -71,36 +71,36 @@ public FMMultiLabelTrainer(MultiLabelObjective objective, StochasticGradientOpti
      * Constructs an SGD trainer for a multi-label factorization machine.
      * <p>
      * Sets the minibatch size to 1.
-     * @param objective The objective function to optimise.
-     * @param optimiser The gradient optimiser to use.
-     * @param epochs The number of epochs (complete passes through the training data).
-     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
-     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     *
+     * @param objective         The objective function to optimise.
+     * @param optimiser         The gradient optimiser to use.
+     * @param epochs            The number of epochs (complete passes through the training data).
+     * @param loggingInterval   Log the loss after this many iterations. If -1 don't log anything.
+     * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
      * @param factorizedDimSize Size of the factorized feature representation.
-     * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
      */
     public FMMultiLabelTrainer(MultiLabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
                                int loggingInterval, long seed,
-                               int factorizedDimSize, double l2, double variance) {
-        this(objective,optimiser,epochs,loggingInterval,1,seed, factorizedDimSize, l2, variance);
+                               int factorizedDimSize, double variance) {
+        this(objective, optimiser, epochs, loggingInterval, 1, seed, factorizedDimSize, variance);
     }
 
     /**
      * Constructs an SGD trainer for a multi-label factorization machine.
      * <p>
      * Sets the minibatch size to 1 and the logging interval to 1000.
-     * @param objective The objective function to optimise.
-     * @param optimiser The gradient optimiser to use.
-     * @param epochs The number of epochs (complete passes through the training data).
-     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     *
+     * @param objective         The objective function to optimise.
+     * @param optimiser         The gradient optimiser to use.
+     * @param epochs            The number of epochs (complete passes through the training data).
+     * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
      * @param factorizedDimSize Size of the factorized feature representation.
-     * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
      */
     public FMMultiLabelTrainer(MultiLabelObjective objective, StochasticGradientOptimiser optimiser, int epochs,
-                               long seed, int factorizedDimSize, double l2, double variance) {
-        this(objective,optimiser,epochs,1000,1,seed, factorizedDimSize, l2, variance);
+                               long seed, int factorizedDimSize, double variance) {
+        this(objective, optimiser, epochs, 1000, 1, seed, factorizedDimSize, variance);
     }
 
     /**
@@ -122,7 +122,7 @@ protected SGDObjective<SGDVector> getObjective() {
 
     @Override
     protected FMMultiLabelModel createModel(String name, ModelProvenance provenance, ImmutableFeatureMap featureMap, ImmutableOutputInfo<MultiLabel> outputInfo, FMParameters parameters) {
-        return new FMMultiLabelModel(name,provenance,featureMap,outputInfo,parameters,objective.getNormalizer(),objective.isProbabilistic(),objective.threshold());
+        return new FMMultiLabelModel(name, provenance, featureMap, outputInfo, parameters, objective.getNormalizer(), objective.isProbabilistic(), objective.threshold());
     }
 
     @Override
@@ -134,7 +134,7 @@ protected String getModelClassName() {
     public String toString() {
         return "FMMultiLabelTrainer(objective=" + objective.toString() + ",optimiser=" + optimiser.toString() +
                 ",epochs=" + epochs + ",minibatchSize=" + minibatchSize + ",seed=" + seed +
-                ",factorizedDimSize=" + factorizedDimSize + ",l2=" + l2 + ",variance=" + variance +
+                ",factorizedDimSize=" + factorizedDimSize + ",variance=" + variance +
                 ")";
     }
 
diff --git a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
index dbed244d3..391beefe0 100644
--- a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
+++ b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
@@ -43,8 +43,8 @@
 
 public class TestFMMultiLabel {
 
-    private static final FMMultiLabelTrainer hinge = new FMMultiLabelTrainer(new Hinge(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED, 5,0.0,0.1);
-    private static final FMMultiLabelTrainer sigmoid = new FMMultiLabelTrainer(new BinaryCrossEntropy(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED,5,0.0,0.1);
+    private static final FMMultiLabelTrainer hinge = new FMMultiLabelTrainer(new Hinge(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED, 5,0.1);
+    private static final FMMultiLabelTrainer sigmoid = new FMMultiLabelTrainer(new BinaryCrossEntropy(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED,5,0.1);
 
     @BeforeAll
     public static void setup() {
@@ -77,7 +77,7 @@ private static void testTrainer(Dataset<MultiLabel> train, Dataset<MultiLabel> t
 
         MultiLabelEvaluation evaluation = (MultiLabelEvaluation) train.getOutputFactory().getEvaluator().evaluate(model,test);
 
-        Assertions.assertEquals(1.0,evaluation.microAveragedRecall());
+        Assertions.assertEquals(1.0, evaluation.microAveragedRecall());
 
         Helpers.testModelSerialization(model, MultiLabel.class);
     }
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
index 8d12e6755..f54faeb15 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionTrainer.java
@@ -43,32 +43,32 @@
  * 2010 IEEE International Conference on Data Mining
  * </pre>
  */
-public class FMRegressionTrainer extends AbstractFMTrainer<Regressor,DenseVector> {
+public class FMRegressionTrainer extends AbstractFMTrainer<Regressor, DenseVector> {
     private static final Logger logger = Logger.getLogger(FMRegressionTrainer.class.getName());
 
-    @Config(mandatory = true,description="The regression objective to use.")
+    @Config(mandatory = true, description = "The regression objective to use.")
     private RegressionObjective objective;
 
-    @Config(mandatory = true,description="Standardise the output variables before fitting the model.")
+    @Config(mandatory = true, description = "Standardise the output variables before fitting the model.")
     private boolean standardise;
 
     /**
      * Constructs an SGD trainer for a factorization machine.
-     * @param objective The objective function to optimise.
-     * @param optimiser The gradient optimiser to use.
-     * @param epochs The number of epochs (complete passes through the training data).
-     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
-     * @param minibatchSize The size of any minibatches.
-     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     *
+     * @param objective         The objective function to optimise.
+     * @param optimiser         The gradient optimiser to use.
+     * @param epochs            The number of epochs (complete passes through the training data).
+     * @param loggingInterval   Log the loss after this many iterations. If -1 don't log anything.
+     * @param minibatchSize     The size of any minibatches.
+     * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
      * @param factorizedDimSize Size of the factorized feature representation.
-     * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
      * @param standardise       Standardise the output regressors before fitting the model.
      */
     public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOptimiser optimiser, int epochs,
                                int loggingInterval, int minibatchSize, long seed,
-                               int factorizedDimSize, double l2, double variance, boolean standardise) {
-        super(optimiser, epochs, loggingInterval, minibatchSize, seed, factorizedDimSize, l2, variance);
+                               int factorizedDimSize, double variance, boolean standardise) {
+        super(optimiser, epochs, loggingInterval, minibatchSize, seed, factorizedDimSize, variance);
         this.objective = objective;
         this.standardise = standardise;
     }
@@ -77,38 +77,38 @@ public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOpti
      * Constructs an SGD trainer for a factorization machine.
      * <p>
      * Sets the minibatch size to 1.
-     * @param objective The objective function to optimise.
-     * @param optimiser The gradient optimiser to use.
-     * @param epochs The number of epochs (complete passes through the training data).
-     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
-     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     *
+     * @param objective         The objective function to optimise.
+     * @param optimiser         The gradient optimiser to use.
+     * @param epochs            The number of epochs (complete passes through the training data).
+     * @param loggingInterval   Log the loss after this many iterations. If -1 don't log anything.
+     * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
      * @param factorizedDimSize Size of the factorized feature representation.
-     * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
      * @param standardise       Standardise the output regressors before fitting the model.
      */
     public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOptimiser optimiser, int epochs,
                                int loggingInterval, long seed,
-                               int factorizedDimSize, double l2, double variance, boolean standardise) {
-        this(objective,optimiser,epochs,loggingInterval,1,seed,factorizedDimSize,l2,variance,standardise);
+                               int factorizedDimSize, double variance, boolean standardise) {
+        this(objective, optimiser, epochs, loggingInterval, 1, seed, factorizedDimSize, variance, standardise);
     }
 
     /**
      * Constructs an SGD trainer for a factorization machine.
      * <p>
      * Sets the minibatch size to 1 and the logging interval to 1000.
-     * @param objective The objective function to optimise.
-     * @param optimiser The gradient optimiser to use.
-     * @param epochs The number of epochs (complete passes through the training data).
-     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
+     *
+     * @param objective         The objective function to optimise.
+     * @param optimiser         The gradient optimiser to use.
+     * @param epochs            The number of epochs (complete passes through the training data).
+     * @param seed              A seed for the random number generator, used to shuffle the examples before each epoch.
      * @param factorizedDimSize Size of the factorized feature representation.
-     * @param l2                The l2 regularisation penalty.
      * @param variance          The variance of the initializer.
      * @param standardise       Standardise the output regressors before fitting the model.
      */
     public FMRegressionTrainer(RegressionObjective objective, StochasticGradientOptimiser optimiser, int epochs,
-                               long seed, int factorizedDimSize, double l2, double variance, boolean standardise) {
-        this(objective,optimiser,epochs,1000,1,seed,factorizedDimSize,l2,variance,standardise);
+                               long seed, int factorizedDimSize, double variance, boolean standardise) {
+        this(objective, optimiser, epochs, 1000, 1, seed, factorizedDimSize, variance, standardise);
     }
 
     /**
@@ -145,7 +145,7 @@ protected FMRegressionModel createModel(String name, ModelProvenance provenance,
             int id = outputInfo.getID(r);
             dimensionNames[id] = r.getNames()[0];
         }
-        return new FMRegressionModel(name,dimensionNames,provenance,featureMap,outputInfo,parameters,standardise);
+        return new FMRegressionModel(name, dimensionNames, provenance, featureMap, outputInfo, parameters, standardise);
     }
 
     @Override
@@ -157,7 +157,7 @@ protected String getModelClassName() {
     public String toString() {
         return "FMRegressionTrainer(objective=" + objective.toString() + ",optimiser=" + optimiser.toString() +
                 ",epochs=" + epochs + ",minibatchSize=" + minibatchSize + ",seed=" + seed +
-                ",factorizedDimSize=" + factorizedDimSize + ",l2=" + l2 + ",variance=" + variance +
+                ",factorizedDimSize=" + factorizedDimSize + ",variance=" + variance +
                 ",standardise=" + standardise + ")";
     }
 }
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
index 3b558d226..ae3e16b8a 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
@@ -147,7 +147,7 @@ public static void main(String[] args) throws IOException {
 
         logger.info("Feature domain - " + train.getFeatureIDMap());
 
-        Trainer<Regressor> trainer = new FMRegressionTrainer(obj,grad,o.epochs,o.loggingInterval,o.minibatchSize,o.general.seed,o.factorSize,o.l2,o.variance,o.standardise);
+        Trainer<Regressor> trainer = new FMRegressionTrainer(obj,grad,o.epochs,o.loggingInterval,o.minibatchSize,o.general.seed,o.factorSize,o.variance,o.standardise);
         logger.info("Training using " + trainer.toString());
         final long trainStart = System.currentTimeMillis();
         Model<Regressor> model = trainer.train(train);
diff --git a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
index 9a5810eec..6fa28e824 100644
--- a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
+++ b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
@@ -46,7 +46,7 @@
 
 public class TestFMRegression {
 
-    private static final FMRegressionTrainer t = new FMRegressionTrainer(new SquaredLoss(), new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED,1,0.0,0.1,true);
+    private static final FMRegressionTrainer t = new FMRegressionTrainer(new SquaredLoss(), new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED,1,0.1,true);
 
     private static final RegressionEvaluator e = new RegressionEvaluator();
 

From 87f4cc8ff19f12ada429a2e836041e3e92e6c632 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sat, 25 Sep 2021 23:31:21 -0400
Subject: [PATCH 16/31] Adding some more overloads to ONNXOperators.build for
 single input and output operators, updating other code to use the new
 overloads.

---
 .../sgd/linear/LinearSGDModel.java            |  2 +-
 .../java/org/tribuo/onnx/ONNXAttribute.java   |  5 ++
 .../java/org/tribuo/onnx/ONNXOperators.java   | 74 +++++++++++++++++--
 .../org/tribuo/math/util/ExpNormalizer.java   |  2 +-
 .../org/tribuo/math/util/NoopNormalizer.java  |  2 +-
 .../java/org/tribuo/math/util/Normalizer.java |  8 +-
 .../tribuo/math/util/SigmoidNormalizer.java   |  2 +-
 .../multilabel/sgd/linear/LinearSGDModel.java |  2 +-
 .../regression/sgd/linear/LinearSGDModel.java |  2 +-
 .../regression/slm/SparseLinearModel.java     | 10 +--
 10 files changed, 88 insertions(+), 21 deletions(-)

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/linear/LinearSGDModel.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/linear/LinearSGDModel.java
index df921e90b..2d929928b 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/linear/LinearSGDModel.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/linear/LinearSGDModel.java
@@ -141,7 +141,7 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
 
         // Make gemm
         String[] gemmInputs = new String[]{inputValueProto.getName(),weightInitializerProto.getName(),biasInitializerProto.getName()};
-        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context,gemmInputs,new String[]{context.generateUniqueName("gemm_output")},Collections.emptyMap());
+        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context,gemmInputs,context.generateUniqueName("gemm_output"));
         graphBuilder.addNode(gemm);
 
         // Make output normalizer
diff --git a/Core/src/main/java/org/tribuo/onnx/ONNXAttribute.java b/Core/src/main/java/org/tribuo/onnx/ONNXAttribute.java
index 9f39c9fdb..7c5c11760 100644
--- a/Core/src/main/java/org/tribuo/onnx/ONNXAttribute.java
+++ b/Core/src/main/java/org/tribuo/onnx/ONNXAttribute.java
@@ -26,6 +26,11 @@
  */
 public final class ONNXAttribute {
 
+    /**
+     * Value used to denote a varaidic input in an ONNX operator
+     */
+    public static final int VARIADIC_INPUT = -1;
+
     private final String name;
     private final OnnxMl.AttributeProto.AttributeType type;
     private final boolean mandatory;
diff --git a/Core/src/main/java/org/tribuo/onnx/ONNXOperators.java b/Core/src/main/java/org/tribuo/onnx/ONNXOperators.java
index e7878a497..d93d1b58b 100644
--- a/Core/src/main/java/org/tribuo/onnx/ONNXOperators.java
+++ b/Core/src/main/java/org/tribuo/onnx/ONNXOperators.java
@@ -26,16 +26,22 @@
 import java.util.Map;
 import java.util.Set;
 
+import static org.tribuo.onnx.ONNXAttribute.VARIADIC_INPUT;
+
 /**
  * The supported ONNX operators.
  */
 public enum ONNXOperators {
-
-
     /**
-     * Identity,
+     * Identity.
      */
     IDENTITY("Identity",1,1),
+    /**
+     * Concatenates tensors.
+     */
+    CONCAT("Constant",VARIADIC_INPUT,1, Collections.singletonList(
+            new ONNXAttribute("axis", OnnxMl.AttributeProto.AttributeType.INT, true)
+    )),
     /**
      * Sigmoid element-wise.
      */
@@ -47,7 +53,7 @@ public enum ONNXOperators {
      * </ul>
      */
     SOFTMAX("Softmax",1,1, Collections.singletonList(
-            new ONNXAttribute("axis", OnnxMl.AttributeProto.AttributeType.INT,false)
+            new ONNXAttribute("axis", OnnxMl.AttributeProto.AttributeType.INT, false)
     )),
     /**
      * Element-wise addition with broadcasting.
@@ -65,6 +71,10 @@ public enum ONNXOperators {
      * Element-wise division with broadcasting.
      */
     DIV("Div",2,1),
+    /**
+     * Element-wise exponentiation with broadcasting.
+     */
+    POW("Pow",2,1),
     /**
      * Compute the minimum along the specified axes of the tensor.
      * <ul>
@@ -83,7 +93,7 @@ public enum ONNXOperators {
      *     <li>{@code keepdims} defaults to 1 which means keep.</li>
      * </ul>
      */
-    REDUCE_SUM("ReduceSum",2,1,Arrays.asList(
+    REDUCE_SUM("ReduceSum",1,1,Arrays.asList(
             new ONNXAttribute("axes", OnnxMl.AttributeProto.AttributeType.INTS, false), //Opset 11
             new ONNXAttribute("keepdims", OnnxMl.AttributeProto.AttributeType.INT, false)
     )),
@@ -169,6 +179,58 @@ private ONNXOperators(String value, int numInputs, int numOutputs, List<ONNXAttr
         this.mandatoryAttributeNames = attributeSet.isEmpty() ? Collections.emptySet() : Collections.unmodifiableSet(attributeSet);
     }
 
+    /**
+     * Builds this node based on the supplied inputs and output.
+     * Throws {@link IllegalArgumentException} if this operator takes more than a single input or output.
+     * @param context The onnx context used to ensure this node has a unique name.
+     * @param input The name of the input.
+     * @param output The name of the output.
+     * @return The NodeProto.
+     */
+    public OnnxMl.NodeProto build(ONNXContext context, String input, String output) {
+        return build(context,new String[]{input},new String[]{output},Collections.emptyMap());
+    }
+
+    /**
+     * Builds this node based on the supplied inputs and output.
+     * Throws {@link IllegalArgumentException} if this operator takes more than a single input or output.
+     * May throw {@link UnsupportedOperationException} if the attribute type is not supported.
+     * @param context The onnx context used to ensure this node has a unique name.
+     * @param input The names of the input.
+     * @param output The name of the output.
+     * @param attributeValues The attribute names and values.
+     * @return The NodeProto.
+     */
+    public OnnxMl.NodeProto build(ONNXContext context, String input, String output, Map<String,Object> attributeValues) {
+        return build(context,new String[]{input},new String[]{output},attributeValues);
+    }
+
+    /**
+     * Builds this node based on the supplied inputs and output.
+     * Throws {@link IllegalArgumentException} if the number of inputs or outputs is wrong.
+     * @param context The onnx context used to ensure this node has a unique name.
+     * @param inputs The names of the inputs.
+     * @param output The name of the output.
+     * @return The NodeProto.
+     */
+    public OnnxMl.NodeProto build(ONNXContext context, String[] inputs, String output) {
+        return build(context,inputs,new String[]{output},Collections.emptyMap());
+    }
+
+    /**
+     * Builds this node based on the supplied inputs and output.
+     * Throws {@link IllegalArgumentException} if the number of inputs, outputs or attributes is wrong.
+     * May throw {@link UnsupportedOperationException} if the attribute type is not supported.
+     * @param context The onnx context used to ensure this node has a unique name.
+     * @param inputs The names of the inputs.
+     * @param output The name of the output.
+     * @param attributeValues The attribute names and values.
+     * @return The NodeProto.
+     */
+    public OnnxMl.NodeProto build(ONNXContext context, String[] inputs, String output, Map<String,Object> attributeValues) {
+        return build(context,inputs,new String[]{output},attributeValues);
+    }
+
     /**
      * Builds this node based on the supplied inputs and outputs.
      * Throws {@link IllegalArgumentException} if the number of inputs or outputs is wrong.
@@ -192,7 +254,7 @@ public OnnxMl.NodeProto build(ONNXContext context, String[] inputs, String[] out
      * @return The NodeProto.
      */
     public OnnxMl.NodeProto build(ONNXContext context, String[] inputs, String[] outputs, Map<String,Object> attributeValues) {
-        if (inputs.length != numInputs) {
+        if ((numInputs != VARIADIC_INPUT) && (inputs.length != numInputs)) {
             throw new IllegalArgumentException("Expected " + numInputs + " inputs, but received " + inputs.length);
         }
         if (outputs.length != numOutputs) {
diff --git a/Math/src/main/java/org/tribuo/math/util/ExpNormalizer.java b/Math/src/main/java/org/tribuo/math/util/ExpNormalizer.java
index 10ca17a32..0e47faf6b 100644
--- a/Math/src/main/java/org/tribuo/math/util/ExpNormalizer.java
+++ b/Math/src/main/java/org/tribuo/math/util/ExpNormalizer.java
@@ -67,7 +67,7 @@ public void normalizeInPlace(double[] input) {
      */
     @Override
     public List<OnnxMl.NodeProto> exportNormalizer(ONNXContext context, String input, String output) {
-       return Collections.singletonList(ONNXOperators.SOFTMAX.build(context,new String[]{input},new String[]{output}, Collections.singletonMap("axis",1)));
+       return Collections.singletonList(ONNXOperators.SOFTMAX.build(context, input, output, Collections.singletonMap("axis",1)));
     }
 
 }
diff --git a/Math/src/main/java/org/tribuo/math/util/NoopNormalizer.java b/Math/src/main/java/org/tribuo/math/util/NoopNormalizer.java
index c0b1d5f84..9c5c885c7 100644
--- a/Math/src/main/java/org/tribuo/math/util/NoopNormalizer.java
+++ b/Math/src/main/java/org/tribuo/math/util/NoopNormalizer.java
@@ -48,6 +48,6 @@ public void normalizeInPlace(double[] input) {}
      */
     @Override
     public List<OnnxMl.NodeProto> exportNormalizer(ONNXContext context, String input, String output) {
-        return Collections.singletonList(ONNXOperators.IDENTITY.build(context,new String[]{input},new String[]{output}));
+        return Collections.singletonList(ONNXOperators.IDENTITY.build(context,input,output));
     }
 }
diff --git a/Math/src/main/java/org/tribuo/math/util/Normalizer.java b/Math/src/main/java/org/tribuo/math/util/Normalizer.java
index 4d0875cfb..597d01470 100644
--- a/Math/src/main/java/org/tribuo/math/util/Normalizer.java
+++ b/Math/src/main/java/org/tribuo/math/util/Normalizer.java
@@ -70,18 +70,18 @@ public List<OnnxMl.NodeProto> exportNormalizer(ONNXContext context, String input
         List<OnnxMl.NodeProto> protos = new ArrayList<>();
 
         String minOutput = context.generateUniqueName("min_output");
-        OnnxMl.NodeProto min = ONNXOperators.REDUCE_MIN.build(context,new String[]{input},new String[]{minOutput},Collections.singletonMap("axes",1));
+        OnnxMl.NodeProto min = ONNXOperators.REDUCE_MIN.build(context,input,minOutput,Collections.singletonMap("axes",1));
         protos.add(min);
 
         String subOutput = context.generateUniqueName("sub_output");
-        OnnxMl.NodeProto sub = ONNXOperators.SUB.build(context,new String[]{input,minOutput},new String[]{subOutput});
+        OnnxMl.NodeProto sub = ONNXOperators.SUB.build(context,new String[]{input,minOutput},subOutput);
         protos.add(sub);
 
         String sumOutput = context.generateUniqueName("sum_output");
-        OnnxMl.NodeProto sum = ONNXOperators.REDUCE_SUM.build(context,new String[]{subOutput},new String[]{sumOutput},Collections.singletonMap("axes",1));
+        OnnxMl.NodeProto sum = ONNXOperators.REDUCE_SUM.build(context,subOutput,sumOutput,Collections.singletonMap("axes",1));
         protos.add(sum);
 
-        OnnxMl.NodeProto div = ONNXOperators.DIV.build(context,new String[]{subOutput,sumOutput},new String[]{output});
+        OnnxMl.NodeProto div = ONNXOperators.DIV.build(context,new String[]{subOutput,sumOutput},output);
         protos.add(div);
 
         return protos;
diff --git a/Math/src/main/java/org/tribuo/math/util/SigmoidNormalizer.java b/Math/src/main/java/org/tribuo/math/util/SigmoidNormalizer.java
index ca5c191f1..1b3ef0c15 100644
--- a/Math/src/main/java/org/tribuo/math/util/SigmoidNormalizer.java
+++ b/Math/src/main/java/org/tribuo/math/util/SigmoidNormalizer.java
@@ -64,7 +64,7 @@ public void normalizeInPlace(double[] input) {
      */
     @Override
     public List<OnnxMl.NodeProto> exportNormalizer(ONNXContext context, String input, String output) {
-        return Collections.singletonList(ONNXOperators.SIGMOID.build(context,new String[]{input},new String[]{output}));
+        return Collections.singletonList(ONNXOperators.SIGMOID.build(context,input,output));
     }
 
 }
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/linear/LinearSGDModel.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/linear/LinearSGDModel.java
index 535eaac3c..88193430b 100644
--- a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/linear/LinearSGDModel.java
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/linear/LinearSGDModel.java
@@ -137,7 +137,7 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
 
         // Make gemm
         String[] gemmInputs = new String[]{inputValueProto.getName(),weightInitializerProto.getName(),biasInitializerProto.getName()};
-        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context,gemmInputs,new String[]{context.generateUniqueName("gemm_output")}, Collections.emptyMap());
+        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context,gemmInputs,context.generateUniqueName("gemm_output"));
         graphBuilder.addNode(gemm);
 
         // Make output normalizer
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDModel.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDModel.java
index 1794fd55d..362f370d8 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDModel.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDModel.java
@@ -123,7 +123,7 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
 
         // Make gemm
         String[] gemmInputs = new String[]{inputValueProto.getName(),weightInitializerProto.getName(),biasInitializerProto.getName()};
-        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context,gemmInputs,new String[]{"output"}, Collections.emptyMap());
+        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context,gemmInputs,"output");
         graphBuilder.addNode(gemm);
 
         return graphBuilder.build();
diff --git a/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java b/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
index ad8ed35c0..0c8a26295 100644
--- a/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
+++ b/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
@@ -328,10 +328,10 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
 
         // Scale features
         String featureMeanOutput = context.generateUniqueName("feature_mean_scale_output");
-        OnnxMl.NodeProto subFeatureMean = ONNXOperators.SUB.build(context,new String[]{inputValueProto.getName(),featureMeanProto.getName()},new String[]{featureMeanOutput});
+        OnnxMl.NodeProto subFeatureMean = ONNXOperators.SUB.build(context,new String[]{inputValueProto.getName(),featureMeanProto.getName()},featureMeanOutput);
         graphBuilder.addNode(subFeatureMean);
         String featureVarianceOutput = context.generateUniqueName("feature_var_scale_output");
-        OnnxMl.NodeProto divFeatureVariance = ONNXOperators.DIV.build(context,new String[]{subFeatureMean.getOutput(0),featureVarianceProto.getName()},new String[]{featureVarianceOutput});
+        OnnxMl.NodeProto divFeatureVariance = ONNXOperators.DIV.build(context,new String[]{subFeatureMean.getOutput(0),featureVarianceProto.getName()},featureVarianceOutput);
         graphBuilder.addNode(divFeatureVariance);
 
         // Make gemm
@@ -339,16 +339,16 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
                 weightInitializerProto.getName(),
                 biasInitializerProto.getName()};
         String gemmOutput = context.generateUniqueName("gemm_output");
-        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context, gemmInputs, new String[]{gemmOutput});
+        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context, gemmInputs, gemmOutput);
         graphBuilder.addNode(gemm);
 
         // Scale outputs
         String varianceOutput = context.generateUniqueName("y_var_scale_output");
-        OnnxMl.NodeProto varianceScale = ONNXOperators.MUL.build(context, new String[]{gemmOutput,outputVarianceProto.getName()}, new String[]{varianceOutput});
+        OnnxMl.NodeProto varianceScale = ONNXOperators.MUL.build(context, new String[]{gemmOutput,outputVarianceProto.getName()}, varianceOutput);
         graphBuilder.addNode(varianceScale);
 
         String meanOutput = "output";
-        OnnxMl.NodeProto meanScale = ONNXOperators.ADD.build(context, new String[]{varianceOutput,outputMeanProto.getName()}, new String[]{meanOutput});
+        OnnxMl.NodeProto meanScale = ONNXOperators.ADD.build(context, new String[]{varianceOutput,outputMeanProto.getName()}, meanOutput);
         graphBuilder.addNode(meanScale);
 
         return graphBuilder.build();

From d3ef792425f8b1c8e1dbc2a041f6cbd6d6730c7a Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sat, 25 Sep 2021 23:32:01 -0400
Subject: [PATCH 17/31] Adding ONNX export to FMClassificationModel.

---
 .../sgd/fm/FMClassificationModel.java         | 136 +++++++++++++++++-
 .../sgd/fm/TestFMClassification.java          |  77 ++++++++++
 .../tribuo/common/sgd/AbstractFMModel.java    |  84 +++++++++++
 .../org/tribuo/common/sgd/FMParameters.java   |   5 +-
 4 files changed, 300 insertions(+), 2 deletions(-)

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
index 7cb96d2c8..3db74031b 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
@@ -16,6 +16,7 @@
 
 package org.tribuo.classification.sgd.fm;
 
+import ai.onnx.proto.OnnxMl;
 import org.tribuo.Example;
 import org.tribuo.ImmutableFeatureMap;
 import org.tribuo.ImmutableOutputInfo;
@@ -23,11 +24,20 @@
 import org.tribuo.classification.Label;
 import org.tribuo.common.sgd.AbstractFMModel;
 import org.tribuo.common.sgd.FMParameters;
+import org.tribuo.math.la.DenseMatrix;
 import org.tribuo.math.la.DenseVector;
+import org.tribuo.math.la.Tensor;
 import org.tribuo.math.util.VectorNormalizer;
+import org.tribuo.onnx.ONNXContext;
+import org.tribuo.onnx.ONNXExportable;
+import org.tribuo.onnx.ONNXOperators;
+import org.tribuo.onnx.ONNXShape;
+import org.tribuo.onnx.ONNXUtils;
 import org.tribuo.provenance.ModelProvenance;
 
+import java.util.Collections;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -40,7 +50,7 @@
  * 2010 IEEE International Conference on Data Mining
  * </pre>
  */
-public class FMClassificationModel extends AbstractFMModel<Label> {
+public class FMClassificationModel extends AbstractFMModel<Label> implements ONNXExportable {
     private static final long serialVersionUID = 1L;
 
     private final VectorNormalizer normalizer;
@@ -93,4 +103,128 @@ protected FMClassificationModel copy(String newName, ModelProvenance newProvenan
     protected String getDimensionName(int index) {
         return outputIDInfo.getOutput(index).getLabel();
     }
+
+    @Override
+    public OnnxMl.ModelProto exportONNXModel(String domain, long modelVersion) {
+        ONNXContext context = new ONNXContext();
+
+        // Build graph
+        OnnxMl.GraphProto graph = exportONNXGraph(context);
+
+        return innerExportONNXModel(graph,domain,modelVersion);
+    }
+
+    @Override
+    public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
+        OnnxMl.GraphProto.Builder graphBuilder = OnnxMl.GraphProto.newBuilder();
+
+        // Make inputs and outputs
+        OnnxMl.TypeProto inputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,featureIDMap.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
+        OnnxMl.ValueInfoProto inputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(inputType).setName("input").build();
+        graphBuilder.addInput(inputValueProto);
+        OnnxMl.TypeProto outputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,outputIDInfo.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
+        OnnxMl.ValueInfoProto outputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(outputType).setName("output").build();
+        graphBuilder.addOutput(outputValueProto);
+
+        Tensor[] modelParams = modelParameters.get();
+
+        // Add constants
+        OnnxMl.TensorProto twoConst = OnnxMl.TensorProto.newBuilder()
+                .setName(context.generateUniqueName("two_const"))
+                .setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber())
+                .addFloatData(2.0f)
+                .build();
+        graphBuilder.addInitializer(twoConst);
+
+        // Add weights
+        OnnxMl.TensorProto weightInitializerProto = matrixBuilder(context, "fm_linear_weights", (DenseMatrix) modelParams[1], true);
+        graphBuilder.addInitializer(weightInitializerProto);
+
+        // Add biases
+        OnnxMl.TensorProto biasInitializerProto = biasBuilder(context);
+        graphBuilder.addInitializer(biasInitializerProto);
+
+        // Add embedding vectors
+        OnnxMl.TensorProto[] embeddingProtos = new OnnxMl.TensorProto[outputIDInfo.size()];
+        for (int i = 0; i < outputIDInfo.size(); i++) {
+            embeddingProtos[i] = matrixBuilder(context, "fm_embedding_"+i, (DenseMatrix) modelParams[i+2], true);
+            graphBuilder.addInitializer(embeddingProtos[i]);
+        }
+
+        // Make gemm
+        String[] gemmInputs = new String[]{inputValueProto.getName(),weightInitializerProto.getName(),biasInitializerProto.getName()};
+        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context,gemmInputs,context.generateUniqueName("gemm_output"));
+        graphBuilder.addNode(gemm);
+
+        // Make feature pow
+        OnnxMl.NodeProto featurePow = ONNXOperators.POW.build(context,new String[]{inputValueProto.getName(),twoConst.getName()},
+                context.generateUniqueName("feature_pow"));
+        graphBuilder.addNode(featurePow);
+
+        // Make interaction terms
+        String[] embeddingOutputs = new String[outputIDInfo.size()];
+        for (int i = 0; i < outputIDInfo.size(); i++) {
+            // Feature matrix * embedding matrix = batch_size, embedding dim
+            OnnxMl.NodeProto gemmFeatureEmb = ONNXOperators.GEMM.build(context,
+                    new String[]{inputValueProto.getName(),embeddingProtos[i].getName()},
+                    context.generateUniqueName("gemm_input_emb"),
+                    Collections.singletonMap("transB",1));
+            graphBuilder.addNode(gemmFeatureEmb);
+            // Square the output
+            OnnxMl.NodeProto powFeatureEmb = ONNXOperators.POW.build(context,
+                    new String[]{gemmFeatureEmb.getOutput(0),twoConst.getName()},
+                    context.generateUniqueName("pow_input_emb"));
+            graphBuilder.addNode(powFeatureEmb);
+            // Square the embeddings
+            OnnxMl.NodeProto powEmb = ONNXOperators.POW.build(context,
+                    new String[]{embeddingProtos[i].getName(),twoConst.getName()},
+                    context.generateUniqueName("pow_emb"));
+            graphBuilder.addNode(powEmb);
+            // squared features * squared embeddings
+            OnnxMl.NodeProto gemmSquaredFeatureSquaredEmb = ONNXOperators.GEMM.build(context,
+                    new String[]{featurePow.getOutput(0),powEmb.getOutput(0)},
+                    context.generateUniqueName("gemm_squared_input_squared_emb"),
+                    Collections.singletonMap("transB",1));
+            graphBuilder.addNode(gemmSquaredFeatureSquaredEmb);
+            // squared product subtract product of squares
+            OnnxMl.NodeProto subtract = ONNXOperators.SUB.build(context,
+                    new String[]{powFeatureEmb.getOutput(0),gemmSquaredFeatureSquaredEmb.getOutput(0)},
+                    context.generateUniqueName("squared_prod_subtract_prod_of_squares"));
+            graphBuilder.addNode(subtract);
+            // sum over embedding dimensions
+            OnnxMl.NodeProto sumOverEmbeddings = ONNXOperators.REDUCE_SUM.build(context,
+                    subtract.getOutput(0),
+                    context.generateUniqueName("sum_over_embeddings"),
+                    Collections.singletonMap("axes",1));
+            graphBuilder.addNode(sumOverEmbeddings);
+            // Divide by 2
+            OnnxMl.NodeProto scaledInteraction = ONNXOperators.DIV.build(context,
+                    new String[]{sumOverEmbeddings.getOutput(0),twoConst.getName()},
+                    context.generateUniqueName("scaled_interaction"));
+            graphBuilder.addNode(scaledInteraction);
+            // Store the output name
+            embeddingOutputs[i] = scaledInteraction.getOutput(0);
+        }
+
+        // Make concat
+        OnnxMl.NodeProto concat = ONNXOperators.CONCAT.build(context, embeddingOutputs, context.generateUniqueName("fm_concat"),
+                    Collections.singletonMap("axes",1)
+                );
+        graphBuilder.addNode(concat);
+
+        // Add to gemm
+        OnnxMl.NodeProto addGemmConcat = ONNXOperators.ADD.build(context, new String[]{gemm.getOutput(0),concat.getOutput(0)}, context.generateUniqueName("fm_output"));
+        graphBuilder.addNode(addGemmConcat);
+
+        // Make output normalizer
+        List<OnnxMl.NodeProto> normalizerProtos = normalizer.exportNormalizer(context,addGemmConcat.getOutput(0),"output");
+        if (normalizerProtos.isEmpty()) {
+            throw new IllegalArgumentException("Normalizer " + normalizer.getClass() + " cannot be exported in ONNX models.");
+        } else {
+            graphBuilder.addAllNode(normalizerProtos);
+        }
+
+        return graphBuilder.build();
+    }
+
 }
diff --git a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
index 3db6682b1..b50f6ef14 100644
--- a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
+++ b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
@@ -16,31 +16,50 @@
 
 package org.tribuo.classification.sgd.fm;
 
+import ai.onnxruntime.OrtEnvironment;
+import ai.onnxruntime.OrtException;
+import ai.onnxruntime.OrtSession;
 import com.oracle.labs.mlrg.olcut.util.Pair;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.tribuo.Dataset;
 import org.tribuo.Model;
+import org.tribuo.Prediction;
 import org.tribuo.Trainer;
+import org.tribuo.VariableIDInfo;
+import org.tribuo.VariableInfo;
 import org.tribuo.classification.Label;
+import org.tribuo.classification.LabelFactory;
 import org.tribuo.classification.evaluation.LabelEvaluation;
 import org.tribuo.classification.evaluation.LabelEvaluator;
 import org.tribuo.classification.example.LabelledDataGenerator;
+import org.tribuo.classification.sgd.linear.LinearSGDModel;
+import org.tribuo.classification.sgd.linear.TestSGDLinear;
 import org.tribuo.classification.sgd.objectives.Hinge;
 import org.tribuo.common.sgd.AbstractFMTrainer;
 import org.tribuo.common.sgd.AbstractSGDTrainer;
+import org.tribuo.interop.onnx.DenseTransformer;
+import org.tribuo.interop.onnx.LabelTransformer;
+import org.tribuo.interop.onnx.ONNXExternalModel;
 import org.tribuo.math.optimisers.AdaGrad;
 import org.tribuo.test.Helpers;
 
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.fail;
 
 public class TestFMClassification {
+    private static final Logger logger = Logger.getLogger(TestFMClassification.class.getName());
 
     private static final FMClassificationTrainer t = new FMClassificationTrainer(new Hinge(),
             new AdaGrad(0.1,0.1),5,1000,Trainer.DEFAULT_SEED,6,0.1);
@@ -103,4 +122,62 @@ public void testEmptyExample() {
             m.predict(LabelledDataGenerator.emptyExample());
         });
     }
+    @Test
+
+    public void testOnnxSerialization() throws IOException, OrtException {
+        Pair<Dataset<Label>,Dataset<Label>> p = LabelledDataGenerator.denseTrainTest();
+        FMClassificationModel model = (FMClassificationModel) t.train(p.getA());
+
+        // Write out model
+        Path onnxFile = Files.createTempFile("tribuo-fm-test",".onnx");
+        model.saveONNXModel("org.tribuo.classification.sgd.fm.test",1,onnxFile);
+
+        // Prep mappings
+        Map<String, Integer> featureMapping = new HashMap<>();
+        for (VariableInfo f : model.getFeatureIDMap()){
+            VariableIDInfo id = (VariableIDInfo) f;
+            featureMapping.put(id.getName(),id.getID());
+        }
+        Map<Label, Integer> outputMapping = new HashMap<>();
+        for (Pair<Integer,Label> l : model.getOutputIDInfo()) {
+            outputMapping.put(l.getB(), l.getA());
+        }
+
+        String arch = System.getProperty("os.arch");
+        if (arch.equalsIgnoreCase("amd64") || arch.equalsIgnoreCase("x86_64")) {
+            // Initialise the OrtEnvironment to load the native library
+            // (as OrtSession.SessionOptions doesn't trigger the static initializer).
+            OrtEnvironment env = OrtEnvironment.getEnvironment();
+            env.close();
+            // Load in via ORT
+            ONNXExternalModel<Label> onnxModel = ONNXExternalModel.createOnnxModel(new LabelFactory(), featureMapping, outputMapping, new DenseTransformer(), new LabelTransformer(), new OrtSession.SessionOptions(), onnxFile, "input");
+
+            // Generate predictions
+            List<Prediction<Label>> nativePredictions = model.predict(p.getB());
+            List<Prediction<Label>> onnxPredictions = onnxModel.predict(p.getB());
+
+            // Assert the predictions are identical
+            for (int i = 0; i < nativePredictions.size(); i++) {
+                Prediction<Label> tribuo = nativePredictions.get(i);
+                Prediction<Label> external = onnxPredictions.get(i);
+                assertEquals(tribuo.getOutput().getLabel(), external.getOutput().getLabel());
+                assertEquals(tribuo.getOutput().getScore(), external.getOutput().getScore(), 1e-6);
+                for (Map.Entry<String, Label> l : tribuo.getOutputScores().entrySet()) {
+                    Label other = external.getOutputScores().get(l.getKey());
+                    if (other == null) {
+                        fail("Failed to find label " + l.getKey() + " in ORT prediction.");
+                    } else {
+                        assertEquals(l.getValue().getScore(), other.getScore(), 1e-6);
+                    }
+                }
+            }
+
+            onnxModel.close();
+        } else {
+            logger.warning("ORT based tests only supported on x86_64, found " + arch);
+        }
+
+        onnxFile.toFile().delete();
+    }
+
 }
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
index bec2675f6..d09f45e4e 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
@@ -16,17 +16,25 @@
 
 package org.tribuo.common.sgd;
 
+import ai.onnx.proto.OnnxMl;
+import com.google.protobuf.ByteString;
 import com.oracle.labs.mlrg.olcut.util.Pair;
 import org.tribuo.Example;
 import org.tribuo.Excuse;
 import org.tribuo.ImmutableFeatureMap;
 import org.tribuo.ImmutableOutputInfo;
 import org.tribuo.Output;
+import org.tribuo.Tribuo;
 import org.tribuo.math.la.DenseMatrix;
 import org.tribuo.math.la.DenseVector;
 import org.tribuo.math.la.Tensor;
+import org.tribuo.onnx.ONNXContext;
+import org.tribuo.onnx.ONNXOperators;
 import org.tribuo.provenance.ModelProvenance;
 
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.FloatBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
@@ -160,4 +168,80 @@ public Optional<Excuse<T>> getExcuse(Example<T> example) {
      */
     protected abstract String getDimensionName(int index);
 
+    /**
+     * Builds the ModelProto according to the standards for this model.
+     * @param graph The model graph.
+     * @param domain The model domain string.
+     * @param modelVersion The model version number.
+     * @return The ModelProto.
+     */
+    protected OnnxMl.ModelProto innerExportONNXModel(OnnxMl.GraphProto graph, String domain, long modelVersion) {
+        // Build model
+        OnnxMl.ModelProto.Builder builder = OnnxMl.ModelProto.newBuilder();
+        builder.setGraph(graph);
+        builder.setDomain(domain);
+        builder.setProducerName("Tribuo");
+        builder.setProducerVersion(Tribuo.VERSION);
+        builder.setModelVersion(modelVersion);
+        builder.setDocString(toString());
+        builder.addOpsetImport(ONNXOperators.getOpsetProto());
+        builder.setIrVersion(6);
+        return builder.build();
+    }
+
+    /**
+     * Builds a TensorProto containing the supplied DenseMatrix.
+     * @param context The ONNX context for naming.
+     * @return The linear weight TensorProto.
+     */
+    protected OnnxMl.TensorProto matrixBuilder(ONNXContext context, String name, DenseMatrix matrix, boolean transpose) {
+        OnnxMl.TensorProto.Builder matrixBuilder = OnnxMl.TensorProto.newBuilder();
+        matrixBuilder.setName(context.generateUniqueName(name));
+        int dim1, dim2;
+        if (transpose) {
+            dim1 = matrix.getDimension2Size();
+            dim2 = matrix.getDimension1Size();
+        } else {
+            dim1 = matrix.getDimension1Size();
+            dim2 = matrix.getDimension2Size();
+        }
+        matrixBuilder.addDims(dim1);
+        matrixBuilder.addDims(dim2);
+        matrixBuilder.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
+        ByteBuffer buffer = ByteBuffer.allocate(dim1 * dim2 * 4).order(ByteOrder.LITTLE_ENDIAN);
+        FloatBuffer floatBuffer = buffer.asFloatBuffer();
+        for (int i = 0; i < dim1; i++) {
+            for (int j = 0; j < dim2; j++) {
+                if (transpose) {
+                    floatBuffer.put((float) matrix.get(j, i));
+                } else {
+                    floatBuffer.put((float) matrix.get(i, j));
+                }
+            }
+        }
+        floatBuffer.rewind();
+        matrixBuilder.setRawData(ByteString.copyFrom(buffer));
+        return matrixBuilder.build();
+    }
+
+    /**
+     * Builds a TensorProto containing the biases for this Factorization Machine.
+     * @param context The ONNX context for naming.
+     * @return The bias TensorProto.
+     */
+    protected OnnxMl.TensorProto biasBuilder(ONNXContext context) {
+        OnnxMl.TensorProto.Builder biasBuilder = OnnxMl.TensorProto.newBuilder();
+        biasBuilder.setName(context.generateUniqueName("fm_biases"));
+        biasBuilder.addDims(outputIDInfo.size());
+        biasBuilder.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
+        ByteBuffer buffer = ByteBuffer.allocate(outputIDInfo.size()*4).order(ByteOrder.LITTLE_ENDIAN);
+        FloatBuffer floatBuffer = buffer.asFloatBuffer();
+        DenseVector biases = (DenseVector) modelParameters.get()[0];
+        for (int i = 0; i < biases.size(); i++) {
+            floatBuffer.put((float)biases.get(i));
+        }
+        floatBuffer.rewind();
+        biasBuilder.setRawData(ByteString.copyFrom(buffer));
+        return biasBuilder.build();
+    }
 }
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
index 32ce4c831..02134a51c 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
@@ -94,6 +94,8 @@ private FMParameters(Tensor[] weights, int numFactors) {
      */
     private void initializeMatrix(SplittableRandom rng, double variance, DenseMatrix matrix) {
         // This is to get a nextGaussian method. In Java 17 we can use the SplittableRandom directly
+        // though that will prevent reproducibility as there will be different draws from the RNG and
+        // the PRNG algorithm is different.
         Random innerRNG = new Random(rng.nextLong());
         int dim1 = matrix.getDimension1Size();
         int dim2 = matrix.getDimension2Size();
@@ -105,7 +107,8 @@ private void initializeMatrix(SplittableRandom rng, double variance, DenseMatrix
     }
 
     /**
-     * Generates an unnormalised prediction by leftMultiply'ing the weights with the incoming features.
+     * Generates an unnormalised prediction by multiplying the weights with the incoming features,
+     * adding the bias and adding the feature factors.
      *
      * @param example A feature vector
      * @return A {@link DenseVector} containing a score for each label.

From a8af4526afac0ebb784853e7f2d7b0a518bcdfaa Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sat, 25 Sep 2021 23:48:45 -0400
Subject: [PATCH 18/31] Adding optional inputs to ONNXOperators to support
 GEMM.

---
 .../java/org/tribuo/onnx/ONNXOperators.java   | 41 ++++++++++++++++---
 1 file changed, 36 insertions(+), 5 deletions(-)

diff --git a/Core/src/main/java/org/tribuo/onnx/ONNXOperators.java b/Core/src/main/java/org/tribuo/onnx/ONNXOperators.java
index d93d1b58b..1c054ab78 100644
--- a/Core/src/main/java/org/tribuo/onnx/ONNXOperators.java
+++ b/Core/src/main/java/org/tribuo/onnx/ONNXOperators.java
@@ -39,7 +39,7 @@ public enum ONNXOperators {
     /**
      * Concatenates tensors.
      */
-    CONCAT("Constant",VARIADIC_INPUT,1, Collections.singletonList(
+    CONCAT("Concat",VARIADIC_INPUT,1, Collections.singletonList(
             new ONNXAttribute("axis", OnnxMl.AttributeProto.AttributeType.INT, true)
     )),
     /**
@@ -98,7 +98,9 @@ public enum ONNXOperators {
             new ONNXAttribute("keepdims", OnnxMl.AttributeProto.AttributeType.INT, false)
     )),
     /**
-     * General Matrix Multiply: alpha*AB + beta*C.
+     * General Matrix Multiply: {@code alpha*AB + beta*C}.
+     * <p>
+     * The {@code C} input is optional, and if not supplied is treated as zero.
      * <ul>
      *     <li>{@code alpha} defaults to 1.0</li>
      *     <li>{@code beta} defaults to 1.0</li>
@@ -106,7 +108,7 @@ public enum ONNXOperators {
      *     <li>{@code transB} defaults to 0 (i.e., not transposed)</li>
      * </ul>
      */
-    GEMM("Gemm",3,1, Arrays.asList(
+    GEMM("Gemm",2,1, 1, Arrays.asList(
             new ONNXAttribute("alpha", OnnxMl.AttributeProto.AttributeType.FLOAT,false),
             new ONNXAttribute("beta", OnnxMl.AttributeProto.AttributeType.FLOAT,false),
             new ONNXAttribute("transA", OnnxMl.AttributeProto.AttributeType.INT,false),
@@ -121,6 +123,10 @@ public enum ONNXOperators {
      * The number of inputs.
      */
     public final int numInputs;
+    /**
+     * The number of optional inputs.
+     */
+    public final int numOptionalInputs;
     /**
      * The number of outputs.
      */
@@ -146,8 +152,20 @@ public enum ONNXOperators {
      * @param numOutputs The number of outputs.
      */
     private ONNXOperators(String value, int numInputs, int numOutputs) {
+        this(value,numInputs,0,numOutputs);
+    }
+
+    /**
+     * Builds an operator without attributes and with optional inputs.
+     * @param value The operator name.
+     * @param numInputs The number of inputs.
+     * @param numOptionalInputs The number of optional inputs.
+     * @param numOutputs The number of outputs.
+     */
+    private ONNXOperators(String value, int numInputs, int numOptionalInputs, int numOutputs) {
         this.opName = value;
         this.numInputs = numInputs;
+        this.numOptionalInputs = numOptionalInputs;
         this.numOutputs = numOutputs;
         this.attributes = Collections.emptyMap();
         this.mandatoryAttributeNames = Collections.emptySet();
@@ -161,8 +179,21 @@ private ONNXOperators(String value, int numInputs, int numOutputs) {
      * @param attributes The attributes.
      */
     private ONNXOperators(String value, int numInputs, int numOutputs, List<ONNXAttribute> attributes) {
+        this(value,numInputs,0,numOutputs,attributes);
+    }
+
+    /**
+     * Builds an operator with attributes and optional inputs.
+     * @param value The operator name.
+     * @param numInputs The number of inputs.
+     * @param numOptionalInputs The number of optional inputs.
+     * @param numOutputs The number of outputs.
+     * @param attributes The attributes.
+     */
+    private ONNXOperators(String value, int numInputs, int numOptionalInputs, int numOutputs, List<ONNXAttribute> attributes) {
         this.opName = value;
         this.numInputs = numInputs;
+        this.numOptionalInputs = numOptionalInputs;
         this.numOutputs = numOutputs;
         Map<String,ONNXAttribute> attributeMap = new HashMap<>();
         Set<String> attributeSet = new HashSet<>();
@@ -254,8 +285,8 @@ public OnnxMl.NodeProto build(ONNXContext context, String[] inputs, String[] out
      * @return The NodeProto.
      */
     public OnnxMl.NodeProto build(ONNXContext context, String[] inputs, String[] outputs, Map<String,Object> attributeValues) {
-        if ((numInputs != VARIADIC_INPUT) && (inputs.length != numInputs)) {
-            throw new IllegalArgumentException("Expected " + numInputs + " inputs, but received " + inputs.length);
+        if ((numInputs != VARIADIC_INPUT) && ((inputs.length < numInputs) || (inputs.length > numInputs + numOptionalInputs))) {
+            throw new IllegalArgumentException("Expected " + numInputs + " inputs, with " + numOptionalInputs + " optional inputs, but received " + inputs.length);
         }
         if (outputs.length != numOutputs) {
             throw new IllegalArgumentException("Expected " + numOutputs + " outputs, but received " + outputs.length);

From f2a9795c633bf5850cd32f2d9c0c8aefb639bf30 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sat, 25 Sep 2021 23:49:13 -0400
Subject: [PATCH 19/31] Fixing FMClassificationModel onnx export so it
 executes. The model produced has not yet been tested.

---
 .../tribuo/classification/sgd/fm/FMClassificationModel.java   | 4 ++--
 Math/src/main/java/org/tribuo/math/util/Normalizer.java       | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
index 3db74031b..34a1f2278 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
@@ -195,7 +195,7 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
             OnnxMl.NodeProto sumOverEmbeddings = ONNXOperators.REDUCE_SUM.build(context,
                     subtract.getOutput(0),
                     context.generateUniqueName("sum_over_embeddings"),
-                    Collections.singletonMap("axes",1));
+                    Collections.singletonMap("axes",new int[]{1}));
             graphBuilder.addNode(sumOverEmbeddings);
             // Divide by 2
             OnnxMl.NodeProto scaledInteraction = ONNXOperators.DIV.build(context,
@@ -208,7 +208,7 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
 
         // Make concat
         OnnxMl.NodeProto concat = ONNXOperators.CONCAT.build(context, embeddingOutputs, context.generateUniqueName("fm_concat"),
-                    Collections.singletonMap("axes",1)
+                    Collections.singletonMap("axis",1)
                 );
         graphBuilder.addNode(concat);
 
diff --git a/Math/src/main/java/org/tribuo/math/util/Normalizer.java b/Math/src/main/java/org/tribuo/math/util/Normalizer.java
index 597d01470..38cf780bc 100644
--- a/Math/src/main/java/org/tribuo/math/util/Normalizer.java
+++ b/Math/src/main/java/org/tribuo/math/util/Normalizer.java
@@ -70,7 +70,7 @@ public List<OnnxMl.NodeProto> exportNormalizer(ONNXContext context, String input
         List<OnnxMl.NodeProto> protos = new ArrayList<>();
 
         String minOutput = context.generateUniqueName("min_output");
-        OnnxMl.NodeProto min = ONNXOperators.REDUCE_MIN.build(context,input,minOutput,Collections.singletonMap("axes",1));
+        OnnxMl.NodeProto min = ONNXOperators.REDUCE_MIN.build(context,input,minOutput,Collections.singletonMap("axes",new int[]{1}));
         protos.add(min);
 
         String subOutput = context.generateUniqueName("sub_output");
@@ -78,7 +78,7 @@ public List<OnnxMl.NodeProto> exportNormalizer(ONNXContext context, String input
         protos.add(sub);
 
         String sumOutput = context.generateUniqueName("sum_output");
-        OnnxMl.NodeProto sum = ONNXOperators.REDUCE_SUM.build(context,subOutput,sumOutput,Collections.singletonMap("axes",1));
+        OnnxMl.NodeProto sum = ONNXOperators.REDUCE_SUM.build(context,subOutput,sumOutput,Collections.singletonMap("axes",new int[]{1}));
         protos.add(sum);
 
         OnnxMl.NodeProto div = ONNXOperators.DIV.build(context,new String[]{subOutput,sumOutput},output);

From c65f6b6c4d862f6f738aef9a30b68bf1ff81e1d8 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 26 Sep 2021 12:13:12 -0400
Subject: [PATCH 20/31] Tidying up javadoc and fixing a broken merge.

---
 .../sgd/fm/FMClassificationOptions.java       | 21 ++++++++++
 .../common/liblinear/LibLinearModel.java      |  2 +-
 .../tribuo/common/sgd/AbstractFMModel.java    |  5 ++-
 .../tribuo/common/sgd/AbstractSGDModel.java   |  6 +++
 .../java/org/tribuo/data/DataOptions.java     | 18 ++++-----
 .../sgd/fm/FMMultiLabelOptions.java           | 39 ++++++++++++++++++-
 .../liblinear/LibLinearRegressionModel.java   |  2 +-
 7 files changed, 78 insertions(+), 15 deletions(-)

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
index f31739094..5fad87098 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
@@ -49,18 +49,39 @@ public enum LossEnum {
 
     public GradientOptimiserOptions sgoOptions;
 
+    /**
+     * Number of SGD epochs.
+     */
     @Option(longName = "fm-epochs", usage = "Number of SGD epochs.")
     public int fmEpochs = 5;
+    /**
+     * Loss function.
+     */
     @Option(longName = "fm-objective", usage = "Loss function.")
     public LossEnum fmObjective = LossEnum.LOG;
+    /**
+     * Log the objective after n examples.
+     */
     @Option(longName = "fm-logging-interval", usage = "Log the objective after <int> examples.")
     public int fmLoggingInterval = 100;
+    /**
+     * Minibatch size.
+     */
     @Option(longName = "fm-minibatch-size", usage = "Minibatch size.")
     public int fmMinibatchSize = 1;
+    /**
+     * Sets the random seed for the FMClassificationTrainer.
+     */
     @Option(longName = "fm-seed", usage = "Sets the random seed for the FMClassificationTrainer.")
     private long fmSeed = Trainer.DEFAULT_SEED;
+    /**
+     * Factor size.
+     */
     @Option(longName = "fm-factor-size", usage = "Factor size.")
     public int fmFactorSize = 6;
+    /**
+     * Variance of the initialization gaussian.
+     */
     @Option(longName = "fm-variance", usage = "Variance of the initialization gaussian.")
     public double fmVariance = 0.1;
 
diff --git a/Common/LibLinear/src/main/java/org/tribuo/common/liblinear/LibLinearModel.java b/Common/LibLinear/src/main/java/org/tribuo/common/liblinear/LibLinearModel.java
index 0382f1dfc..5d2f42c8c 100644
--- a/Common/LibLinear/src/main/java/org/tribuo/common/liblinear/LibLinearModel.java
+++ b/Common/LibLinear/src/main/java/org/tribuo/common/liblinear/LibLinearModel.java
@@ -157,7 +157,7 @@ protected static de.bwaldvogel.liblinear.Model copyModel(de.bwaldvogel.liblinear
      * <p>
      * If it becomes a problem then we could cache the feature weights in the
      * model.
-     * <p>
+     *
      * @param e The example.
      * @param featureWeights The per dimension feature weights.
      * @return An excuse for this example.
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
index d09f45e4e..107c855b0 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
@@ -192,7 +192,10 @@ protected OnnxMl.ModelProto innerExportONNXModel(OnnxMl.GraphProto graph, String
     /**
      * Builds a TensorProto containing the supplied DenseMatrix.
      * @param context The ONNX context for naming.
-     * @return The linear weight TensorProto.
+     * @param name The name for this tensor proto.
+     * @param matrix The matrix to store.
+     * @param transpose Should the matrix be transposed into the tensor?
+     * @return The matrix TensorProto.
      */
     protected OnnxMl.TensorProto matrixBuilder(ONNXContext context, String name, DenseMatrix matrix, boolean transpose) {
         OnnxMl.TensorProto.Builder matrixBuilder = OnnxMl.TensorProto.newBuilder();
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDModel.java
index eb1a9250a..1cd5a132d 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDModel.java
@@ -98,7 +98,13 @@ public FeedForwardParameters getModelParameters() {
      * A nominal tuple used to capture the prediction and the number of active features used by the model.
      */
     protected static final class PredAndActive {
+        /**
+         * The vector prediction.
+         */
         public final DenseVector prediction;
+        /**
+         * The number of active features used in the prediction.
+         */
         public final int numActiveFeatures;
 
         PredAndActive(DenseVector prediction, int numActiveFeatures) {
diff --git a/Data/src/main/java/org/tribuo/data/DataOptions.java b/Data/src/main/java/org/tribuo/data/DataOptions.java
index ad494e3e4..8b64696f1 100644
--- a/Data/src/main/java/org/tribuo/data/DataOptions.java
+++ b/Data/src/main/java/org/tribuo/data/DataOptions.java
@@ -181,21 +181,18 @@ public String getOptionsDescription() {
      */
     @Option(charName = 'v', longName = "testing-file", usage = "Path to the testing file.")
     public Path testingPath;
+    /**
+     * Scales the features to the range 0-1 independently.
+     */
     @Option(longName="scale-features",usage="Scales the features to the range 0-1 independently.")
     public boolean scaleFeatures;
-    @Option(longName="scale-including-zeros",usage="Includes implicit zeros in the scale range calculation")
+    /**
+     * Includes implicit zeros in the scale range calculation.
+     */
+    @Option(longName="scale-including-zeros",usage="Includes implicit zeros in the scale range calculation.")
     public boolean scaleIncZeros;
 
     /**
-<<<<<<< HEAD
-     * Loads the datasets specified in this options.
-     * @param outputFactory The output factory to use.
-     * @param <T> The type of the dataset.
-     * @return A pair of training and test datasets.
-     * @throws IOException If the datasets failed to load.
-     */
-    public <T extends Output<T>> Pair<Dataset<T>, Dataset<T>> load(OutputFactory<T> outputFactory) throws IOException {
-=======
      * Loads the training and testing data from {@link #trainingPath} and {@link #testingPath}
      * according to the other parameters specified in this class.
      * @param outputFactory The output factory to use to process the inputs.
@@ -205,7 +202,6 @@ public <T extends Output<T>> Pair<Dataset<T>, Dataset<T>> load(OutputFactory<T>
      * @throws IOException If the paths could not be loaded.
      */
     public <T extends Output<T>> Pair<Dataset<T>,Dataset<T>> load(OutputFactory<T> outputFactory) throws IOException {
->>>>>>> c459390 (Adding another rescaling option to DataOptions.)
         logger.info(String.format("Loading data from %s", trainingPath));
         Dataset<T> train;
         Dataset<T> test;
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
index e3366a99f..224aeec7f 100644
--- a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
@@ -36,22 +36,55 @@ public class FMMultiLabelOptions implements Options {
     /**
      * Available loss types.
      */
-    public enum LossEnum {HINGE, SIGMOID}
+    public enum LossEnum {
+        /**
+         * Hinge loss.
+         */
+        HINGE,
+        /**
+         * Log loss, i.e., binary cross-entropy.
+         */
+        SIGMOID
+    }
 
+    /**
+     * The gradient descent optimiser options.
+     */
     public GradientOptimiserOptions sgoOptions;
 
+    /**
+     * Number of SGD epochs.
+     */
     @Option(longName = "fm-epochs", usage = "Number of SGD epochs.")
     public int fmEpochs = 5;
+    /**
+     * Loss function.
+     */
     @Option(longName = "fm-objective", usage = "Loss function.")
     public LossEnum fmObjective = LossEnum.SIGMOID;
+    /**
+     * Log the objective after n examples.
+     */
     @Option(longName = "fm-logging-interval", usage = "Log the objective after <int> examples.")
     public int fmLoggingInterval = 100;
+    /**
+     * Minibatch size.
+     */
     @Option(longName = "fm-minibatch-size", usage = "Minibatch size.")
     public int fmMinibatchSize = 1;
+    /**
+     * Sets the random seed for the FMMultiLabelTrainer.
+     */
     @Option(longName = "fm-seed", usage = "Sets the random seed for the FMMultiLabelTrainer.")
     private long fmSeed = Trainer.DEFAULT_SEED;
+    /**
+     * Factor size.
+     */
     @Option(longName = "fm-factor-size", usage = "Factor size.")
     public int fmFactorSize = 6;
+    /**
+     * Variance of the initialization gaussian.
+     */
     @Option(longName = "fm-variance", usage = "Variance of the initialization gaussian.")
     public double fmVariance = 0.1;
 
@@ -70,6 +103,10 @@ public MultiLabelObjective getLoss() {
         }
     }
 
+    /**
+     * Gets the trainer as configured by these options.
+     * @return The trainer.
+     */
     public FMMultiLabelTrainer getTrainer() {
         logger.info(String.format("Set logging interval to %d", fmLoggingInterval));
         return new FMMultiLabelTrainer(getLoss(), sgoOptions.getOptimiser(), fmEpochs, fmLoggingInterval,
diff --git a/Regression/LibLinear/src/main/java/org/tribuo/regression/liblinear/LibLinearRegressionModel.java b/Regression/LibLinear/src/main/java/org/tribuo/regression/liblinear/LibLinearRegressionModel.java
index d6b0b3842..712688f2e 100644
--- a/Regression/LibLinear/src/main/java/org/tribuo/regression/liblinear/LibLinearRegressionModel.java
+++ b/Regression/LibLinear/src/main/java/org/tribuo/regression/liblinear/LibLinearRegressionModel.java
@@ -156,7 +156,7 @@ protected double[][] getFeatureWeights() {
      * <p>
      * If it becomes a problem then we could cache the feature weights in the
      * model.
-     * <p>
+     *
      * @param e The example.
      * @param allFeatureWeights The feature weights.
      * @return An excuse for this example.

From 2d9160099eda7be0b0f298a82be7ef81c1ab5169 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 26 Sep 2021 20:48:37 -0400
Subject: [PATCH 21/31] Adding a name to the generated onnx graph and some
 javadoc cleanup

---
 .../sgd/fm/FMClassificationModel.java         |  1 +
 .../org/tribuo/regression/sgd/TrainTest.java  |  8 ++--
 .../tribuo/regression/sgd/fm/TrainTest.java   | 44 ++++++++++++++-----
 3 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
index 34a1f2278..9d1e94b3d 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
@@ -117,6 +117,7 @@ public OnnxMl.ModelProto exportONNXModel(String domain, long modelVersion) {
     @Override
     public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
         OnnxMl.GraphProto.Builder graphBuilder = OnnxMl.GraphProto.newBuilder();
+        graphBuilder.setName("FMClassificationModel");
 
         // Make inputs and outputs
         OnnxMl.TypeProto inputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,featureIDMap.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/TrainTest.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/TrainTest.java
index f8ea6ca11..b35426abf 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/TrainTest.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/TrainTest.java
@@ -86,22 +86,22 @@ public String getOptionsDescription() {
         /**
          * Number of SGD epochs. Defaults to 5.
          */
-        @Option(charName = 'i', longName = "epochs", usage = "Number of SGD epochs. Defaults to 5.")
+        @Option(charName = 'i', longName = "epochs", usage = "Number of SGD epochs.")
         public int epochs = 5;
         /**
          * Loss function. Defaults to SQUARED.
          */
-        @Option(charName = 'o', longName = "objective", usage = "Loss function. Defaults to SQUARED.")
+        @Option(charName = 'o', longName = "objective", usage = "Loss function.")
         public LossEnum loss = LossEnum.SQUARED;
         /**
          * Log the objective after n examples. Defaults to 100.
          */
-        @Option(charName = 'p', longName = "logging-interval", usage = "Log the objective after <int> examples. Defaults to 100.")
+        @Option(charName = 'p', longName = "logging-interval", usage = "Log the objective after <int> examples.")
         public int loggingInterval = 100;
         /**
          * Minibatch size. Defaults to 1.
          */
-        @Option(charName = 'z', longName = "minibatch-size", usage = "Minibatch size. Defaults to 1.")
+        @Option(charName = 'z', longName = "minibatch-size", usage = "Minibatch size.")
         public int minibatchSize = 1;
     }
 
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
index ae3e16b8a..7f407e594 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
@@ -32,7 +32,6 @@
 import org.tribuo.regression.Regressor;
 import org.tribuo.regression.evaluation.RegressionEvaluation;
 import org.tribuo.regression.sgd.RegressionObjective;
-import org.tribuo.regression.sgd.linear.LinearSGDTrainer;
 import org.tribuo.regression.sgd.objectives.AbsoluteLoss;
 import org.tribuo.regression.sgd.objectives.Huber;
 import org.tribuo.regression.sgd.objectives.SquaredLoss;
@@ -44,7 +43,7 @@
 /**
  * Build and run a regression factorization machine for a standard dataset.
  */
-public class TrainTest {
+public final class TrainTest {
 
     private static final Logger logger = Logger.getLogger(TrainTest.class.getName());
 
@@ -69,28 +68,53 @@ public enum LossEnum {
     /**
      * Command line options.
      */
-    public static class SGDOptions implements Options {
+    public static class FMRegressionOptions implements Options {
         @Override
         public String getOptionsDescription() {
             return "Trains and tests a linear SGD regression model on the specified datasets.";
         }
+        /**
+         * The dataset loading options.
+         */
         public DataOptions general;
+        /**
+         * Gradient descent options.
+         */
         public GradientOptimiserOptions gradientOptions;
 
-        @Option(charName='i',longName="epochs",usage="Number of SGD epochs. Defaults to 5.")
+        /**
+         * Number of SGD epochs. Defaults to 5.
+         */
+        @Option(charName = 'i', longName = "epochs", usage = "Number of SGD epochs.")
         public int epochs = 5;
-        @Option(charName='o',longName="objective",usage="Loss function. Defaults to SQUARED.")
+        /**
+         * Loss function. Defaults to SQUARED.
+         */
+        @Option(charName = 'o', longName = "objective", usage = "Loss function.")
         public LossEnum loss = LossEnum.SQUARED;
-        @Option(charName='p',longName="logging-interval",usage="Log the objective after <int> examples. Defaults to 100.")
+        /**
+         * Log the objective after n examples. Defaults to 100.
+         */
+        @Option(charName = 'p', longName = "logging-interval", usage = "Log the objective after <int> examples.")
         public int loggingInterval = 100;
-        @Option(charName='z',longName="minibatch-size",usage="Minibatch size. Defaults to 1.")
+        /**
+         * Minibatch size. Defaults to 1.
+         */
+        @Option(charName = 'z', longName = "minibatch-size", usage = "Minibatch size.")
         public int minibatchSize = 1;
+        /**
+         * Factor size.
+         */
         @Option(charName = 'd', longName = "factor-size", usage = "Factor size.")
         public int factorSize = 6;
-        @Option(charName = 'k', longName = "l2-penalty", usage = "L2 regularization penalty.")
-        public double l2 = 0.001;
+        /**
+         * Variance of the initialization gaussian.
+         */
         @Option(longName = "variance", usage = "Variance of the initialization gaussian.")
         public double variance = 0.1;
+        /**
+         * Standardise the output regressors before model fitting.
+         */
         @Option(longName = "standardise", usage = "Standardise the output regressors before model fitting.")
         public boolean standardise = false;
     }
@@ -105,7 +129,7 @@ public static void main(String[] args) throws IOException {
         // Use the labs format logging.
         LabsLogFormatter.setAllLogFormatters();
 
-        SGDOptions o = new SGDOptions();
+        FMRegressionOptions o = new FMRegressionOptions();
         ConfigurationManager cm;
         try {
             cm = new ConfigurationManager(args,o);

From bbe2be6ebc14213a52c1321e787e7c1401abfaca Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 26 Sep 2021 21:12:12 -0400
Subject: [PATCH 22/31] Fixing ONNX export so we don't transpose already
 transposed matrices.

---
 .../tribuo/classification/sgd/fm/FMClassificationModel.java | 6 ++----
 .../tribuo/classification/sgd/fm/TestFMClassification.java  | 4 ++--
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
index 9d1e94b3d..68059f58a 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
@@ -168,8 +168,7 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
             // Feature matrix * embedding matrix = batch_size, embedding dim
             OnnxMl.NodeProto gemmFeatureEmb = ONNXOperators.GEMM.build(context,
                     new String[]{inputValueProto.getName(),embeddingProtos[i].getName()},
-                    context.generateUniqueName("gemm_input_emb"),
-                    Collections.singletonMap("transB",1));
+                    context.generateUniqueName("gemm_input_emb"));
             graphBuilder.addNode(gemmFeatureEmb);
             // Square the output
             OnnxMl.NodeProto powFeatureEmb = ONNXOperators.POW.build(context,
@@ -184,8 +183,7 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
             // squared features * squared embeddings
             OnnxMl.NodeProto gemmSquaredFeatureSquaredEmb = ONNXOperators.GEMM.build(context,
                     new String[]{featurePow.getOutput(0),powEmb.getOutput(0)},
-                    context.generateUniqueName("gemm_squared_input_squared_emb"),
-                    Collections.singletonMap("transB",1));
+                    context.generateUniqueName("gemm_squared_input_squared_emb"));
             graphBuilder.addNode(gemmSquaredFeatureSquaredEmb);
             // squared product subtract product of squares
             OnnxMl.NodeProto subtract = ONNXOperators.SUB.build(context,
diff --git a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
index b50f6ef14..4c70429ff 100644
--- a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
+++ b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
@@ -161,13 +161,13 @@ public void testOnnxSerialization() throws IOException, OrtException {
                 Prediction<Label> tribuo = nativePredictions.get(i);
                 Prediction<Label> external = onnxPredictions.get(i);
                 assertEquals(tribuo.getOutput().getLabel(), external.getOutput().getLabel());
-                assertEquals(tribuo.getOutput().getScore(), external.getOutput().getScore(), 1e-6);
+                assertEquals(tribuo.getOutput().getScore(), external.getOutput().getScore(), 1e-5);
                 for (Map.Entry<String, Label> l : tribuo.getOutputScores().entrySet()) {
                     Label other = external.getOutputScores().get(l.getKey());
                     if (other == null) {
                         fail("Failed to find label " + l.getKey() + " in ORT prediction.");
                     } else {
-                        assertEquals(l.getValue().getScore(), other.getScore(), 1e-6);
+                        assertEquals(l.getValue().getScore(), other.getScore(), 1e-5);
                     }
                 }
             }

From 7cd61b815ccd71a4ea93189013597f952bf18d57 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 26 Sep 2021 21:28:06 -0400
Subject: [PATCH 23/31] Moving FM ONNX export to AbstractFMModel and wiring it
 into multilabel and regression FMs.

---
 .../sgd/fm/FMClassificationModel.java         |  95 +----------
 .../tribuo/common/sgd/AbstractFMModel.java    | 155 +++++++++++++++---
 .../multilabel/sgd/fm/FMMultiLabelModel.java  |  45 ++++-
 .../multilabel/sgd/fm/TestFMMultiLabel.java   |  71 ++++++++
 .../multilabel/sgd/linear/TestSGDLinear.java  |   3 +-
 .../regression/sgd/fm/FMRegressionModel.java  |  41 ++++-
 .../regression/sgd/fm/TestFMRegression.java   |  69 +++++++-
 .../regression/sgd/linear/TestSGDLinear.java  |   2 +-
 8 files changed, 359 insertions(+), 122 deletions(-)

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
index 68059f58a..e9a44f29f 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
@@ -24,18 +24,14 @@
 import org.tribuo.classification.Label;
 import org.tribuo.common.sgd.AbstractFMModel;
 import org.tribuo.common.sgd.FMParameters;
-import org.tribuo.math.la.DenseMatrix;
 import org.tribuo.math.la.DenseVector;
-import org.tribuo.math.la.Tensor;
 import org.tribuo.math.util.VectorNormalizer;
 import org.tribuo.onnx.ONNXContext;
 import org.tribuo.onnx.ONNXExportable;
-import org.tribuo.onnx.ONNXOperators;
 import org.tribuo.onnx.ONNXShape;
 import org.tribuo.onnx.ONNXUtils;
 import org.tribuo.provenance.ModelProvenance;
 
-import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -127,96 +123,11 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
         OnnxMl.ValueInfoProto outputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(outputType).setName("output").build();
         graphBuilder.addOutput(outputValueProto);
 
-        Tensor[] modelParams = modelParameters.get();
-
-        // Add constants
-        OnnxMl.TensorProto twoConst = OnnxMl.TensorProto.newBuilder()
-                .setName(context.generateUniqueName("two_const"))
-                .setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber())
-                .addFloatData(2.0f)
-                .build();
-        graphBuilder.addInitializer(twoConst);
-
-        // Add weights
-        OnnxMl.TensorProto weightInitializerProto = matrixBuilder(context, "fm_linear_weights", (DenseMatrix) modelParams[1], true);
-        graphBuilder.addInitializer(weightInitializerProto);
-
-        // Add biases
-        OnnxMl.TensorProto biasInitializerProto = biasBuilder(context);
-        graphBuilder.addInitializer(biasInitializerProto);
-
-        // Add embedding vectors
-        OnnxMl.TensorProto[] embeddingProtos = new OnnxMl.TensorProto[outputIDInfo.size()];
-        for (int i = 0; i < outputIDInfo.size(); i++) {
-            embeddingProtos[i] = matrixBuilder(context, "fm_embedding_"+i, (DenseMatrix) modelParams[i+2], true);
-            graphBuilder.addInitializer(embeddingProtos[i]);
-        }
-
-        // Make gemm
-        String[] gemmInputs = new String[]{inputValueProto.getName(),weightInitializerProto.getName(),biasInitializerProto.getName()};
-        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context,gemmInputs,context.generateUniqueName("gemm_output"));
-        graphBuilder.addNode(gemm);
-
-        // Make feature pow
-        OnnxMl.NodeProto featurePow = ONNXOperators.POW.build(context,new String[]{inputValueProto.getName(),twoConst.getName()},
-                context.generateUniqueName("feature_pow"));
-        graphBuilder.addNode(featurePow);
-
-        // Make interaction terms
-        String[] embeddingOutputs = new String[outputIDInfo.size()];
-        for (int i = 0; i < outputIDInfo.size(); i++) {
-            // Feature matrix * embedding matrix = batch_size, embedding dim
-            OnnxMl.NodeProto gemmFeatureEmb = ONNXOperators.GEMM.build(context,
-                    new String[]{inputValueProto.getName(),embeddingProtos[i].getName()},
-                    context.generateUniqueName("gemm_input_emb"));
-            graphBuilder.addNode(gemmFeatureEmb);
-            // Square the output
-            OnnxMl.NodeProto powFeatureEmb = ONNXOperators.POW.build(context,
-                    new String[]{gemmFeatureEmb.getOutput(0),twoConst.getName()},
-                    context.generateUniqueName("pow_input_emb"));
-            graphBuilder.addNode(powFeatureEmb);
-            // Square the embeddings
-            OnnxMl.NodeProto powEmb = ONNXOperators.POW.build(context,
-                    new String[]{embeddingProtos[i].getName(),twoConst.getName()},
-                    context.generateUniqueName("pow_emb"));
-            graphBuilder.addNode(powEmb);
-            // squared features * squared embeddings
-            OnnxMl.NodeProto gemmSquaredFeatureSquaredEmb = ONNXOperators.GEMM.build(context,
-                    new String[]{featurePow.getOutput(0),powEmb.getOutput(0)},
-                    context.generateUniqueName("gemm_squared_input_squared_emb"));
-            graphBuilder.addNode(gemmSquaredFeatureSquaredEmb);
-            // squared product subtract product of squares
-            OnnxMl.NodeProto subtract = ONNXOperators.SUB.build(context,
-                    new String[]{powFeatureEmb.getOutput(0),gemmSquaredFeatureSquaredEmb.getOutput(0)},
-                    context.generateUniqueName("squared_prod_subtract_prod_of_squares"));
-            graphBuilder.addNode(subtract);
-            // sum over embedding dimensions
-            OnnxMl.NodeProto sumOverEmbeddings = ONNXOperators.REDUCE_SUM.build(context,
-                    subtract.getOutput(0),
-                    context.generateUniqueName("sum_over_embeddings"),
-                    Collections.singletonMap("axes",new int[]{1}));
-            graphBuilder.addNode(sumOverEmbeddings);
-            // Divide by 2
-            OnnxMl.NodeProto scaledInteraction = ONNXOperators.DIV.build(context,
-                    new String[]{sumOverEmbeddings.getOutput(0),twoConst.getName()},
-                    context.generateUniqueName("scaled_interaction"));
-            graphBuilder.addNode(scaledInteraction);
-            // Store the output name
-            embeddingOutputs[i] = scaledInteraction.getOutput(0);
-        }
-
-        // Make concat
-        OnnxMl.NodeProto concat = ONNXOperators.CONCAT.build(context, embeddingOutputs, context.generateUniqueName("fm_concat"),
-                    Collections.singletonMap("axis",1)
-                );
-        graphBuilder.addNode(concat);
-
-        // Add to gemm
-        OnnxMl.NodeProto addGemmConcat = ONNXOperators.ADD.build(context, new String[]{gemm.getOutput(0),concat.getOutput(0)}, context.generateUniqueName("fm_output"));
-        graphBuilder.addNode(addGemmConcat);
+        // Build the output neutral bits of the onnx graph
+        String outputName = generateONNXGraph(context, graphBuilder, inputValueProto.getName());
 
         // Make output normalizer
-        List<OnnxMl.NodeProto> normalizerProtos = normalizer.exportNormalizer(context,addGemmConcat.getOutput(0),"output");
+        List<OnnxMl.NodeProto> normalizerProtos = normalizer.exportNormalizer(context,outputName,"output");
         if (normalizerProtos.isEmpty()) {
             throw new IllegalArgumentException("Normalizer " + normalizer.getClass() + " cannot be exported in ONNX models.");
         } else {
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
index 107c855b0..0c228dcf5 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
@@ -61,11 +61,12 @@ public abstract class AbstractFMModel<T extends Output<T>> extends AbstractSGDMo
 
     /**
      * Constructs a factorization machine model trained via SGD.
-     * @param name The model name.
-     * @param provenance The model provenance.
-     * @param featureIDMap The feature domain.
-     * @param outputIDInfo The output domain.
-     * @param parameters The model parameters.
+     *
+     * @param name                   The model name.
+     * @param provenance             The model provenance.
+     * @param featureIDMap           The feature domain.
+     * @param outputIDInfo           The output domain.
+     * @param parameters             The model parameters.
      * @param generatesProbabilities Does this model generate probabilities?
      */
     protected AbstractFMModel(String name, ModelProvenance provenance,
@@ -80,18 +81,18 @@ public Map<String, List<Pair<String, Double>>> getTopFeatures(int n) {
         DenseMatrix baseWeights = (DenseMatrix) modelParameters.get()[1];
         int maxFeatures = n < 0 ? featureIDMap.size() + 1 : n;
 
-        Comparator<Pair<String,Double>> comparator = Comparator.comparingDouble(p -> Math.abs(p.getB()));
+        Comparator<Pair<String, Double>> comparator = Comparator.comparingDouble(p -> Math.abs(p.getB()));
 
         //
         // Use a priority queue to find the top N features.
         int numClasses = baseWeights.getDimension1Size();
         int numFeatures = baseWeights.getDimension2Size();
-        Map<String, List<Pair<String,Double>>> map = new HashMap<>();
+        Map<String, List<Pair<String, Double>>> map = new HashMap<>();
         for (int i = 0; i < numClasses; i++) {
-            PriorityQueue<Pair<String,Double>> q = new PriorityQueue<>(maxFeatures, comparator);
+            PriorityQueue<Pair<String, Double>> q = new PriorityQueue<>(maxFeatures, comparator);
 
             for (int j = 0; j < numFeatures; j++) {
-                Pair<String,Double> curr = new Pair<>(featureIDMap.get(j).getName(), baseWeights.get(i,j));
+                Pair<String, Double> curr = new Pair<>(featureIDMap.get(j).getName(), baseWeights.get(i, j));
 
                 if (q.size() < maxFeatures) {
                     q.offer(curr);
@@ -100,7 +101,7 @@ public Map<String, List<Pair<String, Double>>> getTopFeatures(int n) {
                     q.offer(curr);
                 }
             }
-            Pair<String,Double> curr = new Pair<>(BIAS_FEATURE, biases.get(i));
+            Pair<String, Double> curr = new Pair<>(BIAS_FEATURE, biases.get(i));
 
             if (q.size() < maxFeatures) {
                 q.offer(curr);
@@ -108,7 +109,7 @@ public Map<String, List<Pair<String, Double>>> getTopFeatures(int n) {
                 q.poll();
                 q.offer(curr);
             }
-            List<Pair<String,Double>> b = new ArrayList<>();
+            List<Pair<String, Double>> b = new ArrayList<>();
             while (q.size() > 0) {
                 b.add(q.poll());
             }
@@ -121,18 +122,20 @@ public Map<String, List<Pair<String, Double>>> getTopFeatures(int n) {
 
     /**
      * Returns a copy of the linear weights.
+     *
      * @return The linear weights.
      */
     public DenseMatrix getLinearWeightsCopy() {
-        return ((DenseMatrix)modelParameters.get()[1]).copy();
+        return ((DenseMatrix) modelParameters.get()[1]).copy();
     }
 
     /**
      * Returns a copy of the output dimension biases.
+     *
      * @return The biases.
      */
     public DenseVector getBiasesCopy() {
-        return ((DenseVector)modelParameters.get()[0]).copy();
+        return ((DenseVector) modelParameters.get()[0]).copy();
     }
 
     /**
@@ -140,19 +143,21 @@ public DenseVector getBiasesCopy() {
      * There is one factor matrix per output dimension.
      * The first factor matrix dimension is the factor dimension,
      * the second is the number of features.
+     *
      * @return The factors.
      */
     public Tensor[] getFactorsCopy() {
         Tensor[] params = modelParameters.get();
-        Tensor[] paramCopy = new Tensor[params.length-2];
+        Tensor[] paramCopy = new Tensor[params.length - 2];
         for (int i = 0; i < paramCopy.length; i++) {
-            paramCopy[i] = params[i+2].copy();
+            paramCopy[i] = params[i + 2].copy();
         }
         return paramCopy;
     }
 
     /**
      * Factorization machines don't provide excuses, use an explainer.
+     *
      * @param example The input example.
      * @return Optional.emtpy.
      */
@@ -163,6 +168,7 @@ public Optional<Excuse<T>> getExcuse(Example<T> example) {
 
     /**
      * Gets the name of the indexed output dimension.
+     *
      * @param index The output dimension index.
      * @return The name of the requested output dimension.
      */
@@ -170,8 +176,9 @@ public Optional<Excuse<T>> getExcuse(Example<T> example) {
 
     /**
      * Builds the ModelProto according to the standards for this model.
-     * @param graph The model graph.
-     * @param domain The model domain string.
+     *
+     * @param graph        The model graph.
+     * @param domain       The model domain string.
      * @param modelVersion The model version number.
      * @return The ModelProto.
      */
@@ -191,9 +198,10 @@ protected OnnxMl.ModelProto innerExportONNXModel(OnnxMl.GraphProto graph, String
 
     /**
      * Builds a TensorProto containing the supplied DenseMatrix.
-     * @param context The ONNX context for naming.
-     * @param name The name for this tensor proto.
-     * @param matrix The matrix to store.
+     *
+     * @param context   The ONNX context for naming.
+     * @param name      The name for this tensor proto.
+     * @param matrix    The matrix to store.
      * @param transpose Should the matrix be transposed into the tensor?
      * @return The matrix TensorProto.
      */
@@ -229,6 +237,7 @@ protected OnnxMl.TensorProto matrixBuilder(ONNXContext context, String name, Den
 
     /**
      * Builds a TensorProto containing the biases for this Factorization Machine.
+     *
      * @param context The ONNX context for naming.
      * @return The bias TensorProto.
      */
@@ -237,14 +246,116 @@ protected OnnxMl.TensorProto biasBuilder(ONNXContext context) {
         biasBuilder.setName(context.generateUniqueName("fm_biases"));
         biasBuilder.addDims(outputIDInfo.size());
         biasBuilder.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
-        ByteBuffer buffer = ByteBuffer.allocate(outputIDInfo.size()*4).order(ByteOrder.LITTLE_ENDIAN);
+        ByteBuffer buffer = ByteBuffer.allocate(outputIDInfo.size() * 4).order(ByteOrder.LITTLE_ENDIAN);
         FloatBuffer floatBuffer = buffer.asFloatBuffer();
         DenseVector biases = (DenseVector) modelParameters.get()[0];
         for (int i = 0; i < biases.size(); i++) {
-            floatBuffer.put((float)biases.get(i));
+            floatBuffer.put((float) biases.get(i));
         }
         floatBuffer.rewind();
         biasBuilder.setRawData(ByteString.copyFrom(buffer));
         return biasBuilder.build();
     }
+
+    /**
+     * Constructs the shared stem of the Factorization Machine, used by all output types.
+     * <p>
+     * Writes into the supplied graph builder.
+     *
+     * @param context      The ONNX context.
+     * @param graphBuilder The graph builder.
+     * @return The name of the output.
+     */
+    protected String generateONNXGraph(ONNXContext context, OnnxMl.GraphProto.Builder graphBuilder, String inputName) {
+        Tensor[] modelParams = modelParameters.get();
+
+        // Add constants
+        OnnxMl.TensorProto twoConst = OnnxMl.TensorProto.newBuilder()
+                .setName(context.generateUniqueName("two_const"))
+                .setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber())
+                .addFloatData(2.0f)
+                .build();
+        graphBuilder.addInitializer(twoConst);
+
+        // Add weights
+        OnnxMl.TensorProto weightInitializerProto = matrixBuilder(context, "fm_linear_weights", (DenseMatrix) modelParams[1], true);
+        graphBuilder.addInitializer(weightInitializerProto);
+
+        // Add biases
+        OnnxMl.TensorProto biasInitializerProto = biasBuilder(context);
+        graphBuilder.addInitializer(biasInitializerProto);
+
+        // Add embedding vectors
+        OnnxMl.TensorProto[] embeddingProtos = new OnnxMl.TensorProto[outputIDInfo.size()];
+        for (int i = 0; i < outputIDInfo.size(); i++) {
+            embeddingProtos[i] = matrixBuilder(context, "fm_embedding_" + i, (DenseMatrix) modelParams[i + 2], true);
+            graphBuilder.addInitializer(embeddingProtos[i]);
+        }
+
+        // Make gemm
+        String[] gemmInputs = new String[]{inputName, weightInitializerProto.getName(), biasInitializerProto.getName()};
+        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context, gemmInputs, context.generateUniqueName("gemm_output"));
+        graphBuilder.addNode(gemm);
+
+        // Make feature pow
+        OnnxMl.NodeProto featurePow = ONNXOperators.POW.build(context, new String[]{inputName, twoConst.getName()},
+                context.generateUniqueName("feature_pow"));
+        graphBuilder.addNode(featurePow);
+
+        // Make interaction terms
+        String[] embeddingOutputs = new String[outputIDInfo.size()];
+        for (int i = 0; i < outputIDInfo.size(); i++) {
+            // Feature matrix * embedding matrix = batch_size, embedding dim
+            OnnxMl.NodeProto gemmFeatureEmb = ONNXOperators.GEMM.build(context,
+                    new String[]{inputName, embeddingProtos[i].getName()},
+                    context.generateUniqueName("gemm_input_emb"));
+            graphBuilder.addNode(gemmFeatureEmb);
+            // Square the output
+            OnnxMl.NodeProto powFeatureEmb = ONNXOperators.POW.build(context,
+                    new String[]{gemmFeatureEmb.getOutput(0), twoConst.getName()},
+                    context.generateUniqueName("pow_input_emb"));
+            graphBuilder.addNode(powFeatureEmb);
+            // Square the embeddings
+            OnnxMl.NodeProto powEmb = ONNXOperators.POW.build(context,
+                    new String[]{embeddingProtos[i].getName(), twoConst.getName()},
+                    context.generateUniqueName("pow_emb"));
+            graphBuilder.addNode(powEmb);
+            // squared features * squared embeddings
+            OnnxMl.NodeProto gemmSquaredFeatureSquaredEmb = ONNXOperators.GEMM.build(context,
+                    new String[]{featurePow.getOutput(0), powEmb.getOutput(0)},
+                    context.generateUniqueName("gemm_squared_input_squared_emb"));
+            graphBuilder.addNode(gemmSquaredFeatureSquaredEmb);
+            // squared product subtract product of squares
+            OnnxMl.NodeProto subtract = ONNXOperators.SUB.build(context,
+                    new String[]{powFeatureEmb.getOutput(0), gemmSquaredFeatureSquaredEmb.getOutput(0)},
+                    context.generateUniqueName("squared_prod_subtract_prod_of_squares"));
+            graphBuilder.addNode(subtract);
+            // sum over embedding dimensions
+            OnnxMl.NodeProto sumOverEmbeddings = ONNXOperators.REDUCE_SUM.build(context,
+                    subtract.getOutput(0),
+                    context.generateUniqueName("sum_over_embeddings"),
+                    Collections.singletonMap("axes", new int[]{1}));
+            graphBuilder.addNode(sumOverEmbeddings);
+            // Divide by 2
+            OnnxMl.NodeProto scaledInteraction = ONNXOperators.DIV.build(context,
+                    new String[]{sumOverEmbeddings.getOutput(0), twoConst.getName()},
+                    context.generateUniqueName("scaled_interaction"));
+            graphBuilder.addNode(scaledInteraction);
+            // Store the output name
+            embeddingOutputs[i] = scaledInteraction.getOutput(0);
+        }
+
+        // Make concat
+        OnnxMl.NodeProto concat = ONNXOperators.CONCAT.build(context, embeddingOutputs, context.generateUniqueName("fm_concat"),
+                Collections.singletonMap("axis", 1)
+        );
+        graphBuilder.addNode(concat);
+
+        // Add to gemm
+        OnnxMl.NodeProto addGemmConcat = ONNXOperators.ADD.build(context, new String[]{gemm.getOutput(0), concat.getOutput(0)}, context.generateUniqueName("fm_output"));
+        graphBuilder.addNode(addGemmConcat);
+
+        return addGemmConcat.getOutput(0);
+    }
+
 }
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelModel.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelModel.java
index a88f48d8f..0242c2c13 100644
--- a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelModel.java
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelModel.java
@@ -16,6 +16,7 @@
 
 package org.tribuo.multilabel.sgd.fm;
 
+import ai.onnx.proto.OnnxMl;
 import org.tribuo.Example;
 import org.tribuo.ImmutableFeatureMap;
 import org.tribuo.ImmutableOutputInfo;
@@ -26,10 +27,15 @@
 import org.tribuo.math.la.DenseVector;
 import org.tribuo.math.util.VectorNormalizer;
 import org.tribuo.multilabel.MultiLabel;
+import org.tribuo.onnx.ONNXContext;
+import org.tribuo.onnx.ONNXExportable;
+import org.tribuo.onnx.ONNXShape;
+import org.tribuo.onnx.ONNXUtils;
 import org.tribuo.provenance.ModelProvenance;
 
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -43,7 +49,7 @@
  * 2010 IEEE International Conference on Data Mining
  * </pre>
  */
-public class FMMultiLabelModel extends AbstractFMModel<MultiLabel> {
+public class FMMultiLabelModel extends AbstractFMModel<MultiLabel> implements ONNXExportable {
     private static final long serialVersionUID = 2L;
 
     private final VectorNormalizer normalizer;
@@ -96,4 +102,41 @@ protected String getDimensionName(int index) {
     protected FMMultiLabelModel copy(String newName, ModelProvenance newProvenance) {
         return new FMMultiLabelModel(newName,newProvenance,featureIDMap,outputIDInfo,(FMParameters)modelParameters.copy(),normalizer,generatesProbabilities,threshold);
     }
+
+    @Override
+    public OnnxMl.ModelProto exportONNXModel(String domain, long modelVersion) {
+        ONNXContext context = new ONNXContext();
+
+        // Build graph
+        OnnxMl.GraphProto graph = exportONNXGraph(context);
+
+        return innerExportONNXModel(graph,domain,modelVersion);
+    }
+
+    @Override
+    public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
+        OnnxMl.GraphProto.Builder graphBuilder = OnnxMl.GraphProto.newBuilder();
+        graphBuilder.setName("FMMultiLabelModel");
+
+        // Make inputs and outputs
+        OnnxMl.TypeProto inputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,featureIDMap.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
+        OnnxMl.ValueInfoProto inputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(inputType).setName("input").build();
+        graphBuilder.addInput(inputValueProto);
+        OnnxMl.TypeProto outputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,outputIDInfo.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
+        OnnxMl.ValueInfoProto outputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(outputType).setName("output").build();
+        graphBuilder.addOutput(outputValueProto);
+
+        // Build the output neutral bits of the onnx graph
+        String outputName = generateONNXGraph(context, graphBuilder, inputValueProto.getName());
+
+        // Make output normalizer
+        List<OnnxMl.NodeProto> normalizerProtos = normalizer.exportNormalizer(context,outputName,"output");
+        if (normalizerProtos.isEmpty()) {
+            throw new IllegalArgumentException("Normalizer " + normalizer.getClass() + " cannot be exported in ONNX models.");
+        } else {
+            graphBuilder.addAllNode(normalizerProtos);
+        }
+
+        return graphBuilder.build();
+    }
 }
diff --git a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
index 391beefe0..dc9638c91 100644
--- a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
+++ b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
@@ -16,6 +16,9 @@
 
 package org.tribuo.multilabel.sgd.fm;
 
+import ai.onnxruntime.OrtEnvironment;
+import ai.onnxruntime.OrtException;
+import ai.onnxruntime.OrtSession;
 import com.oracle.labs.mlrg.olcut.util.Pair;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
@@ -24,24 +27,36 @@
 import org.tribuo.Model;
 import org.tribuo.Prediction;
 import org.tribuo.Trainer;
+import org.tribuo.VariableIDInfo;
+import org.tribuo.VariableInfo;
 import org.tribuo.common.sgd.AbstractFMTrainer;
 import org.tribuo.common.sgd.AbstractSGDTrainer;
+import org.tribuo.interop.onnx.DenseTransformer;
+import org.tribuo.interop.onnx.MultiLabelTransformer;
+import org.tribuo.interop.onnx.ONNXExternalModel;
 import org.tribuo.math.optimisers.AdaGrad;
 import org.tribuo.multilabel.MultiLabel;
+import org.tribuo.multilabel.MultiLabelFactory;
 import org.tribuo.multilabel.evaluation.MultiLabelEvaluation;
 import org.tribuo.multilabel.example.MultiLabelDataGenerator;
 import org.tribuo.multilabel.sgd.objectives.BinaryCrossEntropy;
 import org.tribuo.multilabel.sgd.objectives.Hinge;
 import org.tribuo.test.Helpers;
 
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.fail;
 
 public class TestFMMultiLabel {
+    private static final Logger logger = Logger.getLogger(TestFMMultiLabel.class.getName());
 
     private static final FMMultiLabelTrainer hinge = new FMMultiLabelTrainer(new Hinge(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED, 5,0.1);
     private static final FMMultiLabelTrainer sigmoid = new FMMultiLabelTrainer(new BinaryCrossEntropy(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED,5,0.1);
@@ -82,4 +97,60 @@ private static void testTrainer(Dataset<MultiLabel> train, Dataset<MultiLabel> t
         Helpers.testModelSerialization(model, MultiLabel.class);
     }
 
+    @Test
+    public void testOnnxSerialization() throws IOException, OrtException {
+        Dataset<MultiLabel> train = MultiLabelDataGenerator.generateTrainData();
+        Dataset<MultiLabel> test = MultiLabelDataGenerator.generateTestData();
+        FMMultiLabelModel model = (FMMultiLabelModel) sigmoid.train(train);
+
+        // Write out model
+        Path onnxFile = Files.createTempFile("tribuo-fm-test",".onnx");
+        model.saveONNXModel("org.tribuo.multilabel.sgd.fm.test",1,onnxFile);
+
+        // Prep mappings
+        Map<String, Integer> featureMapping = new HashMap<>();
+        for (VariableInfo f : model.getFeatureIDMap()){
+            VariableIDInfo id = (VariableIDInfo) f;
+            featureMapping.put(id.getName(),id.getID());
+        }
+        Map<MultiLabel, Integer> outputMapping = new HashMap<>();
+        for (Pair<Integer,MultiLabel> l : model.getOutputIDInfo()) {
+            outputMapping.put(l.getB(), l.getA());
+        }
+
+        String arch = System.getProperty("os.arch");
+        if (arch.equalsIgnoreCase("amd64") || arch.equalsIgnoreCase("x86_64")) {
+            // Initialise the OrtEnvironment to load the native library
+            // (as OrtSession.SessionOptions doesn't trigger the static initializer).
+            OrtEnvironment env = OrtEnvironment.getEnvironment();
+            env.close();
+            // Load in via ORT
+            ONNXExternalModel<MultiLabel> onnxModel = ONNXExternalModel.createOnnxModel(new MultiLabelFactory(),featureMapping,outputMapping,new DenseTransformer(),new MultiLabelTransformer(),new OrtSession.SessionOptions(),onnxFile,"input");
+
+            // Generate predictions
+            List<Prediction<MultiLabel>> nativePredictions = model.predict(test);
+            List<Prediction<MultiLabel>> onnxPredictions = onnxModel.predict(test);
+
+            // Assert the predictions are identical
+            for (int i = 0; i < nativePredictions.size(); i++) {
+                Prediction<MultiLabel> tribuo = nativePredictions.get(i);
+                Prediction<MultiLabel> external = onnxPredictions.get(i);
+                assertEquals(tribuo.getOutput().getLabelSet(), external.getOutput().getLabelSet());
+                for (Map.Entry<String,MultiLabel> l : tribuo.getOutputScores().entrySet()) {
+                    MultiLabel other = external.getOutputScores().get(l.getKey());
+                    if (other == null) {
+                        fail("Failed to find label " + l.getKey() + " in ORT prediction.");
+                    } else {
+                        assertEquals(l.getValue().getScore(),other.getScore(),1e-5);
+                    }
+                }
+            }
+
+            onnxModel.close();
+        } else {
+            logger.warning("ORT based tests only supported on x86_64, found " + arch);
+        }
+
+        onnxFile.toFile().delete();
+    }
 }
diff --git a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java
index 94a1ab92a..e65568e69 100644
--- a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java
+++ b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java
@@ -96,6 +96,7 @@ private static void testTrainer(Dataset<MultiLabel> train, Dataset<MultiLabel> t
 
         Helpers.testModelSerialization(model, MultiLabel.class);
     }
+
     @Test
     public void testOnnxSerialization() throws IOException, OrtException {
         Dataset<MultiLabel> train = MultiLabelDataGenerator.generateTrainData();
@@ -104,7 +105,7 @@ public void testOnnxSerialization() throws IOException, OrtException {
 
         // Write out model
         Path onnxFile = Files.createTempFile("tribuo-sgd-test",".onnx");
-        model.saveONNXModel("org.tribuo.classification.sgd.linear.test",1,onnxFile);
+        model.saveONNXModel("org.tribuo.multilabel.sgd.linear.test",1,onnxFile);
 
         // Prep mappings
         Map<String, Integer> featureMapping = new HashMap<>();
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
index f0fa55ba1..7665351eb 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
@@ -16,15 +16,20 @@
 
 package org.tribuo.regression.sgd.fm;
 
+import ai.onnx.proto.OnnxMl;
 import org.tribuo.Example;
 import org.tribuo.ImmutableFeatureMap;
 import org.tribuo.ImmutableOutputInfo;
 import org.tribuo.Prediction;
 import org.tribuo.common.sgd.AbstractFMModel;
 import org.tribuo.common.sgd.FMParameters;
+import org.tribuo.onnx.ONNXContext;
+import org.tribuo.onnx.ONNXExportable;
+import org.tribuo.onnx.ONNXOperators;
+import org.tribuo.onnx.ONNXShape;
+import org.tribuo.onnx.ONNXUtils;
 import org.tribuo.provenance.ModelProvenance;
 import org.tribuo.regression.ImmutableRegressionInfo;
-import org.tribuo.regression.RegressionInfo;
 import org.tribuo.regression.Regressor;
 
 import java.util.Arrays;
@@ -41,7 +46,7 @@
  * 2010 IEEE International Conference on Data Mining
  * </pre>
  */
-public class FMRegressionModel extends AbstractFMModel<Regressor> {
+public class FMRegressionModel extends AbstractFMModel<Regressor> implements ONNXExportable {
     private static final long serialVersionUID = 3L;
 
     private final String[] dimensionNames;
@@ -100,4 +105,36 @@ protected String getDimensionName(int index) {
         return dimensionNames[index];
     }
 
+    @Override
+    public OnnxMl.ModelProto exportONNXModel(String domain, long modelVersion) {
+        ONNXContext context = new ONNXContext();
+
+        // Build graph
+        OnnxMl.GraphProto graph = exportONNXGraph(context);
+
+        return innerExportONNXModel(graph,domain,modelVersion);
+    }
+
+    @Override
+    public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
+        OnnxMl.GraphProto.Builder graphBuilder = OnnxMl.GraphProto.newBuilder();
+        graphBuilder.setName("FMMultiLabelModel");
+
+        // Make inputs and outputs
+        OnnxMl.TypeProto inputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,featureIDMap.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
+        OnnxMl.ValueInfoProto inputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(inputType).setName("input").build();
+        graphBuilder.addInput(inputValueProto);
+        OnnxMl.TypeProto outputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,outputIDInfo.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
+        OnnxMl.ValueInfoProto outputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(outputType).setName("output").build();
+        graphBuilder.addOutput(outputValueProto);
+
+        // Build the output neutral bits of the onnx graph
+        String outputName = generateONNXGraph(context, graphBuilder, inputValueProto.getName());
+
+        // Link up the output to the graph output
+        OnnxMl.NodeProto output = ONNXOperators.IDENTITY.build(context,outputName,"output");
+        graphBuilder.addNode(output);
+
+        return graphBuilder.build();
+    }
 }
diff --git a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
index 6fa28e824..33a45b4c5 100644
--- a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
+++ b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
@@ -16,19 +16,26 @@
 
 package org.tribuo.regression.sgd.fm;
 
+import ai.onnxruntime.OrtEnvironment;
+import ai.onnxruntime.OrtException;
+import ai.onnxruntime.OrtSession;
 import com.oracle.labs.mlrg.olcut.util.Pair;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.tribuo.Dataset;
 import org.tribuo.Model;
+import org.tribuo.Prediction;
 import org.tribuo.Trainer;
-import org.tribuo.common.sgd.AbstractFMModel;
+import org.tribuo.VariableIDInfo;
+import org.tribuo.VariableInfo;
 import org.tribuo.common.sgd.AbstractFMTrainer;
 import org.tribuo.common.sgd.AbstractSGDTrainer;
-import org.tribuo.math.la.DenseMatrix;
-import org.tribuo.math.la.DenseVector;
+import org.tribuo.interop.onnx.DenseTransformer;
+import org.tribuo.interop.onnx.ONNXExternalModel;
+import org.tribuo.interop.onnx.RegressorTransformer;
 import org.tribuo.math.optimisers.AdaGrad;
+import org.tribuo.regression.RegressionFactory;
 import org.tribuo.regression.Regressor;
 import org.tribuo.regression.evaluation.RegressionEvaluation;
 import org.tribuo.regression.evaluation.RegressionEvaluator;
@@ -36,15 +43,21 @@
 import org.tribuo.regression.sgd.objectives.SquaredLoss;
 import org.tribuo.test.Helpers;
 
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class TestFMRegression {
+    private static final Logger logger = Logger.getLogger(TestFMRegression.class.getName());
 
     private static final FMRegressionTrainer t = new FMRegressionTrainer(new SquaredLoss(), new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED,1,0.1,true);
 
@@ -119,4 +132,54 @@ public void testMultiEmptyExample() {
             m.predict(RegressionDataGenerator.emptyMultiDimExample());
         });
     }
+
+    @Test
+    public void testOnnxSerialization() throws IOException, OrtException {
+        Pair<Dataset<Regressor>,Dataset<Regressor>> p = RegressionDataGenerator.denseTrainTest();
+        FMRegressionModel model = (FMRegressionModel) t.train(p.getA());
+
+        // Write out model
+        Path onnxFile = Files.createTempFile("tribuo-fm-test",".onnx");
+        model.saveONNXModel("org.tribuo.regression.sgd.fm.test",1,onnxFile);
+
+        // Prep mappings
+        Map<String, Integer> featureMapping = new HashMap<>();
+        for (VariableInfo f : model.getFeatureIDMap()){
+            VariableIDInfo id = (VariableIDInfo) f;
+            featureMapping.put(id.getName(),id.getID());
+        }
+        Map<Regressor, Integer> outputMapping = new HashMap<>();
+        for (Pair<Integer,Regressor> l : model.getOutputIDInfo()) {
+            outputMapping.put(l.getB(), l.getA());
+        }
+
+        String arch = System.getProperty("os.arch");
+        if (arch.equalsIgnoreCase("amd64") || arch.equalsIgnoreCase("x86_64")) {
+            // Initialise the OrtEnvironment to load the native library
+            // (as OrtSession.SessionOptions doesn't trigger the static initializer).
+            OrtEnvironment env = OrtEnvironment.getEnvironment();
+            env.close();
+            // Load in via ORT
+            ONNXExternalModel<Regressor> onnxModel = ONNXExternalModel.createOnnxModel(new RegressionFactory(),featureMapping,outputMapping,new DenseTransformer(),new RegressorTransformer(),new OrtSession.SessionOptions(),onnxFile,"input");
+
+            // Generate predictions
+            List<Prediction<Regressor>> nativePredictions = model.predict(p.getB());
+            List<Prediction<Regressor>> onnxPredictions = onnxModel.predict(p.getB());
+
+            // Assert the predictions are identical
+            for (int i = 0; i < nativePredictions.size(); i++) {
+                Prediction<Regressor> tribuo = nativePredictions.get(i);
+                Prediction<Regressor> external = onnxPredictions.get(i);
+                assertArrayEquals(tribuo.getOutput().getNames(),external.getOutput().getNames());
+                assertArrayEquals(tribuo.getOutput().getValues(),external.getOutput().getValues(),1e-5);
+            }
+
+            onnxModel.close();
+        } else {
+            logger.warning("ORT based tests only supported on x86_64, found " + arch);
+        }
+
+        onnxFile.toFile().delete();
+    }
+
 }
diff --git a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java
index cbce5b3ba..c782c8f66 100644
--- a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java
+++ b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java
@@ -100,7 +100,7 @@ public void testOnnxSerialization() throws IOException, OrtException {
 
         // Write out model
         Path onnxFile = Files.createTempFile("tribuo-sgd-test",".onnx");
-        model.saveONNXModel("org.tribuo.classification.sgd.linear.test",1,onnxFile);
+        model.saveONNXModel("org.tribuo.regression.sgd.linear.test",1,onnxFile);
 
         // Prep mappings
         Map<String, Integer> featureMapping = new HashMap<>();

From e82812826a15894b5a7d85acf5bc6d4becd4c5fe Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 26 Sep 2021 21:46:44 -0400
Subject: [PATCH 24/31] Adding graph names to the existing ONNX exports. The
 graph name is mandatory but ORT doesn't appear to care.

---
 .../org/tribuo/classification/sgd/linear/LinearSGDModel.java     | 1 +
 .../java/org/tribuo/multilabel/sgd/linear/LinearSGDModel.java    | 1 +
 .../java/org/tribuo/regression/sgd/linear/LinearSGDModel.java    | 1 +
 .../main/java/org/tribuo/regression/slm/SparseLinearModel.java   | 1 +
 4 files changed, 4 insertions(+)

diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/linear/LinearSGDModel.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/linear/LinearSGDModel.java
index 2d929928b..344d83a1e 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/linear/LinearSGDModel.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/linear/LinearSGDModel.java
@@ -122,6 +122,7 @@ public OnnxMl.ModelProto exportONNXModel(String domain, long modelVersion) {
     @Override
     public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
         OnnxMl.GraphProto.Builder graphBuilder = OnnxMl.GraphProto.newBuilder();
+        graphBuilder.setName("Classification-LinearSGDModel");
 
         // Make inputs and outputs
         OnnxMl.TypeProto inputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,featureIDMap.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/linear/LinearSGDModel.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/linear/LinearSGDModel.java
index 88193430b..98f28a26f 100644
--- a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/linear/LinearSGDModel.java
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/linear/LinearSGDModel.java
@@ -118,6 +118,7 @@ public OnnxMl.ModelProto exportONNXModel(String domain, long modelVersion) {
     @Override
     public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
         OnnxMl.GraphProto.Builder graphBuilder = OnnxMl.GraphProto.newBuilder();
+        graphBuilder.setName("MultiLabel-LinearSGDModel");
 
         // Make inputs and outputs
         OnnxMl.TypeProto inputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,featureIDMap.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDModel.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDModel.java
index 362f370d8..801a2c0f9 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDModel.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/linear/LinearSGDModel.java
@@ -104,6 +104,7 @@ public OnnxMl.ModelProto exportONNXModel(String domain, long modelVersion) {
     @Override
     public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
         OnnxMl.GraphProto.Builder graphBuilder = OnnxMl.GraphProto.newBuilder();
+        graphBuilder.setName("Regression-LinearSGDModel");
 
         // Make inputs and outputs
         OnnxMl.TypeProto inputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,featureIDMap.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
diff --git a/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java b/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
index 0c8a26295..23718351a 100644
--- a/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
+++ b/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
@@ -295,6 +295,7 @@ protected OnnxMl.TensorProto biasBuilder(ONNXContext context) {
     @Override
     public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
         OnnxMl.GraphProto.Builder graphBuilder = OnnxMl.GraphProto.newBuilder();
+        graphBuilder.setName("Regression-SparseLinearModel");
 
         // Make inputs and outputs
         OnnxMl.TypeProto inputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1, featureIDMap.size()}, new String[]{"batch", null}), OnnxMl.TensorProto.DataType.FLOAT);

From a159bb938ccb451174430bc97c140bc0aa1e2033 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 26 Sep 2021 21:54:53 -0400
Subject: [PATCH 25/31] Tidying up the bias tensor proto generation.

---
 .../tribuo/common/sgd/AbstractFMModel.java    | 31 ++++++++++---------
 1 file changed, 16 insertions(+), 15 deletions(-)

diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
index 0c228dcf5..138ff01ea 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
@@ -205,7 +205,7 @@ protected OnnxMl.ModelProto innerExportONNXModel(OnnxMl.GraphProto graph, String
      * @param transpose Should the matrix be transposed into the tensor?
      * @return The matrix TensorProto.
      */
-    protected OnnxMl.TensorProto matrixBuilder(ONNXContext context, String name, DenseMatrix matrix, boolean transpose) {
+    protected static OnnxMl.TensorProto matrixBuilder(ONNXContext context, String name, DenseMatrix matrix, boolean transpose) {
         OnnxMl.TensorProto.Builder matrixBuilder = OnnxMl.TensorProto.newBuilder();
         matrixBuilder.setName(context.generateUniqueName(name));
         int dim1, dim2;
@@ -236,25 +236,26 @@ protected OnnxMl.TensorProto matrixBuilder(ONNXContext context, String name, Den
     }
 
     /**
-     * Builds a TensorProto containing the biases for this Factorization Machine.
+     * Builds a TensorProto containing the supplied dense vector.
      *
      * @param context The ONNX context for naming.
-     * @return The bias TensorProto.
+     * @param name    The name for this tensor proto.
+     * @param vector  The vector to store.
+     * @return The vector TensorProto.
      */
-    protected OnnxMl.TensorProto biasBuilder(ONNXContext context) {
-        OnnxMl.TensorProto.Builder biasBuilder = OnnxMl.TensorProto.newBuilder();
-        biasBuilder.setName(context.generateUniqueName("fm_biases"));
-        biasBuilder.addDims(outputIDInfo.size());
-        biasBuilder.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
-        ByteBuffer buffer = ByteBuffer.allocate(outputIDInfo.size() * 4).order(ByteOrder.LITTLE_ENDIAN);
+    protected static OnnxMl.TensorProto vectorBuilder(ONNXContext context, String name, DenseVector vector) {
+        OnnxMl.TensorProto.Builder vectorBuilder = OnnxMl.TensorProto.newBuilder();
+        vectorBuilder.setName(context.generateUniqueName(name));
+        vectorBuilder.addDims(vector.size());
+        vectorBuilder.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
+        ByteBuffer buffer = ByteBuffer.allocate(vector.size() * 4).order(ByteOrder.LITTLE_ENDIAN);
         FloatBuffer floatBuffer = buffer.asFloatBuffer();
-        DenseVector biases = (DenseVector) modelParameters.get()[0];
-        for (int i = 0; i < biases.size(); i++) {
-            floatBuffer.put((float) biases.get(i));
+        for (int i = 0; i < vector.size(); i++) {
+            floatBuffer.put((float) vector.get(i));
         }
         floatBuffer.rewind();
-        biasBuilder.setRawData(ByteString.copyFrom(buffer));
-        return biasBuilder.build();
+        vectorBuilder.setRawData(ByteString.copyFrom(buffer));
+        return vectorBuilder.build();
     }
 
     /**
@@ -282,7 +283,7 @@ protected String generateONNXGraph(ONNXContext context, OnnxMl.GraphProto.Builde
         graphBuilder.addInitializer(weightInitializerProto);
 
         // Add biases
-        OnnxMl.TensorProto biasInitializerProto = biasBuilder(context);
+        OnnxMl.TensorProto biasInitializerProto = vectorBuilder(context, "fm_biases", (DenseVector) modelParams[0]);
         graphBuilder.addInitializer(biasInitializerProto);
 
         // Add embedding vectors

From 4ea0252e2ae337d9cef4d60260104919d4ce3585 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 26 Sep 2021 22:15:24 -0400
Subject: [PATCH 26/31] Adding standardisation to ONNX export from
 FMRegressionModel.

---
 .../regression/sgd/fm/FMRegressionModel.java  | 35 ++++++++++++++++---
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
index 7665351eb..973be0250 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/FMRegressionModel.java
@@ -124,16 +124,41 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
         OnnxMl.TypeProto inputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,featureIDMap.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
         OnnxMl.ValueInfoProto inputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(inputType).setName("input").build();
         graphBuilder.addInput(inputValueProto);
+        String outputName = "output";
         OnnxMl.TypeProto outputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,outputIDInfo.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
-        OnnxMl.ValueInfoProto outputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(outputType).setName("output").build();
+        OnnxMl.ValueInfoProto outputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(outputType).setName(outputName).build();
         graphBuilder.addOutput(outputValueProto);
 
         // Build the output neutral bits of the onnx graph
-        String outputName = generateONNXGraph(context, graphBuilder, inputValueProto.getName());
+        String fmOutputName = generateONNXGraph(context, graphBuilder, inputValueProto.getName());
 
-        // Link up the output to the graph output
-        OnnxMl.NodeProto output = ONNXOperators.IDENTITY.build(context,outputName,"output");
-        graphBuilder.addNode(output);
+        if (standardise) {
+            // standardise the FM output
+            ImmutableRegressionInfo info = (ImmutableRegressionInfo) outputIDInfo;
+            double[] means = new double[outputIDInfo.size()];
+            double[] variances = new double[outputIDInfo.size()];
+            for (int i = 0; i < means.length; i++) {
+                means[i] = info.getMean(i);
+                variances[i] = info.getVariance(i);
+            }
+
+            // Create mean and variance initializers
+            OnnxMl.TensorProto outputMeanProto = ONNXUtils.arrayBuilder(context,context.generateUniqueName("y_mean"),means);
+            graphBuilder.addInitializer(outputMeanProto);
+            OnnxMl.TensorProto outputVarianceProto = ONNXUtils.arrayBuilder(context, context.generateUniqueName("y_var"),variances);
+            graphBuilder.addInitializer(outputVarianceProto);
+
+            // Add standardisation operations
+            String varianceOutput = context.generateUniqueName("y_var_scale_output");
+            OnnxMl.NodeProto varianceScale = ONNXOperators.MUL.build(context, new String[]{fmOutputName,outputVarianceProto.getName()}, varianceOutput);
+            graphBuilder.addNode(varianceScale);
+            OnnxMl.NodeProto meanScale = ONNXOperators.ADD.build(context, new String[]{varianceOutput,outputMeanProto.getName()}, outputName);
+            graphBuilder.addNode(meanScale);
+        } else {
+            // Not standardised, so link up the FM output to the graph output
+            OnnxMl.NodeProto output = ONNXOperators.IDENTITY.build(context, fmOutputName, outputName);
+            graphBuilder.addNode(output);
+        }
 
         return graphBuilder.build();
     }

From e4d1c2830cd29c5bed11bef2bf9479bb7475778e Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Sun, 26 Sep 2021 22:21:20 -0400
Subject: [PATCH 27/31] Relaxing the similarity test in FMRegressionModel as
 the standardisation is quite different in floats.

---
 .../java/org/tribuo/regression/sgd/fm/TestFMRegression.java     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
index 33a45b4c5..09e8e8fd1 100644
--- a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
+++ b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
@@ -171,7 +171,7 @@ public void testOnnxSerialization() throws IOException, OrtException {
                 Prediction<Regressor> tribuo = nativePredictions.get(i);
                 Prediction<Regressor> external = onnxPredictions.get(i);
                 assertArrayEquals(tribuo.getOutput().getNames(),external.getOutput().getNames());
-                assertArrayEquals(tribuo.getOutput().getValues(),external.getOutput().getValues(),1e-5);
+                assertArrayEquals(tribuo.getOutput().getValues(),external.getOutput().getValues(),1e-4);
             }
 
             onnxModel.close();

From 9b500afa2d234f2bb79d2f5c0b7a475ded69e18a Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Thu, 7 Oct 2021 09:20:37 -0400
Subject: [PATCH 28/31] Fixing a bug in AbstractFMTrainer when instantiated
 from configuration.

---
 .../sgd/fm/TestFMClassification.java          | 23 +++++++++++++++++--
 .../tribuo/common/sgd/AbstractFMTrainer.java  |  2 +-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
index 4c70429ff..2d0b44dde 100644
--- a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
+++ b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
@@ -19,6 +19,9 @@
 import ai.onnxruntime.OrtEnvironment;
 import ai.onnxruntime.OrtException;
 import ai.onnxruntime.OrtSession;
+import com.oracle.labs.mlrg.olcut.config.ConfigurationData;
+import com.oracle.labs.mlrg.olcut.config.ConfigurationManager;
+import com.oracle.labs.mlrg.olcut.provenance.ProvenanceUtil;
 import com.oracle.labs.mlrg.olcut.util.Pair;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
@@ -55,6 +58,7 @@
 import java.util.logging.Logger;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -78,10 +82,10 @@ public static Model<Label> testFMClassification(Pair<Dataset<Label>,Dataset<Labe
         LabelEvaluator e = new LabelEvaluator();
         LabelEvaluation evaluation = e.evaluate(m,p.getB());
         Map<String, List<Pair<String,Double>>> features = m.getTopFeatures(3);
-        Assertions.assertNotNull(features);
+        assertNotNull(features);
         Assertions.assertFalse(features.isEmpty());
         features = m.getTopFeatures(-1);
-        Assertions.assertNotNull(features);
+        assertNotNull(features);
         Assertions.assertFalse(features.isEmpty());
         return m;
     }
@@ -180,4 +184,19 @@ public void testOnnxSerialization() throws IOException, OrtException {
         onnxFile.toFile().delete();
     }
 
+    @Test
+    public void fmRoundTripTest() {
+        Pair<Dataset<Label>,Dataset<Label>> p = LabelledDataGenerator.denseTrainTest();
+        Model<Label> model = testFMClassification(p);
+
+        ConfigurationManager cm = new ConfigurationManager();
+        List<ConfigurationData> provConfig = ProvenanceUtil.extractConfiguration(model.getProvenance());
+        cm.addConfiguration(provConfig);
+
+        FMClassificationTrainer trainer = (FMClassificationTrainer) cm.lookup("fmclassificationtrainer-0");
+        Model<Label> newModel = trainer.train(p.getA());
+
+        assertNotNull(newModel);
+    }
+
 }
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
index 07f66afce..7fcf94834 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMTrainer.java
@@ -68,7 +68,7 @@ protected AbstractFMTrainer(StochasticGradientOptimiser optimiser, int epochs, i
      * For olcut.
      */
     protected AbstractFMTrainer() {
-        super(true);
+        super(false);
     }
 
     @Override

From db5dd86cef2124dc51be02c353cbc3bbf5eb7086 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Fri, 8 Oct 2021 18:30:09 -0400
Subject: [PATCH 29/31] Adding example configs for factorization machines

---
 .../configs/mnb-example.xml                   |  2 +-
 .../classification-factorization-machine.xml  | 43 +++++++++++++++++++
 .../multilabel-factorization-machine.xml      | 43 +++++++++++++++++++
 .../regression-factorization-machine.xml      | 43 +++++++++++++++++++
 4 files changed, 130 insertions(+), 1 deletion(-)
 create mode 100644 Classification/SGD/configs/classification-factorization-machine.xml
 create mode 100644 MultiLabel/SGD/configs/multilabel-factorization-machine.xml
 create mode 100644 Regression/SGD/configs/regression-factorization-machine.xml

diff --git a/Classification/MultinomialNaiveBayes/configs/mnb-example.xml b/Classification/MultinomialNaiveBayes/configs/mnb-example.xml
index fed33c3d8..94e7a63f2 100644
--- a/Classification/MultinomialNaiveBayes/configs/mnb-example.xml
+++ b/Classification/MultinomialNaiveBayes/configs/mnb-example.xml
@@ -22,7 +22,7 @@
 -->
 
 <config>
-    <component name="cart" type="org.tribuo.classification.mnb.MultinomialNaiveBayesTrainer">
+    <component name="mnb" type="org.tribuo.classification.mnb.MultinomialNaiveBayesTrainer">
         <property name="alpha" value="1.0"/>
     </component>
 </config>
diff --git a/Classification/SGD/configs/classification-factorization-machine.xml b/Classification/SGD/configs/classification-factorization-machine.xml
new file mode 100644
index 000000000..a396a8cc7
--- /dev/null
+++ b/Classification/SGD/configs/classification-factorization-machine.xml
@@ -0,0 +1,43 @@
+<?xml version="1.0" encoding="UTF-8"?>
+
+<!--
+  ~ Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~     http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<!--
+    Description:
+        An example configuration file for a factorization machine trained using AdaGrad.
+-->
+
+<config>
+    <component name="fm" type="org.tribuo.classification.sgd.fm.LinearSGDTrainer">
+        <property name="objective" value="log"/>
+        <property name="optimiser" value="adagrad"/>
+        <property name="epochs" value="10"/>
+        <property name="loggingInterval" value="1000"/>
+        <property name="minibatchSize" value="1"/>
+        <property name="seed" value="1"/>
+        <property name="standardise" value="true"/>
+        <property name="factoredDimSize" value="5"/>
+        <property name="variance" value="0.5"/>
+    </component>
+
+    <component name="log" type="org.tribuo.classification.sgd.objectives.LogMulticlass"/>
+
+    <component name="adagrad" type="org.tribuo.math.optimisers.AdaGrad">
+        <property name="initialLearningRate" value="1.0"/>
+        <property name="epsilon" value="0.1"/>
+    </component>
+</config>
diff --git a/MultiLabel/SGD/configs/multilabel-factorization-machine.xml b/MultiLabel/SGD/configs/multilabel-factorization-machine.xml
new file mode 100644
index 000000000..58c820194
--- /dev/null
+++ b/MultiLabel/SGD/configs/multilabel-factorization-machine.xml
@@ -0,0 +1,43 @@
+<?xml version="1.0" encoding="UTF-8"?>
+
+<!--
+  ~ Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~     http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<!--
+    Description:
+        An example configuration file for a factorization machine trained using AdaGrad.
+-->
+
+<config>
+    <component name="fm" type="org.tribuo.multilabel.sgd.fm.LinearSGDTrainer">
+        <property name="objective" value="bce"/>
+        <property name="optimiser" value="adagrad"/>
+        <property name="epochs" value="10"/>
+        <property name="loggingInterval" value="1000"/>
+        <property name="minibatchSize" value="1"/>
+        <property name="seed" value="1"/>
+        <property name="standardise" value="true"/>
+        <property name="factoredDimSize" value="5"/>
+        <property name="variance" value="0.5"/>
+    </component>
+
+    <component name="bce" type="org.tribuo.multilabel.sgd.objectives.BinaryCrossEntropy"/>
+
+    <component name="adagrad" type="org.tribuo.math.optimisers.AdaGrad">
+        <property name="initialLearningRate" value="1.0"/>
+        <property name="epsilon" value="0.1"/>
+    </component>
+</config>
diff --git a/Regression/SGD/configs/regression-factorization-machine.xml b/Regression/SGD/configs/regression-factorization-machine.xml
new file mode 100644
index 000000000..99768948e
--- /dev/null
+++ b/Regression/SGD/configs/regression-factorization-machine.xml
@@ -0,0 +1,43 @@
+<?xml version="1.0" encoding="UTF-8"?>
+
+<!--
+  ~ Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~     http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<!--
+    Description:
+        An example configuration file for a factorization machine trained using AdaGrad.
+-->
+
+<config>
+    <component name="fm" type="org.tribuo.regression.sgd.fm.LinearSGDTrainer">
+        <property name="objective" value="squared"/>
+        <property name="optimiser" value="adagrad"/>
+        <property name="epochs" value="10"/>
+        <property name="loggingInterval" value="1000"/>
+        <property name="minibatchSize" value="1"/>
+        <property name="seed" value="1"/>
+        <property name="standardise" value="true"/>
+        <property name="factoredDimSize" value="5"/>
+        <property name="variance" value="0.5"/>
+    </component>
+
+    <component name="squared" type="org.tribuo.regression.sgd.objectives.SquaredLoss"/>
+
+    <component name="adagrad" type="org.tribuo.math.optimisers.AdaGrad">
+        <property name="initialLearningRate" value="1.0"/>
+        <property name="epsilon" value="0.1"/>
+    </component>
+</config>

From d5859a32ecb85ee6c2c62c2e6e8c8e7e5c3e8317 Mon Sep 17 00:00:00 2001
From: Adam Pocock <craigacp@gmail.com>
Date: Tue, 12 Oct 2021 13:54:05 -0400
Subject: [PATCH 30/31] Apply Jack's suggestions from code review

Co-authored-by: Jack Sullivan <john.t.sullivan@gmail.com>
---
 .../SGD/configs/classification-factorization-machine.xml  | 3 +--
 .../main/java/org/tribuo/common/sgd/AbstractFMModel.java  | 2 +-
 .../java/org/tribuo/common/sgd/AbstractSGDTrainer.java    | 4 ++--
 .../src/main/java/org/tribuo/common/sgd/FMParameters.java | 8 ++++----
 Math/src/main/java/org/tribuo/math/la/SparseVector.java   | 4 +---
 .../SGD/configs/multilabel-factorization-machine.xml      | 3 +--
 .../SGD/configs/regression-factorization-machine.xml      | 2 +-
 7 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/Classification/SGD/configs/classification-factorization-machine.xml b/Classification/SGD/configs/classification-factorization-machine.xml
index a396a8cc7..02ba89bab 100644
--- a/Classification/SGD/configs/classification-factorization-machine.xml
+++ b/Classification/SGD/configs/classification-factorization-machine.xml
@@ -22,14 +22,13 @@
 -->
 
 <config>
-    <component name="fm" type="org.tribuo.classification.sgd.fm.LinearSGDTrainer">
+    <component name="fm" type="org.tribuo.classification.sgd.fm.FMClassificationTrainer">
         <property name="objective" value="log"/>
         <property name="optimiser" value="adagrad"/>
         <property name="epochs" value="10"/>
         <property name="loggingInterval" value="1000"/>
         <property name="minibatchSize" value="1"/>
         <property name="seed" value="1"/>
-        <property name="standardise" value="true"/>
         <property name="factoredDimSize" value="5"/>
         <property name="variance" value="0.5"/>
     </component>
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
index 138ff01ea..f888ab8a0 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
@@ -159,7 +159,7 @@ public Tensor[] getFactorsCopy() {
      * Factorization machines don't provide excuses, use an explainer.
      *
      * @param example The input example.
-     * @return Optional.emtpy.
+     * @return Optional.empty.
      */
     @Override
     public Optional<Excuse<T>> getExcuse(Example<T> example) {
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java
index 9f45905ba..390baae4c 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractSGDTrainer.java
@@ -176,8 +176,8 @@ public V train(Dataset<T> examples, Map<String, Provenance> runProvenance) {
             n++;
         }
         logger.info(String.format("Training SGD model with %d examples", n));
-        logger.fine("Average feature size = " + featureSize / (double)n);
-        logger.fine("Dense count = " + denseCount);
+        logger.fine("Mean number of active features = " + featureSize / (double)n);
+        logger.fine("Number of dense examples = " + denseCount);
         logger.info("Outputs - " + outputIDInfo.toReadableString());
 
         X parameters = createParameters(featureIDMap.size(), outputIDInfo.size(), localRNG);
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
index 02134a51c..869db9058 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/FMParameters.java
@@ -150,7 +150,7 @@ public DenseVector predict(SGDVector example) {
      * Generate the gradients for a particular feature vector given
      * the loss and the per output gradients.
      * <p>
-     * This parameters returns a {@link Tensor} array with numLabels + 2 elements.
+     * This method returns a {@link Tensor} array with numLabels + 2 elements.
      *
      * @param score    The Pair returned by the objective.
      * @param features The feature vector.
@@ -266,12 +266,12 @@ public Tensor[] merge(Tensor[][] gradients, int size) {
                     gradients[0][i].intersectAndAddInPlace(gradients[j][i]);
                 }
                 output[i] = gradients[0][i];
-            } else if (gradients[0][0] instanceof DenseMatrix) {
+            } else if (gradients[0][i] instanceof DenseMatrix) {
                 for (int j = 1; j < size; j++) {
                     gradients[0][i].intersectAndAddInPlace(gradients[j][i]);
                 }
                 output[i] = gradients[0][i];
-            } else if (gradients[0][0] instanceof DenseSparseMatrix) {
+            } else if (gradients[0][i] instanceof DenseSparseMatrix) {
                 DenseSparseMatrix[] updates = new DenseSparseMatrix[size];
                 for (int j = 0; j < updates.length; j++) {
                     updates[j] = (DenseSparseMatrix) gradients[j][0];
@@ -281,7 +281,7 @@ public Tensor[] merge(Tensor[][] gradients, int size) {
 
                 output[i] = update;
             } else {
-                throw new IllegalStateException("Unexpected gradient type, expected DenseMatrix or DenseSparseMatrix, received " + gradients[0][0].getClass().getName());
+                throw new IllegalStateException("Unexpected gradient type, expected DenseVector, DenseMatrix or DenseSparseMatrix, received " + gradients[0][i].getClass().getName());
             }
         }
         return output;
diff --git a/Math/src/main/java/org/tribuo/math/la/SparseVector.java b/Math/src/main/java/org/tribuo/math/la/SparseVector.java
index acbf408a0..e6a0ca85f 100644
--- a/Math/src/main/java/org/tribuo/math/la/SparseVector.java
+++ b/Math/src/main/java/org/tribuo/math/la/SparseVector.java
@@ -472,9 +472,7 @@ public void hadamardProductInPlace(Tensor other, DoubleUnaryOperator f) {
     }
 
     /**
-     * Applies a {@link ToDoubleBiFunction} elementwise to this {@link SGDVector}.
-     * <p>
-     * The first argument to the function is the index, the second argument is the current value.
+     * Applies a {@link DoubleUnaryOperator} elementwise to this {@link SGDVector}.
      * <p>
      * Only applies the function to the elements which are present.
      * <p>
diff --git a/MultiLabel/SGD/configs/multilabel-factorization-machine.xml b/MultiLabel/SGD/configs/multilabel-factorization-machine.xml
index 58c820194..2d5649476 100644
--- a/MultiLabel/SGD/configs/multilabel-factorization-machine.xml
+++ b/MultiLabel/SGD/configs/multilabel-factorization-machine.xml
@@ -22,14 +22,13 @@
 -->
 
 <config>
-    <component name="fm" type="org.tribuo.multilabel.sgd.fm.LinearSGDTrainer">
+    <component name="fm" type="org.tribuo.multilabel.sgd.fm.FMMultiLabelTrainer">
         <property name="objective" value="bce"/>
         <property name="optimiser" value="adagrad"/>
         <property name="epochs" value="10"/>
         <property name="loggingInterval" value="1000"/>
         <property name="minibatchSize" value="1"/>
         <property name="seed" value="1"/>
-        <property name="standardise" value="true"/>
         <property name="factoredDimSize" value="5"/>
         <property name="variance" value="0.5"/>
     </component>
diff --git a/Regression/SGD/configs/regression-factorization-machine.xml b/Regression/SGD/configs/regression-factorization-machine.xml
index 99768948e..f455f5b7f 100644
--- a/Regression/SGD/configs/regression-factorization-machine.xml
+++ b/Regression/SGD/configs/regression-factorization-machine.xml
@@ -22,7 +22,7 @@
 -->
 
 <config>
-    <component name="fm" type="org.tribuo.regression.sgd.fm.LinearSGDTrainer">
+    <component name="fm" type="org.tribuo.regression.sgd.fm.FMRegressionTrainer">
         <property name="objective" value="squared"/>
         <property name="optimiser" value="adagrad"/>
         <property name="epochs" value="10"/>

From 2185b4ac5230bae9d85640ceadd93970110a2095 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 12 Oct 2021 17:24:17 -0400
Subject: [PATCH 31/31] Fixing PR comments.

---
 .../configs/classification-factorization-machine.xml  |  4 ++--
 .../sgd/fm/FMClassificationOptions.java               |  4 ++--
 .../java/org/tribuo/common/sgd/AbstractFMModel.java   | 11 +++++++++++
 Math/src/main/java/org/tribuo/math/Parameters.java    |  3 +++
 .../SGD/configs/multilabel-factorization-machine.xml  |  4 ++--
 .../tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java |  4 ++--
 .../SGD/configs/regression-factorization-machine.xml  |  4 ++--
 .../java/org/tribuo/regression/sgd/fm/TrainTest.java  |  4 ++--
 8 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/Classification/SGD/configs/classification-factorization-machine.xml b/Classification/SGD/configs/classification-factorization-machine.xml
index 02ba89bab..ee4862bc5 100644
--- a/Classification/SGD/configs/classification-factorization-machine.xml
+++ b/Classification/SGD/configs/classification-factorization-machine.xml
@@ -25,8 +25,8 @@
     <component name="fm" type="org.tribuo.classification.sgd.fm.FMClassificationTrainer">
         <property name="objective" value="log"/>
         <property name="optimiser" value="adagrad"/>
-        <property name="epochs" value="10"/>
-        <property name="loggingInterval" value="1000"/>
+        <property name="epochs" value="5"/>
+        <property name="loggingInterval" value="100"/>
         <property name="minibatchSize" value="1"/>
         <property name="seed" value="1"/>
         <property name="factoredDimSize" value="5"/>
diff --git a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
index 5fad87098..8874698b9 100644
--- a/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
+++ b/Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationOptions.java
@@ -78,12 +78,12 @@ public enum LossEnum {
      * Factor size.
      */
     @Option(longName = "fm-factor-size", usage = "Factor size.")
-    public int fmFactorSize = 6;
+    public int fmFactorSize = 5;
     /**
      * Variance of the initialization gaussian.
      */
     @Option(longName = "fm-variance", usage = "Variance of the initialization gaussian.")
-    public double fmVariance = 0.1;
+    public double fmVariance = 0.5;
 
     /**
      * Returns the loss function specified in the arguments.
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
index f888ab8a0..df4396422 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
@@ -75,6 +75,17 @@ protected AbstractFMModel(String name, ModelProvenance provenance,
         super(name, provenance, featureIDMap, outputIDInfo, parameters, generatesProbabilities, false);
     }
 
+    /**
+     * Gets the top {@code n} features for each output dimension.
+     * <p>
+     * Note that the feature rankings are based only off the linear portion of the
+     * factorization machine.
+     * @param n The number of features to return. If this value is less than 0,
+     * all features are returned for each class.
+     * @return A map from string outputs to an ordered list of pairs of
+     * feature names and weights associated with that feature in the factorization machine.
+     */
+    // TODO investigate using the factorized representation magnitude as an additional feature weight.
     @Override
     public Map<String, List<Pair<String, Double>>> getTopFeatures(int n) {
         DenseVector biases = (DenseVector) modelParameters.get()[0];
diff --git a/Math/src/main/java/org/tribuo/math/Parameters.java b/Math/src/main/java/org/tribuo/math/Parameters.java
index 877918fd4..7a19e6cb8 100644
--- a/Math/src/main/java/org/tribuo/math/Parameters.java
+++ b/Math/src/main/java/org/tribuo/math/Parameters.java
@@ -65,6 +65,9 @@ public interface Parameters extends Serializable {
      * Merge together an array of gradient arrays. Assumes the first dimension
      * is the number of gradient arrays and the second dimension is the
      * number of parameter {@link Tensor}s.
+     * <p>
+     * For performance reasons this call may mutate the input gradient array, and
+     * may return a subset of those elements as the merge output.
      * @param gradients An array of gradient update arrays.
      * @param size The number of elements of gradients to merge. Allows gradients to have unused elements.
      * @return A single {@link Tensor} array of the summed gradients.
diff --git a/MultiLabel/SGD/configs/multilabel-factorization-machine.xml b/MultiLabel/SGD/configs/multilabel-factorization-machine.xml
index 2d5649476..78ef8bf6c 100644
--- a/MultiLabel/SGD/configs/multilabel-factorization-machine.xml
+++ b/MultiLabel/SGD/configs/multilabel-factorization-machine.xml
@@ -25,8 +25,8 @@
     <component name="fm" type="org.tribuo.multilabel.sgd.fm.FMMultiLabelTrainer">
         <property name="objective" value="bce"/>
         <property name="optimiser" value="adagrad"/>
-        <property name="epochs" value="10"/>
-        <property name="loggingInterval" value="1000"/>
+        <property name="epochs" value="5"/>
+        <property name="loggingInterval" value="100"/>
         <property name="minibatchSize" value="1"/>
         <property name="seed" value="1"/>
         <property name="factoredDimSize" value="5"/>
diff --git a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
index 224aeec7f..f3948aea8 100644
--- a/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
+++ b/MultiLabel/SGD/src/main/java/org/tribuo/multilabel/sgd/fm/FMMultiLabelOptions.java
@@ -81,12 +81,12 @@ public enum LossEnum {
      * Factor size.
      */
     @Option(longName = "fm-factor-size", usage = "Factor size.")
-    public int fmFactorSize = 6;
+    public int fmFactorSize = 5;
     /**
      * Variance of the initialization gaussian.
      */
     @Option(longName = "fm-variance", usage = "Variance of the initialization gaussian.")
-    public double fmVariance = 0.1;
+    public double fmVariance = 0.5;
 
     /**
      * Returns the loss function specified in the arguments.
diff --git a/Regression/SGD/configs/regression-factorization-machine.xml b/Regression/SGD/configs/regression-factorization-machine.xml
index f455f5b7f..8a319d06c 100644
--- a/Regression/SGD/configs/regression-factorization-machine.xml
+++ b/Regression/SGD/configs/regression-factorization-machine.xml
@@ -25,8 +25,8 @@
     <component name="fm" type="org.tribuo.regression.sgd.fm.FMRegressionTrainer">
         <property name="objective" value="squared"/>
         <property name="optimiser" value="adagrad"/>
-        <property name="epochs" value="10"/>
-        <property name="loggingInterval" value="1000"/>
+        <property name="epochs" value="5"/>
+        <property name="loggingInterval" value="100"/>
         <property name="minibatchSize" value="1"/>
         <property name="seed" value="1"/>
         <property name="standardise" value="true"/>
diff --git a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
index 7f407e594..c99d99a99 100644
--- a/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
+++ b/Regression/SGD/src/main/java/org/tribuo/regression/sgd/fm/TrainTest.java
@@ -106,12 +106,12 @@ public String getOptionsDescription() {
          * Factor size.
          */
         @Option(charName = 'd', longName = "factor-size", usage = "Factor size.")
-        public int factorSize = 6;
+        public int factorSize = 5;
         /**
          * Variance of the initialization gaussian.
          */
         @Option(longName = "variance", usage = "Variance of the initialization gaussian.")
-        public double variance = 0.1;
+        public double variance = 0.5;
         /**
          * Standardise the output regressors before model fitting.
          */