Skip to content

fix temperature flag so it takes effect #981

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
RaviTriv wants to merge 1 commit into
base: dev
Choose a base branch
from
+82 −35
Open

fix temperature flag so it takes effect #981

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 24 additions & 20 deletions ops/ops-inl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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<float>;
Expand All @@ -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
Expand Down
73 changes: 58 additions & 15 deletions ops/ops_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<T>(rng);
e[i] = x[i];
}
for (const float temperature : {1.0f, 2.0f}) {
for (size_t i = 0; i < count; ++i) {
x[i] = Random<T>(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;
Expand All @@ -347,6 +350,45 @@ void TestAllSoftmax() {
hn::ForPartialVectors<ForeachCountAndMisalign<TestSoftmax>>()(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 <class D>
Expand Down Expand Up @@ -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);
Expand Down