Large numerical drift between PyTorch and TFLite when toggling enable_hlfb in ai_edge_torch EmbeddingGemma300

[deeptanshusekhri]

Description

I am exporting the EmbeddingGemma-300M example model to TFLite using ai_edge_torch.convert(). The example PyTorch model matches the reference Hugging Face SentenceTransformer outputs (~1.0 cosine similarity), but TFLite outputs drift heavily depending on whether enable_hlfb is enabled in the internal attention modules.

• With enable_hlfb=True: PyTorch ↔ TFLite similarity is near 1.0
• With enable_hlfb=False: PyTorch ↔ TFLite similarity collapses to ~0.32 cosine on the same inputs
  (while PyTorch outputs remain identical)

This suggests there is some difference in export / lowering / runtime behavior of the non-HLFB path.

---

What I am trying to understand

1. What semantic / numerical differences enable_hlfb introduces
2. Why disabling it changes TFLite numerics so drastically
3. Whether the drift is due to graph construction / lowering differences or TFLite runtime kernel behavior
4. What is the intended way to get “pure TFL ops" (i.e., avoid STABLEHLO_COMPOSITE) without losing numerical correctness

---

Minimal reproduction

The following script:

• Loads EmbeddingGemma from the example implementation
• Patches enable_hlfb
• Exports two TFLite models: hlfb_on and hlfb_off
• Compares TFLite vs PyTorch cosine similarity on a few short sentences
• (PyTorch vs SentenceTransformer is stable and included only as a sanity check)

import os, numpy as np, torch, tensorflow as tf, ai_edge_torch as at
from sentence_transformers import SentenceTransformer
from ai_edge_torch.generative.examples.embedding_gemma.embedding_gemma import build_model
from ai_edge_torch.generative.layers import attention as aiet_attention
from ai_edge_torch.generative.layers import scaled_dot_product_attention as sdpa

CKPT="/path/to/embeddinggemma-300m"
SEQ=256
TEXTS=[
  "Cow went over the moon",
  "The quick brown fox jumps over the lazy dog",
  "Neural networks can approximate complex functions",
]

class Wrap(torch.nn.Module):
  def __init__(s, m): super().__init__(); s.m=m
  def forward(s, t, am): return s.m(t.long(), am.long())

def patch_hlfb(m, on: bool):
  for x in m.modules():
    if isinstance(x, aiet_attention.CausalSelfAttentionBase) or hasattr(x, "enable_hlfb"):
      try: x.enable_hlfb = on
      except: pass
    if hasattr(x, "sdpa_func"):
      x.sdpa_func = sdpa.scaled_dot_product_attention_with_hlfb if on else sdpa.scaled_dot_product_attention

def export_tfl(m, tag):
  path=f"./out/eg_{tag}.tflite"
  os.makedirs(os.path.dirname(path), exist_ok=True)
  samp=(torch.ones(1, SEQ, dtype=torch.int32), torch.ones(1, SEQ, dtype=torch.int32))
  at.convert(m, samp, strict_export=True).export(path)
  return path

def run_tfl(path, ids, am):
  itp=tf.lite.Interpreter(model_path=path)
  itp.allocate_tensors()
  ins=itp.get_input_details()
  out=itp.get_output_details()[0]["index"]
  itp.set_tensor(ins[0]["index"], ids)
  itp.set_tensor(ins[1]["index"], am)
  itp.invoke()
  return itp.get_tensor(out)

def cos(a, b):
  a=a.reshape(-1).astype(np.float32)
  b=b.reshape(-1).astype(np.float32)
  return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12))

st = SentenceTransformer(CKPT, device="cpu")
m = Wrap(build_model(CKPT).eval()).eval()

for on in (True, False):
  tag = "hlfb_on" if on else "hlfb_off"
  patch_hlfb(m, on)
  tfl = export_tfl(m, tag)
  sims = []
  for t in TEXTS:
    enc = st.tokenizer(t, padding="max_length", truncation=True, max_length=SEQ, return_tensors="np")
    ids, am = enc["input_ids"].astype(np.int32), enc["attention_mask"].astype(np.int32)
    with torch.no_grad():
      pt = m(torch.from_numpy(ids), torch.from_numpy(am)).cpu().numpy()
    tf_out = run_tfl(tfl, ids, am)
    sims.append(cos(tf_out, pt))
  print(tag, "avg/min/max", float(np.mean(sims)), float(np.min(sims)), float(np.max(sims)))

