Skip to content

Test transformer loss decreases during training #2855

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
GargiGupta-io wants to merge 2 commits into
base: master
Choose a base branch
from
+14 −5
Open

Test transformer loss decreases during training #2855

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cfe685c
Test transformer loss decreases during training
GargiGupta-io Jul 22, 2026
3ba4ab2
Make transformer training assertion more robust
GargiGupta-io Jul 22, 2026
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
19 changes: 14 additions & 5 deletions integration_tests/basic_usage_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@

class BasicUsageTest(unittest.TestCase):
def test_transformer(self):
keras.utils.set_random_seed(1337)

# Tokenize some inputs with a binary label.
vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
sentences = ["The quick brown fox jumped.", "The fox slept."]
Expand All @@ -32,12 +34,19 @@ def test_transformer(self):
outputs = keras.layers.Dense(1, activation="sigmoid")(outputs)
model = keras.Model(inputs, outputs)

# Run a single batch of gradient descent.
model.compile(loss="binary_crossentropy")
loss = model.train_on_batch(x, y)
# Train the model end-to-end.
model.compile(
optimizer=keras.optimizers.Adam(learning_rate=0.01),
loss="binary_crossentropy",
)
initial_loss = model.evaluate(x, y, verbose=0)
model.fit(x, y, epochs=5, verbose=0)
final_loss = model.evaluate(x, y, verbose=0)

# Make sure we have a valid loss.
self.assertGreater(loss, 0)
# Make sure training produces a finite, decreasing loss.
self.assertTrue(np.isfinite(initial_loss))
self.assertTrue(np.isfinite(final_loss))
self.assertLess(final_loss, initial_loss)

@gemini-code-assist gemini-code-assist Bot Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Asserting a strict decrease in loss after only 5 epochs (5 gradient steps) on a tiny dataset of 2 samples can be fragile and prone to flakiness. Due to differences in weight initialization and random number generation across Keras backends (TensorFlow, JAX, PyTorch), the initial weights and training trajectories will differ even with a set seed. With Adam's momentum initialization, 5 steps might not be enough to guarantee a strict decrease in loss on all backends.

Increasing the number of epochs to 15 or 20 ensures that the model has enough steps to consistently minimize the loss and overfit the 2 samples, making the test much more robust across all backends without adding noticeable overhead.

Suggested change
model.fit(x, y, epochs=5, verbose=0)
final_loss = model.evaluate(x, y, verbose=0)
# Make sure we have a valid loss.
self.assertGreater(loss, 0)
# Make sure training produces a finite, decreasing loss.
self.assertTrue(np.isfinite(initial_loss))
self.assertTrue(np.isfinite(final_loss))
self.assertLess(final_loss, initial_loss)
model.fit(x, y, epochs=15, verbose=0)
final_loss = model.evaluate(x, y, verbose=0)
# Make sure training produces a finite, decreasing loss.
self.assertTrue(np.isfinite(initial_loss))
self.assertTrue(np.isfinite(final_loss))
self.assertLess(final_loss, initial_loss)
References
  1. Demand Robustness: Do not accept fragile code. If the proposed code is not robust enough or lacks proper error handling, explicitly tell the author why the current approach is brittle and what must be done to reinforce it. (link)


def test_quickstart(self):
"""This roughly matches the quick start example in our base README."""
Expand Down
Loading