[tmva][sofie] Fix LogSoftmax returning -inf for inputs far apart from the maximum

tmva/sofie/inc/TMVA/ROperator_Softmax.hxx

@@ -104,12 +104,21 @@ public:
          out << SP << SP << SP << "sum += y_ptr[j];\n";
          out << SP << SP << "}\n";
 
-         out << SP << SP << fType << " inv_sum = 1.0f / sum;\n";
-         out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n";
-         out << SP << SP << SP << "y_ptr[j] *= inv_sum;\n";
-         if (fLogSoftmax)
-            out << SP << SP << SP << "y_ptr[j] = " << logFunction << "(y_ptr[j]);\n";
-         out << SP << SP << "}\n";
+         if (fLogSoftmax) {
+            // evaluate (x - vmax) - log(sum(exp(x - vmax))) instead of
+            // log(exp(x - vmax) / sum): for inputs far apart from the maximum
+            // exp(x - vmax) underflows to 0 and its log is -inf, while the
+            // correct result is a large negative number (see #23547)
+            out << SP << SP << fType << " log_sum = " << logFunction << "(sum);\n";
+            out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n";
+            out << SP << SP << SP << "y_ptr[j] = (x_ptr[j] - vmax) - log_sum;\n";
+            out << SP << SP << "}\n";
+         } else {
+            out << SP << SP << fType << " inv_sum = 1.0f / sum;\n";
+            out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n";
+            out << SP << SP << SP << "y_ptr[j] *= inv_sum;\n";
+            out << SP << SP << "}\n";
+         }
          out << SP << "}\n";
 
       } else {
@@ -166,20 +175,44 @@ public:
          for (size_t j = 0; j < size-1; j++) out << SP;
          out << "}\n";
          // normalize
-         for (size_t j = 0; j < size-1; j++) out << SP;
-         out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n";
-         for (size_t j = 0; j < size; j++) out << SP;
-         out << "size_t id = index + i";
-         if (stride[axis].GetVal() != "1") out << "*(" << stride[axis] << ")";
-         out << ";\n";
-         for (size_t j = 0; j < size; j++) out << SP;
-         out << "tensor_" << fNY << "[id] /= sum;\n";
          if (fLogSoftmax) {
-            for (size_t j = 0; j < size; j++) out << SP;
-            out << "tensor_" << fNY << "[id] = " << logFunction << "(tensor_" << fNY << "[id]);\n";
+            // same numerically stable form as in the last-axis case above:
+            // (x - vmax) - log(sum(exp(x - vmax))), see #23547
+            for (size_t j = 0; j < size - 1; j++)
+               out << SP;
+            out << fType << " log_sum = " << logFunction << "(sum);\n";
+            for (size_t j = 0; j < size - 1; j++)
+               out << SP;
+            out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n";
+            for (size_t j = 0; j < size; j++)
+               out << SP;
+            out << "size_t id = index + i";
+            if (stride[axis].GetVal() != "1")
+               out << "*(" << stride[axis] << ")";
+            out << ";\n";
+            for (size_t j = 0; j < size; j++)
+               out << SP;
+            out << "tensor_" << fNY << "[id] = (tensor_" << fNX << "[id] - vmax) - log_sum;\n";
+            for (size_t j = 0; j < size - 1; j++)
+               out << SP;
+            out << "}\n";
+         } else {
+            for (size_t j = 0; j < size - 1; j++)
+               out << SP;
+            out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n";
+            for (size_t j = 0; j < size; j++)
+               out << SP;
+            out << "size_t id = index + i";
+            if (stride[axis].GetVal() != "1")
+               out << "*(" << stride[axis] << ")";
+            out << ";\n";
+            for (size_t j = 0; j < size; j++)
+               out << SP;
+            out << "tensor_" << fNY << "[id] /= sum;\n";
+            for (size_t j = 0; j < size - 1; j++)
+               out << SP;
+            out << "}\n";
          }
-         for (size_t j = 0; j < size-1; j++) out << SP;
-         out << "}\n";
          //end loops
          for (int i = static_cast<int>(k) - 1; i >= 0; i--) {
             for (int j = 0; j < i; j++) out << SP;

tmva/sofie/test/TestCustomModelsFromONNX.cxx

@@ -1024,6 +1024,56 @@ TEST(ONNX, Softmax4d)
    expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
 }
 
+TEST(ONNX, LogSoftmax1d)
+{
+   SofieReference ref = readReference("LogSoftmax1d");
+
+   ASSERT_INCLUDE_AND_RUN(std::vector<float>, "LogSoftmax1d", ref.f32("input0"));
+
+   expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
+}
+
+TEST(ONNX, LogSoftmax2d)
+{
+   SofieReference ref = readReference("LogSoftmax2d");
+
+   ASSERT_INCLUDE_AND_RUN(std::vector<float>, "LogSoftmax2d", ref.f32("input0"));
+
+   expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
+}
+
+TEST(ONNX, LogSoftmax3d)
+{
+   SofieReference ref = readReference("LogSoftmax3d");
+
+   ASSERT_INCLUDE_AND_RUN(std::vector<float>, "LogSoftmax3d", ref.f32("input0"));
+
+   expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
+}
+
+// Inputs more than ~88 apart from the maximum make exp(x - max) underflow to 0
+// in float32, so log(softmax(x)) returns -inf where the stable form
+// (x - max) - log(sum(exp(x - max))) returns a large negative number
+TEST(ONNX, LogSoftmaxLargeRange)
+{
+   SofieReference ref = readReference("LogSoftmaxLargeRange");
+
+   ASSERT_INCLUDE_AND_RUN(std::vector<float>, "LogSoftmaxLargeRange", ref.f32("input0"));
+
+   expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
+}
+
+// same as above, but reducing along the first axis, which takes the generic
+// code path of the SOFIE Softmax operator
+TEST(ONNX, LogSoftmaxLargeRangeAxis0)
+{
+   SofieReference ref = readReference("LogSoftmaxLargeRangeAxis0");
+
+   ASSERT_INCLUDE_AND_RUN(std::vector<float>, "LogSoftmaxLargeRangeAxis0", ref.f32("input0"));
+
+   expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
+}
+
 TEST(ONNX, ConvTranspose1d)
 {
    SofieReference ref = readReference("ConvTranspose1d");

tmva/sofie/test/generate_input_models.py

@@ -3532,6 +3532,103 @@ def make_Log():
     return _model(graph, opset=14, ir_version=7, producer_name='pytorch', producer_version='1.13.1')
 
 
+def make_LogSoftmax1d():
+    """Ops: LogSoftmax"""
+    nodes = [
+        helper.make_node("LogSoftmax", ["X"], ["Y"]),
+    ]
+    graph = helper.make_graph(
+        nodes,
+        "LogSoftmax",
+        inputs=[
+            _vi("X", FLOAT, [3]),
+        ],
+        outputs=[
+            _vi("Y", FLOAT, [3]),
+        ],
+    )
+    return _model(graph, opset=17, ir_version=8)
+
+
+def make_LogSoftmax2d():
+    """Ops: LogSoftmax"""
+    nodes = [
+        helper.make_node("LogSoftmax", ["X"], ["Y"], axis=-1),
+    ]
+    graph = helper.make_graph(
+        nodes,
+        "LogSoftmax",
+        inputs=[
+            _vi("X", FLOAT, [2, 3]),
+        ],
+        outputs=[
+            _vi("Y", FLOAT, [2, 3]),
+        ],
+    )
+    return _model(graph, opset=17, ir_version=8)
+
+
+def make_LogSoftmax3d():
+    """Ops: LogSoftmax"""
+    nodes = [
+        helper.make_node("LogSoftmax", ["X"], ["Y"], axis=1),
+    ]
+    graph = helper.make_graph(
+        nodes,
+        "LogSoftmax",
+        inputs=[
+            _vi("X", FLOAT, [2, 3, 4]),
+        ],
+        outputs=[
+            _vi("Y", FLOAT, [2, 3, 4]),
+        ],
+    )
+    return _model(graph, opset=17, ir_version=8)
+
+
+def make_LogSoftmaxLargeRange():
+    """Ops: LogSoftmax
+
+    The inputs of each row are 100 apart, so exp(x - max) underflows to 0 in
+    float32 for all but the largest one. log(softmax(x)) is -inf there, while
+    the stable form (x - max) - log(sum(exp(x - max))) stays finite."""
+    nodes = [
+        helper.make_node("LogSoftmax", ["X"], ["Y"], axis=-1),
+    ]
+    graph = helper.make_graph(
+        nodes,
+        "LogSoftmax",
+        inputs=[
+            _vi("X", FLOAT, [2, 3]),
+        ],
+        outputs=[
+            _vi("Y", FLOAT, [2, 3]),
+        ],
+    )
+    return _model(graph, opset=17, ir_version=8)
+
+
+def make_LogSoftmaxLargeRangeAxis0():
+    """Ops: LogSoftmax
+
+    Same as LogSoftmaxLargeRange but with the reduction along the first axis,
+    which takes the generic code path of the SOFIE Softmax operator."""
+    nodes = [
+        helper.make_node("LogSoftmax", ["X"], ["Y"], axis=0),
+    ]
+    graph = helper.make_graph(
+        nodes,
+        "LogSoftmax",
+        inputs=[
+            _vi("X", FLOAT, [3, 2]),
+        ],
+        outputs=[
+            _vi("Y", FLOAT, [3, 2]),
+        ],
+    )
+    return _model(graph, opset=17, ir_version=8)
+
+
 def make_MatMul_1D_Constant():
     """Ops: MatMul"""
     nodes = [
@@ -5955,6 +6052,11 @@ def make_IdentityWeightOutput():
     "Linear_32": make_Linear_32,
     "Linear_64": make_Linear_64,
     "Log": make_Log,
+    "LogSoftmax1d": make_LogSoftmax1d,
+    "LogSoftmax2d": make_LogSoftmax2d,
+    "LogSoftmax3d": make_LogSoftmax3d,
+    "LogSoftmaxLargeRange": make_LogSoftmaxLargeRange,
+    "LogSoftmaxLargeRangeAxis0": make_LogSoftmaxLargeRangeAxis0,
     "MatMul_1D_Constant": make_MatMul_1D_Constant,
     "MatMul_Stacked": make_MatMul_Stacked,
     "MatMul_Stacked2": make_MatMul_Stacked2,
@@ -6862,6 +6964,13 @@ def rand_f32(seed, shape):
     "Linear_32": [f32(np.full(3200, 1.0), (32, 100))],
     "Linear_64": [f32(np.full(6400, 1.0), (64, 100))],
     "Log": [f32([1.0, 2.0, 3.0, 4.0], (4,))],
+    "LogSoftmax1d": [f32([-1.0, 0.0, 1.0], (3,))],
+    "LogSoftmax2d": [f32([-1.0, 0.0, 1.0, 2.0, 1.0, 0.0], (2, 3))],
+    "LogSoftmax3d": [rand_f32(11, (2, 3, 4))],
+    # 100 apart per row: exp(x - max) underflows to 0 in float32 away from the
+    # maximum, which log(softmax(x)) turns into -inf
+    "LogSoftmaxLargeRange": [f32([0.0, 100.0, 200.0, 200.0, 100.0, 0.0], (2, 3))],
+    "LogSoftmaxLargeRangeAxis0": [f32([0.0, 100.0, 100.0, 200.0, 200.0, 300.0], (3, 2))],
     "Max": [
         f32([1.0, 2.0, -1.0], (1, 3)),
         f32([3.0, 0.0, 4.0], (1, 3)),
@@ -7698,6 +7807,23 @@ def resolve(name):
     return [y.astype(np.float32)]
 
 
+def _logsoftmax_reference(model, feeds):
+    """LogSoftmax evaluated with the numerically stable formula in double
+    precision.
+
+    The ReferenceEvaluator computes log(softmax(x)), whose exp(x - max)
+    underflows to 0 in float32 once the inputs are more than ~88 apart from the
+    maximum, so it yields -inf where the correct result is a large negative
+    number (see root-project/root#23547)."""
+    node = next(n for n in model.graph.node if n.op_type == "LogSoftmax")
+    axis = next((a.i for a in node.attribute if a.name == "axis"), -1)
+    x = np.asarray(feeds[node.input[0]], dtype=np.float64)
+    axis = axis % x.ndim
+    shifted = x - x.max(axis=axis, keepdims=True)
+    y = shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True))
+    return [y.astype(np.float32)]
+
+
 EXPECTED_OVERRIDES = {
     "GRUBidirectional": _recurrent_reference,
     "LSTMBidirectional": _recurrent_reference,
@@ -7712,6 +7838,8 @@ def resolve(name):
     "IdentityWeightBatchNorm": _batchnorm_reference,
     "BatchNormEpsilon": _batchnorm_reference,
     "BatchNormReluEpsilon": _batchnorm_reference,
+    "LogSoftmaxLargeRange": _logsoftmax_reference,
+    "LogSoftmaxLargeRangeAxis0": _logsoftmax_reference,
 }