diff --git a/ops/ops-inl.h b/ops/ops-inl.h index 3a1537c7..ce0a28b2 100644 --- a/ops/ops-inl.h +++ b/ops/ops-inl.h @@ -1334,10 +1334,8 @@ HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileUpTo8_BF16( // Computes softmax probabilities for the given logits, normalizing in-place. // The calculation is numerically stable, using the max-subtraction trick to // compute exp(logits[i] - max(logits)) before normalizing by the sum. -// If temperature is provided and not 1.0, each intermediate exp() result is -// divided by temperature before normalization; however, this division by -// temperature cancels out during the final normalization step, meaning -// temperature currently has no effect on the output probabilities. +// Temperature divides the logits before exponentiating; lower values +// concentrate the mass on the max, and zero puts all of it there (ties share). // @param logits In-out: on input, contains logits; on output, overwritten with // probabilities. // @param ctx Input: threading context for parallelism and profiling. @@ -1352,6 +1350,7 @@ static HWY_NOINLINE void Softmax(Logits logits, ThreadingContext& ctx, const SMOptions& sm_options = {}) { GCPP_ZONE(ctx, worker, Zones::kOpsSoftmax); HWY_DASSERT(logits.size() != 0); + HWY_ASSERT(temperature >= 0.0f); namespace hn = hwy::HWY_NAMESPACE; using D = hn::ScalableTag; @@ -1366,24 +1365,29 @@ static HWY_NOINLINE void Softmax(Logits logits, ThreadingContext& ctx, HWY_ATTR { *pmax = hn::Max(*pmax, value); }); vmax = hn::MaxOfLanes(d, vmax); - // Subtract max (avoid precision loss for large exponents) and exponentiate. - hn::Transform(d, logits.data(), logits.size(), - [pmax](const auto d, const V value) HWY_ATTR { - if constexpr (HWY_TARGET & HWY_ALL_SVE) { - // Workaround for buggy SVE codegen: avoid inlined - // FastExpMinusOrZero(). - return hn::CallFastExpMinusOrZero(d, hn::Sub(value, *pmax)); - } else { - return hn::FastExpMinusOrZero(d, hn::Sub(value, *pmax)); - } - }); - - if (temperature != 1.0f) { - const float temperature_inv = 1.0f / temperature; + if (temperature == 0.0f) { hn::Transform(d, logits.data(), logits.size(), - [temperature_inv](const auto d, const V value) HWY_ATTR { - return hn::Mul(value, hn::Set(d, temperature_inv)); + [pmax](const auto d, const V value) HWY_ATTR { + return hn::IfThenElseZero(hn::Eq(value, *pmax), + hn::Set(d, 1.0f)); }); + } else { + // Subtract max to avoid precision loss for large exponents, divide by the + // temperature, and exponentiate. + const float temperature_inv = 1.0f / temperature; + hn::Transform( + d, logits.data(), logits.size(), + [pmax, temperature_inv](const auto d, const V value) HWY_ATTR { + const V scaled = + hn::Mul(hn::Sub(value, *pmax), hn::Set(d, temperature_inv)); + if constexpr (HWY_TARGET & HWY_ALL_SVE) { + // Workaround for buggy SVE codegen: avoid inlined + // FastExpMinusOrZero(). + return hn::CallFastExpMinusOrZero(d, scaled); + } else { + return hn::FastExpMinusOrZero(d, scaled); + } + }); } // Normalize to probability distribution. The exact sum seems like it should diff --git a/ops/ops_test.cc b/ops/ops_test.cc index 2e21ec28..769a713e 100644 --- a/ops/ops_test.cc +++ b/ops/ops_test.cc @@ -309,31 +309,34 @@ class TestSoftmax { T* x = px.get() + misalign_a; T* e = pe.get() + misalign_a; - for (size_t i = 0; i < count; ++i) { - x[i] = Random(rng); - e[i] = x[i]; - } + for (const float temperature : {1.0f, 2.0f}) { + for (size_t i = 0; i < count; ++i) { + x[i] = Random(rng); + e[i] = x[i]; + } - SimpleSoftmax(e, count); - Softmax(Logits(x, count), Ctx(), /*worker=*/0); + SimpleSoftmax(e, count, temperature); + Softmax(Logits(x, count), Ctx(), /*worker=*/0, temperature); - T sum = 0.0f; - for (size_t i = 0; i < count; ++i) { - sum += x[i]; - double rel = std::abs(x[i] - e[i]) / e[i]; - ASSERT_LT(rel, 2e-5) << "Mismatch on coordinate " << i << " out of " - << count; + T sum = 0.0f; + for (size_t i = 0; i < count; ++i) { + sum += x[i]; + double rel = std::abs(x[i] - e[i]) / e[i]; + ASSERT_LT(rel, 2e-5) << "Mismatch on coordinate " << i << " out of " + << count << " at temperature " << temperature; + } + ASSERT_NEAR(sum, 1.0, 2e-5); } - ASSERT_NEAR(sum, 1.0, 2e-5); } private: - static HWY_NOINLINE void SimpleSoftmax(float* HWY_RESTRICT x, size_t size) { + static HWY_NOINLINE void SimpleSoftmax(float* HWY_RESTRICT x, size_t size, + float temperature) { HWY_DASSERT(size != 0); float sum = 0.0; const float maxval = *std::max_element(x, x + size); for (size_t i = 0; i < size; ++i) { - x[i] = std::exp(x[i] - maxval); + x[i] = std::exp((x[i] - maxval) / temperature); sum += x[i]; } const float scale = 1.0f / sum; @@ -347,6 +350,45 @@ void TestAllSoftmax() { hn::ForPartialVectors>()(float()); } +void TestSoftmaxTemperature() { + constexpr size_t kNum = 4; + const float kLogits[kNum] = {2.0f, 1.0f, 0.0f, -1.0f}; + + for (const float temperature : {0.5f, 1.0f, 2.0f}) { + float x[kNum]; + double expected[kNum]; + double sum = 0.0; + for (size_t i = 0; i < kNum; ++i) { + x[i] = kLogits[i]; + expected[i] = std::exp((kLogits[i] - kLogits[0]) / temperature); + sum += expected[i]; + } + Softmax(Logits(x, kNum), Ctx(), /*worker=*/0, temperature); + for (size_t i = 0; i < kNum; ++i) { + EXPECT_NEAR(x[i], expected[i] / sum, 1e-5) + << "Mismatch on coordinate " << i << " at temperature " + << temperature; + } + } + + // Zero temperature puts all the mass on the max. + float zero[kNum]; + std::copy(kLogits, kLogits + kNum, zero); + Softmax(Logits(zero, kNum), Ctx(), /*worker=*/0, /*temperature=*/0.0f); + EXPECT_FLOAT_EQ(zero[0], 1.0f); + for (size_t i = 1; i < kNum; ++i) { + EXPECT_FLOAT_EQ(zero[i], 0.0f) << "Mismatch on coordinate " << i; + } + + // Ties at zero temperature share the mass evenly. + float ties[kNum] = {2.0f, 2.0f, 0.0f, -1.0f}; + Softmax(Logits(ties, kNum), Ctx(), /*worker=*/0, /*temperature=*/0.0f); + EXPECT_FLOAT_EQ(ties[0], 0.5f); + EXPECT_FLOAT_EQ(ties[1], 0.5f); + EXPECT_FLOAT_EQ(ties[2], 0.0f); + EXPECT_FLOAT_EQ(ties[3], 0.0f); +} + class TestSoftmaxState { public: template @@ -866,6 +908,7 @@ HWY_EXPORT_AND_TEST_P(OpsTest, TestAllMulByConst); HWY_EXPORT_AND_TEST_P(OpsTest, TestAllMulByConstTo); HWY_EXPORT_AND_TEST_P(OpsTest, TestAllMulByConstAndAdd); HWY_EXPORT_AND_TEST_P(OpsTest, TestAllSoftmax); +HWY_EXPORT_AND_TEST_P(OpsTest, TestSoftmaxTemperature); HWY_EXPORT_AND_TEST_P(OpsTest, TestAllSoftmaxState); HWY_EXPORT_AND_TEST_P(OpsTest, TestAllCreateDistribution); HWY_EXPORT_AND_TEST_P(OpsTest, TestAllSigmoid);