---

Observed behavior

Sanity check

• PyTorch vs SentenceTransformer: ~1.0 cosine

Main issue

• TFLite vs PyTorch

  • hlfb_on: good / expected (near 1.0 or close)
  • hlfb_off: ~0.32 cosine average (large drift)

This is consistent across multiple samples.

---

Questions and request for clarification

1. What exactly does enable_hlfb change semantically?
   Is it only inserting a StableHLO composite boundary, or does it cause numerical changes (e.g., different kernel selection, different softmax/logits handling, different masking conventions)?

2. Why would hlfb_off produce large TFLite drift while PyTorch stays stable?

   • Is the non-HLFB path lowered differently into ops that behave differently in LiteRT / TFLite?
   • or could this be related to something else entirely ex softmax stability, or dtype lowering?

3. Is the discrepancy primarily:

   • a graph creation / lowering issue, or
   • a runtime kernel issue in TFLite / LiteRT (XNNPACK / reference kernels)?

4. Is there an intended supported path to avoid STABLEHLO_COMPOSITE while retaining numerics?
   I need “pure TFL ops" for my downstream pipeline, but disabling HLFB currently breaks accuracy.

---

Environment

• ai_edge_torch / litert-torch: 0.7.1
• PyTorch: 2.9.1
• TensorFlow: 2.20.0
• OS & CPU: Ubuntu 24.04.2 LTS (x86_64)

---

[aravindmurali22]

Hello @deeptanshusekhri,

Thanks for filing this issue. My Analysis:

• enable_hlfb=True: This setting wraps the operation in a composite that maps to a fused, numerically stable TFLite custom op (odml.scaled_dot_product_attention).
• enable_hlfb=False: This decomposes torch.nn.functional.scaled_dot_product_attention into elementary ops. It appears that the drift results from numerical instability within this decomposition (likely due to precision loss in the Softmax layer or how the -inf mask is handled).

Suggestion:

To achieve "pure TFL ops" while avoiding this drift, I recommend manually implementing SDPA using standard PyTorch primitives rather than relying on F.sdpa. This approach should help bypass the unstable decomposition.

Please consider using the following replacement in your patch_hlfb for the False case:

import math
import torch

def manual_sdpa(q, k, v, mask=None):
    scale = 1.0 / math.sqrt(q.size(-1))
    scores = torch.matmul(q, k.transpose(-2, -1)) * scale
    if mask is not None:
        scores = scores + mask
    # Cast to float32 for Softmax stability
    attn = torch.softmax(scores.float(), dim=-1).type_as(q)
    return torch.matmul(attn, v)

Let me know if you face any issues. Thanks.

[deeptanshusekhri]

Hi @aravindmurali22 ,

Thanks for the analysis and the suggestion.

I tried replacing the sdpa_func in the enable_hlfb=False path with your manual_sdpa implementation. Unfortunately, I see essentially the same drift as before.

Results (TFLite vs PyTorch cosine similarity):

• manual_sdpa avg/min/max: 0.3175 / 0.2589 / 0.3822

So, while the decomposition stability hypothesis is plausible, this manual implementation does not improve the numerics in my case.

Given that:

• PyTorch outputs remain identical with enable_hlfb toggled, and
• only the exported TFLite model diverges heavily when HLFB is disabled,

to me this still suggests the issue is likely in export/lowering or runtime behavior of the non-HLFB graph, rather than strictly decomposition instability in PyTorch.

could you advise on any other solutions?

[terryheo]

@sirakiin @qukhan any idea on this?

[deeptanshusekhri]

I dug into the SDPA path in litert_torch/generative/layers/scaled_dot_product_attention.py, and the drift in the non-HLFB export appears to be triggered by the GQA head expansion via repeat_interleave.

In the GQA case, the implementation currently uses repeat_interleave to expand KV heads to match Q heads:

k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)

This is from:
https://github.com/google-ai-edge/litert-torch/blob/main/litert_torch/generative/layers/scaled_dot_product_attention.py#L67-L68

I replaced this with the equivalent broadcast-based formulation and the TFLite vs PyTorch cosine similarity comes out to be ~1.0 and the numerical drift disappears.

This probably suggests a loweringissue with repeat_interleave in this specific pattern

Happy to share a minimal patch if useful.