diff --git a/requirements.txt b/requirements.txt
index 8595248..4d787cf 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -11,7 +11,7 @@
 # Core JAX ecosystem
 jax>=0.4.0
 jaxlib>=0.4.0
-flax>=0.7.0
+flax>=0.8.5
 
 # Optimizers (for training)
 optax>=0.1.0