Release latentASR adapter, code, docs, and reproducibility artifacts

262fa3f verified 3 months ago

90.5 kB

	"""LatentQwenASR model for lr_whisper.

	The ``LatentQwenASR`` class wraps a frozen Qwen3-ASR base model and adds a
	small controller network that emits per-step delta vectors (latent tokens).
	"""

	import math
	import re
	import string
	from typing import Any, Dict, List, Optional, Tuple

	import numpy as np
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from tqdm import tqdm
	from jiwer import wer

	from data import build_audio_token_seq, _coerce_feat_attention_mask

	def _detach_kv(kv_tuple: Any) -> Any:
	"""Detach past_key_values to prevent excessive BPTT exploding gradients."""
	if kv_tuple is None:
	return None
	# If it's a raw tuple (older transformers standard)
	if isinstance(kv_tuple, tuple):
	return tuple(tuple(t.detach() for t in layer) for layer in kv_tuple)

	# Transformers >= 4.38 uses Cache objects like DynamicCache
	if hasattr(kv_tuple, "to_legacy_cache"):
	legacy_tuple = kv_tuple.to_legacy_cache()
	detached_tuple = tuple(tuple(t.detach() for t in layer) for layer in legacy_tuple)
	return type(kv_tuple).from_legacy_cache(detached_tuple)

	# Fallback to cloning attributes if the specific class doesn't match
	import copy
	new_cache = copy.copy(kv_tuple)
	if hasattr(new_cache, "key_cache"):
	new_cache.key_cache = [k.detach() for k in getattr(kv_tuple, "key_cache", [])]
	new_cache.value_cache = [v.detach() for v in getattr(kv_tuple, "value_cache", [])]
	return new_cache


	class LatentQwenASR(nn.Module):
	"""Wrapper around Qwen3-ASR that supports latent prompt injection.

	The class keeps the underlying ASR model frozen and introduces a small
	controller network that emits delta vectors for each latent token. These
	deltas are normalised and scaled, then added to the token embeddings at
	the NT-token positions. During training, the deltas are inserted into the
	decoder embeddings at the appropriate positions; during generation, they
	are applied just before decoding.

	Design note - float32 parameters
	-----------------------------------
	``log_scale`` is kept in float32 even when the rest of the model is in bfloat16
	to preserve gradient precision for the per-step scales.
	"""

	@staticmethod
	def _resolve_text_model(thinker: nn.Module) -> nn.Module:
	"""Resolve the inner text decoder module across wrapped backends (e.g. PEFT)."""
	queue: List[object] = [thinker]
	seen: set = set()
	while queue:
	node = queue.pop(0)
	if node is None:
	continue
	node_id = id(node)
	if node_id in seen:
	continue
	seen.add(node_id)
	emb = getattr(node, "embed_tokens", None)
	if isinstance(emb, nn.Module):
	return node
	for attr in ("model", "base_model", "module"):
	child = getattr(node, attr, None)
	if child is not None and child is not node:
	queue.append(child)
	raise AttributeError(
	"Unable to resolve text decoder with `embed_tokens` from thinker module."
	)

	@staticmethod
	def _resolve_embed_tokens(text_model: nn.Module) -> nn.Module:
	emb = getattr(text_model, "embed_tokens", None)
	if isinstance(emb, nn.Module):
	return emb
	if hasattr(text_model, "get_input_embeddings"):
	emb = text_model.get_input_embeddings()
	if isinstance(emb, nn.Module):
	return emb
	raise AttributeError("Unable to resolve input embedding module.")

	def __init__(
	self,
	asr_model: nn.Module,
	processor: Any,
	n_latent: int,
	nt_token_id: int,
	lang_token_id: int,
	transcribe_token_id: int,
	freeze_base: bool = True,
	use_latent: bool = True,
	use_soft_prompt: bool = False,
	soft_prompt_init_mode: str = "text",
	soft_prompt_init_text: str = "",
	user_prompt_text: str = "Transcribe the audio into text.",

	delta_tanh_c: float = 5.0,
	scale_max: float = 3.0,
	scale_init: float = 0.2,
	thought_mode: str = "prefix",
	thought_group_size: int = 1,
	halt_threshold: float = 0.0,
	latent_drop_prob: float = 0.0,
	latent_input_noise_std: float = 0.0,
	latent_use_bounded_delta: bool = True,
	latent_use_injection_gate: bool = True,
	latent_use_embedding_anchor: bool = True,
	freeze_audio_stack: bool = True,
	) -> None:
	super().__init__()
	self.base_model = asr_model
	self.processor = processor
	self.config = getattr(asr_model, "config", None)
	self.n_latent = int(n_latent)
	self.use_latent = bool(use_latent and self.n_latent > 0)
	self.nt_token_id = int(nt_token_id)
	self.use_soft_prompt = bool(
	use_soft_prompt
	and (not self.use_latent)
	and self.n_latent > 0
	and self.nt_token_id >= 0
	)
	self.soft_prompt_init_mode = (soft_prompt_init_mode or "text").strip().lower()
	self.soft_prompt_init_text = soft_prompt_init_text or ""
	self.lang_token_id = int(lang_token_id)
	self.transcribe_token_id = int(transcribe_token_id)

	# Hyperparams stored from constructor args (originally read from env globals).

	self.delta_tanh_c = float(delta_tanh_c)
	self.halt_threshold = float(halt_threshold)
	self.latent_drop_prob = float(latent_drop_prob)
	self.latent_input_noise_std = float(latent_input_noise_std)
	self.latent_use_bounded_delta = bool(latent_use_bounded_delta)
	self.latent_use_injection_gate = bool(latent_use_injection_gate)
	self.latent_use_embedding_anchor = bool(latent_use_embedding_anchor)
	self.freeze_audio_stack = bool(freeze_audio_stack)

	# Qwen3-ASR audio token ID
	_audio_tok = None
	if self.config and hasattr(self.config, "audio_token_id") and self.config.audio_token_id is not None:
	_audio_tok = int(self.config.audio_token_id)
	if _audio_tok is None:
	_thinker_cfg = getattr(self.base_model, "config", None) or getattr(
	getattr(self.base_model, "thinker", None), "config", None
	)
	if _thinker_cfg is not None and hasattr(_thinker_cfg, "audio_token_id") and _thinker_cfg.audio_token_id is not None:
	_audio_tok = int(_thinker_cfg.audio_token_id)
	if _audio_tok is None:
	for _name in ("<\|AUDIO\|>", "<\|audio\|>", "<\|audio_content\|>", "<\|audio_pad\|>"):
	target_tok = self.processor.tokenizer if hasattr(self.processor, "tokenizer") else self.processor
	if hasattr(target_tok, "convert_tokens_to_ids"):
	_try = target_tok.convert_tokens_to_ids(_name)
	unk_id = getattr(target_tok, "unk_token_id", None)
	if _try is not None and _try != unk_id:
	_audio_tok = int(_try)
	break
	if _audio_tok is None:
	_audio_tok = getattr(self.processor, "audio_token_id", None)
	if _audio_tok is None:
	_audio_tok = 151646 # last-resort fallback
	self.audio_token_id = _audio_tok

	target_tok = self.processor.tokenizer if hasattr(self.processor, "tokenizer") else self.processor
	_decoded = target_tok.convert_ids_to_tokens(self.audio_token_id) if hasattr(target_tok, "convert_ids_to_tokens") else str(self.audio_token_id)
	print(f"[init] audio_token_id={self.audio_token_id} decodes_to={_decoded!r}")

	# Freeze all parameters in the underlying model for latent-adapter training.
	if freeze_base:
	for p in self.base_model.parameters():
	p.requires_grad = False

	if self.freeze_audio_stack:
	audio_tower = getattr(self.base_model, "thinker", self.base_model)
	audio_tower = getattr(audio_tower, "audio_tower", None)
	if audio_tower is not None:
	n_frozen = 0
	for p in audio_tower.parameters():
	if p.requires_grad:
	p.requires_grad = False
	n_frozen += 1
	if n_frozen > 0:
	print(f"[init] Froze {n_frozen} audio_tower parameters.")
	# Also freeze multi_modal_projector (audio-to-text projection).
	projector = getattr(getattr(self.base_model, "thinker", self.base_model), "multi_modal_projector", None)
	if projector is not None:
	n_frozen_proj = 0
	for p in projector.parameters():
	if p.requires_grad:
	p.requires_grad = False
	n_frozen_proj += 1
	if n_frozen_proj > 0:
	print(f"[init] Froze {n_frozen_proj} multi_modal_projector parameters.")

	self.thinker = self.base_model.thinker
	self.text_model = self._resolve_text_model(self.thinker)
	self.embed_tokens = self._resolve_embed_tokens(self.text_model)

	if self.config and hasattr(self.config, "text_config"):
	d = self.config.text_config.hidden_size
	elif self.config and hasattr(self.config, "hidden_size"):
	d = self.config.hidden_size
	else:
	if hasattr(self.embed_tokens, "embedding_dim"):
	d = int(self.embed_tokens.embedding_dim)
	else:
	d = int(self.embed_tokens.weight.size(-1))

	print(f"[LatentQwenASR] Using hidden_size={d}")

	target_dtype = self.thinker.dtype if hasattr(self.thinker, "dtype") else torch.float32

	self.init_proj = nn.Linear(d, d, bias=True).to(dtype=target_dtype)
	self.delta_proj = nn.Linear(d, d, bias=True).to(dtype=target_dtype)

	self.step_embed = nn.Parameter(torch.zeros(self.n_latent, d, dtype=target_dtype))
	nn.init.normal_(self.step_embed, mean=0.0, std=0.02)
	self.step_proj = nn.Linear(d, d, bias=False).to(dtype=target_dtype)
	self.scale_max = float(scale_max)



	scale_init_val = float(scale_init)
	if scale_init_val <= 0:
	scale_init_val = 1.0
	scale_init_unconstrained = math.log(math.expm1(scale_init_val))
	self.log_scale = nn.Parameter(
	torch.full((self.n_latent,), float(scale_init_unconstrained), dtype=torch.float32)
	)

	# Stability: LayerNorm for the recurrent thought loop
	self.thought_ln = nn.LayerNorm(d, elementwise_affine=False).to(dtype=target_dtype)

	# Thought Quality Analyzer (Value Head) predicting expected CE error
	self.value_head = nn.Linear(d, 1).to(dtype=target_dtype)

	# Soft Gated Injection (Scheme 4)
	self.injection_gate = nn.Sequential(
	nn.Linear(d * 2, d),
	nn.Sigmoid()
	).to(dtype=target_dtype)
	# Initialize gate weights to 0 so sigmoid(0) = 0.5 (neutral start)
	nn.init.zeros_(self.injection_gate[0].weight)
	nn.init.zeros_(self.injection_gate[0].bias)
	self._hidden_size = d

	# Front-token prompt tuning
	self.soft_prompt_embed = nn.Parameter(
	torch.zeros(self.n_latent, d, dtype=target_dtype),
	requires_grad=self.use_soft_prompt,
	)
	if self.use_soft_prompt:
	init_prompt = None
	if self.soft_prompt_init_mode == "text" and self.soft_prompt_init_text.strip():
	init_ids = self.processor.tokenizer.encode(
	self.soft_prompt_init_text.strip(),
	add_special_tokens=False,
	)
	if init_ids:
	with torch.no_grad():
	emb_device = self.embed_tokens.weight.device
	init_ids_t = torch.tensor(init_ids, dtype=torch.long, device=emb_device)
	base = self.embed_tokens(init_ids_t).detach()
	if base.numel() > 0:
	if base.size(0) < self.n_latent:
	reps = (self.n_latent + base.size(0) - 1) // base.size(0)
	base = base.repeat((reps, 1))
	init_prompt = base[: self.n_latent]
	if init_prompt is None:
	init_prompt = torch.empty((self.n_latent, d), dtype=target_dtype)
	nn.init.normal_(init_prompt, mean=0.0, std=0.02)
	self.soft_prompt_embed.data.copy_(
	init_prompt.to(device=self.soft_prompt_embed.device, dtype=self.soft_prompt_embed.dtype)
	)
	else:
	self.soft_prompt_embed.requires_grad = False

	# Cast trainable modules to backbone dtype
	target_dtype = self.base_model.dtype if hasattr(self.base_model, "dtype") else torch.float32
	self.init_proj.to(target_dtype)
	self.delta_proj.to(target_dtype)
	self.step_embed.data = self.step_embed.data.to(target_dtype)
	self.step_proj.to(target_dtype)
	self.soft_prompt_embed.data = self.soft_prompt_embed.data.to(target_dtype)
	# log_scale intentionally stays float32 for per-step scale gradient precision.

	# Qwen3-ASR Special Audio Tokens
	self.audio_bos_token_id = self.processor.tokenizer.convert_tokens_to_ids("<\|audio_start\|>")
	self.audio_eos_token_id = self.processor.tokenizer.convert_tokens_to_ids("<\|audio_end\|>")
	if self.audio_bos_token_id is None:
	self.audio_bos_token_id = self.processor.tokenizer.bos_token_id
	if self.audio_eos_token_id is None:
	self.audio_eos_token_id = self.processor.tokenizer.eos_token_id

	# Chat-template token pieces
	self.im_start_id = self.processor.tokenizer.convert_tokens_to_ids("<\|im_start\|>")
	if self.im_start_id is None or self.im_start_id == self.processor.tokenizer.unk_token_id:
	self.im_start_id = self.lang_token_id
	_im_end_id = self.processor.tokenizer.convert_tokens_to_ids("<\|im_end\|>")
	if _im_end_id is None or _im_end_id == self.processor.tokenizer.unk_token_id:
	_im_end_id = self.processor.tokenizer.eos_token_id
	self.im_end_id = _im_end_id
	# Unified stop-token set for generation AND eval truncation.
	_stop_ids = {int(self.im_end_id)}
	_eot_id = self.processor.tokenizer.convert_tokens_to_ids("<\|endoftext\|>")
	if _eot_id is not None and _eot_id != self.processor.tokenizer.unk_token_id:
	_stop_ids.add(int(_eot_id))
	if self.processor.tokenizer.eos_token_id is not None:
	_stop_ids.add(int(self.processor.tokenizer.eos_token_id))
	self.stop_ids = sorted(_stop_ids)
	self.user_nl_ids = self.processor.tokenizer.encode("user\n", add_special_tokens=False)
	self.asst_nl_ids = self.processor.tokenizer.encode("assistant\n", add_special_tokens=False)
	self.asr_prefix_ids = self.processor.tokenizer.encode("language English<asr_text>", add_special_tokens=False)
	self.nl_ids = self.processor.tokenizer.encode("\n", add_special_tokens=False)
	self.user_prompt_text = user_prompt_text
	if self.user_prompt_text:
	self.user_prompt_ids = self.processor.tokenizer.encode(
	self.user_prompt_text, add_special_tokens=False
	)
	else:
	self.user_prompt_ids = []
	self._warned_generate_fallback = False
	self.thought_mode = thought_mode.strip().lower()
	self.thought_group_size = max(1, int(thought_group_size))
	print(
	"[latent-ablation] "
	f"bounded_delta={self.latent_use_bounded_delta} "
	f"injection_gate={self.latent_use_injection_gate} "
	f"embedding_anchor={self.latent_use_embedding_anchor}"
	)

	# In non-latent modes, keep latent controller out of optimization.
	if not self.use_latent:
	for p in self.init_proj.parameters():
	p.requires_grad = False
	for p in self.delta_proj.parameters():
	p.requires_grad = False
	for p in self.step_proj.parameters():
	p.requires_grad = False
	for p in self.thought_ln.parameters():
	p.requires_grad = False
	self.step_embed.requires_grad = False

	self.log_scale.requires_grad = False
	if not self.use_soft_prompt:
	self.soft_prompt_embed.requires_grad = False

	# ------------------------------------------------------------------
	# Internal helpers
	# ------------------------------------------------------------------

	def _build_chat_template_tensors(
	self,
	device: torch.device,
	language: str = "English",
	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
	"""Build shared chat-template pieces for train/eval consistency."""
	system_nl_ids = self.processor.tokenizer.encode("system\n", add_special_tokens=False)
	system_turn = torch.tensor(
	[self.im_start_id] + system_nl_ids + [self.im_end_id] + self.nl_ids,
	dtype=torch.long,
	device=device,
	)
	user_prefix = torch.tensor(
	[self.im_start_id] + self.user_nl_ids,
	dtype=torch.long,
	device=device,
	)
	if self.user_prompt_ids:
	user_prompt = torch.tensor(self.user_prompt_ids, dtype=torch.long, device=device)
	else:
	user_prompt = torch.empty((0,), dtype=torch.long, device=device)
	user_suffix = torch.tensor(
	[self.im_end_id] + self.nl_ids,
	dtype=torch.long,
	device=device,
	)
	asr_prefix_ids = self.processor.tokenizer.encode(f"language {language}<asr_text>", add_special_tokens=False)
	assistant_prefix = torch.tensor(
	[self.im_start_id] + self.asst_nl_ids + asr_prefix_ids,
	dtype=torch.long,
	device=device,
	)
	return system_turn, user_prefix, user_prompt, user_suffix, assistant_prefix

	def _build_full_sequence(
	self,
	audio_toks: torch.Tensor,
	text_all: torch.Tensor,
	system_turn: torch.Tensor,
	user_prefix: torch.Tensor,
	user_prompt: torch.Tensor,
	user_suffix: torch.Tensor,
	label_pad: int = -100,
	n_extra_mask: int = 0,
	mask_all_nt: bool = False,
	) -> Tuple[torch.Tensor, torch.Tensor]:
	"""Build full input_ids and aligned labels for one sample.

	The sequence layout is::

	[system_turn, user_prefix, audio_toks, user_prompt,
	user_suffix, text_all]

	Labels clone input_ids, then mask everything up to (and including)
	the assistant control tokens, plus an optional ``n_extra_mask``
	additional positions (e.g. NT tokens in latent prefix mode).

	When ``mask_all_nt=True``, all NT-token positions anywhere in the
	sequence are additionally masked. Use this for interleaved mode where
	NT tokens are scattered throughout the assistant turn.

	Args:
	audio_toks: 1-D token tensor for the audio segment.
	text_all: 1-D token tensor for the target text (assistant turn).
	system_turn: Shared system-turn token tensor.
	user_prefix: Shared user-prefix token tensor.
	user_prompt: Shared user-prompt token tensor.
	user_suffix: Shared user-suffix token tensor.
	label_pad: Padding value for ignored label positions.
	n_extra_mask: Number of additional positions after the assistant
	prefix to mask (set to ``self.n_latent`` for latent prefix
	mode, 0 for native/baseline/interleaved).
	mask_all_nt: If True, additionally mask every NT-token position
	in the full sequence (for interleaved thought mode).

	Returns:
	``(full_ids, full_labels)`` - 1-D LongTensors.
	"""
	full_ids = torch.cat(
	[system_turn, user_prefix, audio_toks, user_prompt, user_suffix, text_all]
	)
	n_prefix = (
	system_turn.size(0)
	+ user_prefix.size(0)
	+ audio_toks.size(0)
	+ user_prompt.size(0)
	+ user_suffix.size(0)
	)
	# Mask: <\|im_start\|> + "assistant\n" + n_extra_mask (e.g. NT tokens)
	n_asst_control = 1 + len(self.asst_nl_ids) + n_extra_mask
	n_total_mask = n_prefix + n_asst_control
	full_labels = full_ids.clone()
	full_labels[:n_total_mask] = label_pad
	if mask_all_nt and self.nt_token_id >= 0:
	full_labels[full_ids == self.nt_token_id] = label_pad
	return full_ids, full_labels

	def _inject_latent_deltas(
	self,
	inputs_embeds: torch.Tensor,
	input_ids: torch.LongTensor,
	deltas: torch.Tensor,
	) -> torch.Tensor:
	"""Inject latent delta embeddings at NT-token positions.

	Replaces the NT-token positions in inputs_embeds with::

	base_nt_embedding + delta

	Deltas are always injected at full strength. The Value Head controls
	hard routing (N=0 skip / early halt) instead of soft scaling.
	"""
	B = inputs_embeds.size(0)
	K = deltas.size(1)
	device = inputs_embeds.device
	nt_mask = (input_ids == self.nt_token_id)
	nt_counts = nt_mask.sum(dim=1)

	max_nt = int(nt_counts.max().item())
	if K == 1 and max_nt > 1:
	deltas = deltas.expand(B, max_nt, -1)
	K = max_nt

	assert (nt_counts == K).all(), (
	f"NT token count mismatch! {nt_counts} vs expected {K}"
	)
	nt_token_ids = torch.full(
	(B, K), self.nt_token_id, dtype=torch.long, device=device
	)
	base_nt = self.embed_tokens(nt_token_ids).to(dtype=inputs_embeds.dtype)
	delta_scaled = deltas.float().to(dtype=inputs_embeds.dtype)

	if self.latent_use_injection_gate:
	gate_input = torch.cat([base_nt, delta_scaled], dim=-1)
	gating_weight = self.injection_gate(gate_input)
	else:
	gating_weight = torch.ones_like(delta_scaled)
	gated_delta = delta_scaled * gating_weight
	if self.latent_use_embedding_anchor:
	nt_embeds = base_nt + gated_delta
	else:
	nt_embeds = gated_delta
	batch_idx, seq_idx = torch.nonzero(nt_mask, as_tuple=True)
	result = inputs_embeds.clone()
	result[batch_idx, seq_idx, :] = nt_embeds.to(dtype=inputs_embeds.dtype).view(-1, self._hidden_size)
	return result

	def _bounded_delta_from_raw(
	self,
	delta_raw: torch.Tensor,
	step_idx: int,
	scales: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""Convert a raw projected state into the configured latent delta."""
	if not self.latent_use_bounded_delta:
	return delta_raw
	if scales is None:
	scales = self.step_scales()
	delta_raw_f32 = delta_raw.float()
	norm_val = delta_raw_f32.norm(dim=-1, keepdim=True)
	delta_dir = (delta_raw_f32 / (norm_val + 1e-6)).to(dtype=delta_raw.dtype)
	scale = scales[min(int(step_idx), scales.size(0) - 1)]
	return delta_dir * scale

	def _compute_prefix_state(
	self,
	inputs_embeds: torch.Tensor,
	input_ids: torch.LongTensor,
	strict: bool = True,
	language: str = "English",
	) -> Tuple[torch.Tensor, object, torch.LongTensor]:
	"""Compute prefix state and KV cache up to the first NT token."""
	B, _, _ = inputs_embeds.shape
	device = inputs_embeds.device
	text_model = self.text_model

	nt_mask = (input_ids == self.nt_token_id)
	nt_counts = nt_mask.sum(dim=1)
	if strict:
	assert (nt_counts == self.n_latent).all(), (
	f"NT token count mismatch! {nt_counts} vs {self.n_latent}"
	)

	first_nt = nt_mask.float().argmax(dim=1)
	if strict:
	asr_prefix_ids = self.processor.tokenizer.encode(
	f"language {language}<asr_text>", add_special_tokens=False
	)
	expected = torch.tensor(
	[self.im_start_id] + self.asst_nl_ids + asr_prefix_ids, dtype=torch.long, device=device
	)
	exp_len = expected.numel()
	for i in range(B):
	idx = int(first_nt[i].item())
	start = idx - exp_len
	assert start >= 0, (
	f"Prompt too short to contain assistant prefix (idx={idx}, exp_len={exp_len})."
	)
	actual = input_ids[i, start:idx]
	assert torch.equal(actual, expected), (
	"Assistant prefix mismatch before NT. "
	f"expected={expected.tolist()} actual={actual.tolist()}"
	)

	prefix_lens = first_nt
	if strict:
	assert (prefix_lens > 0).all(), f"Invalid prefix length(s): {prefix_lens}"

	max_prefix = int(prefix_lens.max().item())
	prefix_embeds = inputs_embeds[:, :max_prefix, :]
	prefix_attn = (
	torch.arange(max_prefix, device=device).unsqueeze(0) < prefix_lens.unsqueeze(1)
	).long()
	if prefix_attn.numel() > 0:
	prefix_embeds = prefix_embeds * prefix_attn.unsqueeze(-1)

	with torch.no_grad():
	out = text_model(
	inputs_embeds=prefix_embeds,
	attention_mask=prefix_attn,
	use_cache=True,
	return_dict=True,
	)
	last_positions = torch.clamp(prefix_lens - 1, min=0)
	batch_idx = torch.arange(B, device=device)
	last = out.last_hidden_state[batch_idx, last_positions].detach()
	return last, out.past_key_values, prefix_attn

	def _embed_scale(self) -> float:
	# Qwen3-ASR does NOT scale embeddings by sqrt(d_model).
	return 1.0


	def step_scales(self) -> torch.Tensor:
	s = F.softplus(self.log_scale)
	s = torch.clamp(s, max=self.scale_max)
	return s[:self.n_latent]

	def forced_decoder_ids_latent(self) -> List[Tuple[int, int]]:
	ids: List[Tuple[int, int]] = [(0, self.im_start_id)]
	cur = 1
	for tid in self.asst_nl_ids:
	ids.append((cur, tid))
	cur += 1
	for i in range(self.n_latent):
	ids.append((cur + i, self.nt_token_id))
	return ids

	def forced_decoder_ids_baseline(self) -> List[Tuple[int, int]]:
	ids: List[Tuple[int, int]] = [(0, self.im_start_id)]
	cur = 1
	for tid in self.asst_nl_ids:
	ids.append((cur, tid))
	cur += 1
	return ids

	def _shift_right(self, labels: torch.LongTensor) -> torch.LongTensor:
	pad_id = self.processor.tokenizer.eos_token_id
	decoder_input_ids = labels.new_full(labels.shape, pad_id)
	decoder_input_ids[:, 0] = (
	self.processor.tokenizer.bos_token_id
	if self.processor.tokenizer.bos_token_id is not None
	else self.lang_token_id
	)
	shifted = labels[:, :-1].clone()
	shifted = shifted.masked_fill(shifted == -100, pad_id)
	decoder_input_ids[:, 1:] = shifted
	return decoder_input_ids

	def _encode_audio(
	self,
	input_features: torch.FloatTensor,
	feature_attention_mask: torch.LongTensor,
	) -> Tuple[torch.Tensor, List[int]]:
	"""Manually encode audio features to get their exact lengths per sample."""
	if feature_attention_mask is not None:
	feature_lens = torch.sum(feature_attention_mask, dim=1)
	else:
	feature_lens = torch.full(
	(input_features.size(0),), input_features.size(2),
	device=input_features.device, dtype=torch.long,
	)

	audio_features = []
	lengths = []
	audio_tower = self.thinker.audio_tower
	projector = getattr(self.thinker, "multi_modal_projector", None)

	for i in range(len(input_features)):
	feat = input_features[i]
	l = feature_lens[i]
	inp = feat[:, :l]
	length_tensor = l.unsqueeze(0)
	out = audio_tower(inp, feature_lens=length_tensor).last_hidden_state
	if projector is not None:
	out = projector(out)
	out = out.squeeze(0)
	audio_features.append(out)
	lengths.append(out.size(0))

	return torch.cat(audio_features, dim=0), lengths

	def _get_audio_lengths(
	self,
	input_features: torch.FloatTensor,
	feature_attention_mask: torch.LongTensor,
	) -> List[int]:
	"""Compute audio output lengths using the exact same formula as the thinker."""
	if feature_attention_mask is not None:
	feature_lens = torch.sum(feature_attention_mask, dim=1)
	else:
	feature_lens = torch.full(
	(input_features.size(0),), input_features.size(2),
	device=input_features.device, dtype=torch.long,
	)

	input_lengths_leave = feature_lens % 100
	feat_lengths = (input_lengths_leave - 1) // 2 + 1
	output_lengths = (
	((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1
	+ (feature_lens // 100) * 13
	)
	return output_lengths.tolist()

	def _get_native_audio_embeds(
	self,
	input_ids: torch.LongTensor,
	input_features: torch.FloatTensor,
	feature_attention_mask: torch.LongTensor,
	attention_mask: torch.LongTensor,
	) -> torch.Tensor:
	"""Capture the thinker's native audio-fused embeddings via a forward hook."""
	captured: Dict[str, torch.Tensor] = {}

	def _hook(module: nn.Module, args: Any, kwargs: Any) -> None:
	embeds = kwargs.get("inputs_embeds", None)
	if embeds is None and len(args) > 0:
	embeds = args[0]
	if embeds is not None:
	captured["embeds"] = embeds
	return None

	hook = self.text_model.register_forward_pre_hook(_hook, with_kwargs=True)
	try:
	with torch.no_grad():
	self.thinker(
	input_ids=input_ids,
	attention_mask=attention_mask,
	input_features=input_features,
	feature_attention_mask=feature_attention_mask,
	labels=None,
	return_dict=True,
	)
	finally:
	hook.remove()

	if "embeds" not in captured:
	raise RuntimeError(
	"_get_native_audio_embeds: forward hook did not capture embeddings. "
	"The thinker's internal forward may use a different code path."
	)
	return captured["embeds"].detach().clone()

	def _compute_continuous_thoughts(
	self,
	initial_state: torch.Tensor,
	prefix_past_key_values: Optional[object] = None,
	prefix_attention_mask: Optional[torch.LongTensor] = None,
	halt_threshold: Optional[float] = None,
	) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
	"""Compute recurrent latent states + per-step deltas using DEQ fixed-point iteration."""
	B = initial_state.size(0)
	D = initial_state.size(-1)
	device = initial_state.device
	text_model = self.text_model

	initial_norm = initial_state / (initial_state.norm(dim=-1, keepdim=True) + 1e-6)
	h = self.init_proj(initial_norm)

	# Remove random corruption as teacher forcing artificially lowers CE
	_corrupted = False

	scales = self.step_scales()
	thought_embeds: List[torch.Tensor] = []
	state_embeds: List[torch.Tensor] = []
	raw_norms: List[torch.Tensor] = []
	scaled_norms: List[torch.Tensor] = []
	v_preds: List[torch.Tensor] = []
	past_kv = prefix_past_key_values
	running_attn = prefix_attention_mask

	# ---- Dynamic Causal Thinking Loop ----
	# Instead of fixed-point iteration (which doesn't converge on frozen LLMs),
	# we generate a sequence of thought tokens.
	max_iter = self.n_latent
	if halt_threshold is None:
	halt_threshold = getattr(self, "halt_threshold", 0.0)

	current_embed = h

	# ---- N=0 Skip Check (Tanh) ----
	# v_pred ∈ [-1,1]: expected delta CE (positive means LR helps).
	# If v_pred < threshold → model is confident LR is NOT needed → skip.
	if not self.training:
	v_init = torch.tanh(self.value_head(h.unsqueeze(1)).squeeze(-1).squeeze(-1)) # (B,)
	if (v_init < halt_threshold).all():
	# Return empty deltas → generate() will remove all NT tokens
	empty_deltas = torch.zeros(B, 0, D, device=device, dtype=h.dtype)
	empty_states = torch.zeros(B, 0, D, device=device, dtype=h.dtype)
	stats = {
	"predicted_value": v_init.detach(),
	"deq_iters": torch.tensor(0.0, device=device),
	"raw_norm_mean": torch.zeros(0, device=device),
	"raw_norm_std": torch.zeros(0, device=device),
	"scaled_norm_mean": torch.zeros(0, device=device),
	"scaled_norm_std": torch.zeros(0, device=device),
	"cos_mean": torch.zeros(0, device=device),
	"cos_std": torch.zeros(0, device=device),
	"scales": scales.detach(),
	"diff_norm": torch.tensor(0.0, device=device),
	"step_cos": torch.tensor(0.0, device=device),
	"v_preds": v_init.unsqueeze(1).detach(),
	"gate": v_init.unsqueeze(1).detach(),
	"skipped": True,
	}
	return empty_deltas, empty_states, stats

	with torch.no_grad():
	for k in range(max_iter):
	# The input for step k is step_embed[k] projected, added to the normalized state
	step_idx = min(k, self.step_embed.size(0) - 1)
	step = self.step_proj(self.step_embed[step_idx]).view(1, 1, -1).expand(B, 1, -1)

	h_norm = self.thought_ln(current_embed)
	h_input = h_norm.unsqueeze(1) + step

	step_attn = None
	if running_attn is not None:
	# Causal: attention mask GROWS by 1 at each step
	one = torch.ones((B, 1), dtype=running_attn.dtype, device=device)
	step_attn = torch.cat([running_attn, one], dim=1)

	out = text_model(
	inputs_embeds=h_input,
	past_key_values=past_kv,
	attention_mask=step_attn,
	use_cache=True,
	return_dict=True,
	)

	# Next state
	h_next = out.last_hidden_state.squeeze(1)

	# Compute delta to inject into this token's output representation
	delta_raw = self.delta_proj(h_next)
	delta = self._bounded_delta_from_raw(delta_raw, step_idx, scales)

	# Early Halt Check (Value Head, tanh) — BEFORE committing delta
	v_raw = self.value_head(h_next.unsqueeze(1)).squeeze(-1).squeeze(-1) # (B,)
	v_pred = torch.tanh(v_raw)

	# Halt when v_pred drops below threshold (model says this step is harmful)
	if not self.training:
	if (v_pred < halt_threshold).all():
	break

	# Only commit the delta if we did NOT halt
	thought_embeds.append(delta)
	state_embeds.append(h_next)
	raw_norms.append(delta_raw.norm(dim=-1))
	scaled_norms.append(delta.norm(dim=-1))
	v_preds.append(v_pred)

	# Update context for the next iteration
	past_kv = out.past_key_values
	running_attn = step_attn
	current_embed = h_next

	# Grab the past_kv produced exactly prior to the final step
	with torch.no_grad():
	# We need to re-run up to final_step_idx-1 to get the exact past_kv?
	# Actually, no. We can just run from initial_state and re-accumulate.
	# BUT since max_iter <= n_latent is small (e.g. 4), we can just re-run
	# the entire sequence with gradients enabled!
	pass

	# To make things simple and correct for training with such a small sequence (N=4),
	# we can just re-run the FULL causal generation with gradients enabled and overwrite.
	# This gives exact gradients through all thinking steps (BPTT).
	if self.training:
	thought_embeds_grad = []
	state_embeds_grad = []
	current_embed_grad = h # Keep grad flow to init_proj
	running_attn_grad = prefix_attention_mask
	past_kv_grad = prefix_past_key_values

	# Re-run for exactly the number of steps we decided to take in the no_grad pass
	for k in range(len(thought_embeds)):
	step_idx = min(k, self.step_embed.size(0) - 1)
	step = self.step_proj(self.step_embed[step_idx]).view(1, 1, -1).expand(B, 1, -1)

	h_norm = self.thought_ln(current_embed_grad)
	h_input = h_norm.unsqueeze(1) + step

	step_attn = None
	if running_attn_grad is not None:
	one = torch.ones((B, 1), dtype=running_attn_grad.dtype, device=device)
	step_attn = torch.cat([running_attn_grad, one], dim=1)

	out = text_model(
	inputs_embeds=h_input,
	past_key_values=past_kv_grad, # Evolve KV cache like inference!
	attention_mask=step_attn,
	use_cache=True,
	return_dict=True,
	)

	h_next_g = out.last_hidden_state.squeeze(1)

	delta_raw_g = self.delta_proj(h_next_g)
	delta_g = self._bounded_delta_from_raw(delta_raw_g, step_idx, scales)

	thought_embeds_grad.append(delta_g)
	state_embeds_grad.append(h_next_g)

	# Detach state so CE-loss gradients don't flow back to previous steps.
	# Each step is independently trained on its own delta (1-step truncated BPTT).
	past_kv_grad = out.past_key_values
	running_attn_grad = step_attn
	current_embed_grad = h_next_g.detach()

	# Override the no_grad lists with the gradient-tracked ones
	thought_embeds = thought_embeds_grad
	state_embeds = state_embeds_grad

	# Stack sequences — guard against empty list when all steps halted
	if not thought_embeds:
	# All steps rejected by Value Head → return empty (K=0)
	deltas = torch.zeros((B, 0, D), device=device, dtype=h.dtype)
	states = torch.zeros((B, 0, D), device=device, dtype=h.dtype)
	# Use initial state for value prediction
	predicted_value = torch.tanh(self.value_head(h.unsqueeze(1)))
	stats = {
	"predicted_value": predicted_value.squeeze(-1).squeeze(-1),
	"deq_iters": torch.tensor(0.0, device=device),
	"raw_norm_mean": torch.zeros(self.n_latent, device=device),
	"raw_norm_std": torch.zeros(self.n_latent, device=device),
	"scaled_norm_mean": torch.zeros(self.n_latent, device=device),
	"scaled_norm_std": torch.zeros(self.n_latent, device=device),
	"cos_mean": torch.zeros(self.n_latent, device=device),
	"cos_std": torch.zeros(self.n_latent, device=device),
	"scales": scales.detach(),
	"diff_norm": torch.tensor(0.0, device=device),
	"step_cos": torch.tensor(0.0, device=device),
	"v_preds": None,
	"gate": torch.zeros((B, 1), device=device),
	"corrupted": _corrupted,
	"skipped": True,
	}
	return deltas, states, stats

	deltas = torch.stack(thought_embeds, dim=1) # (B, K, D)
	states = torch.stack(state_embeds, dim=1) # (B, K, D)

	# predicted_value: tanh [-1, 1] estimating LR's impact on CE at
	# every latent state. Training supervises each step because inference
	# may halt from any of these value predictions.
	predicted_value = torch.tanh(self.value_head(states)).squeeze(-1) # (B, K)

	# Gate: average tanh v_pred across all steps, scaled to [0, 1] for gated delta injection
	v_preds_stacked = torch.stack(v_preds, dim=1) if v_preds else None # (B, K)
	if v_preds_stacked is not None:
	gate = (v_preds_stacked.mean(dim=1, keepdim=True) + 1.0) / 2.0 # (B, 1) — mean confidence mapped to [0, 1]
	else:
	gate = (predicted_value.squeeze(-1) + 1.0) / 2.0 # fallback, scaled to [0, 1]

	deq_iters = len(thought_embeds)

	raw_tensor = torch.stack(raw_norms, dim=1)
	scaled_tensor = torch.stack(scaled_norms, dim=1)
	scales_val = scales.detach()

	with torch.no_grad():
	flat_states = states.view(-1, D)
	flat_states_norm = F.normalize(flat_states, dim=-1)

	word_emb_weight = self.embed_tokens.weight
	vocab_size = word_emb_weight.size(0)
	if vocab_size > 10000:
	idx = torch.randperm(vocab_size, device=device)[:10000]
	word_emb_sub = word_emb_weight[idx]
	else:
	word_emb_sub = word_emb_weight

	word_emb_norm = F.normalize(word_emb_sub, dim=-1)
	sim_matrix = torch.matmul(flat_states_norm, word_emb_norm.t())
	if sim_matrix.numel() > 0:
	max_sim, _ = sim_matrix.max(dim=-1)
	cos_mean_val = max_sim.mean()
	cos_std_val = max_sim.std(unbiased=False)
	else:
	cos_mean_val = flat_states_norm.new_tensor(0.0)
	cos_std_val = flat_states_norm.new_tensor(0.0)

	if deltas.size(1) > 1:
	diffs = deltas[:, 1:] - deltas[:, :-1]
	diff_norm_val = diffs.norm(dim=-1).mean()
	step_cos_val = F.cosine_similarity(deltas[:, 1:], deltas[:, :-1], dim=-1).mean()
	else:
	diff_norm_val = deltas.new_tensor(0.0)
	step_cos_val = deltas.new_tensor(0.0)

	stats = {
	"predicted_value": predicted_value, # (B, K) for per-step MSE against per-sample target
	"deq_iters": torch.tensor(deq_iters, dtype=torch.float32, device=device),
	"raw_norm_mean": raw_tensor.mean(dim=0).detach() if raw_tensor.ndim > 1 else raw_tensor.detach(),
	"raw_norm_std": raw_tensor.std(dim=0, unbiased=False).detach() if raw_tensor.ndim > 1 else raw_tensor.new_zeros(raw_tensor.size(-1)),
	"scaled_norm_mean": scaled_tensor.mean(dim=0).detach() if scaled_tensor.ndim > 1 else scaled_tensor.detach(),
	"scaled_norm_std": scaled_tensor.std(dim=0, unbiased=False).detach() if scaled_tensor.ndim > 1 else scaled_tensor.new_zeros(scaled_tensor.size(-1)),
	"cos_mean": torch.full((self.n_latent,), cos_mean_val.item(), device=device),
	"cos_std": torch.full((self.n_latent,), cos_std_val.item(), device=device),
	"scales": scales_val,
	"diff_norm": diff_norm_val.detach(),
	"step_cos": step_cos_val.detach(),
	"v_preds": v_preds_stacked.detach() if v_preds_stacked is not None else None,
	"gate": gate.detach(),
	"corrupted": _corrupted,
	}
	return deltas, states, stats

	def _compute_causal_thoughts(
	self,
	inputs_embeds: torch.Tensor,
	input_ids: torch.LongTensor,
	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
	"""Two-pass causal thought injection for interleaved mode.

	Each NT token at position ``k`` receives a thought delta computed from
	the hidden state at position ``k-1`` (the token immediately before it).
	This makes each thought causal: it can attend to all previous context
	including the word generated just before it, which aligns with ASR's
	left-to-right generation structure.

	Pass 1 (``no_grad``): forward the full sequence with base NT
	embeddings to capture all hidden states ``h_all``.

	Pass 2 (with grad): for each NT at position k, compute a delta from
	``h_all[:, k-1, :]`` via ``init_proj → delta_proj → normalize →
	scale``. Inject the modified NT embedding and return the updated
	``inputs_embeds``.

	Args:
	inputs_embeds: ``(B, SeqLen, D)`` embedding tensor (base NT
	embeddings already placed at NT positions by
	``_get_native_audio_embeds``).
	input_ids: ``(B, SeqLen)`` token id tensor to locate NT positions.

	Returns:
	``(modified_embeds, deltas_padded, states_padded, stats)`` where
	``deltas_padded`` and ``states_padded`` have shape
	``(B, max_nt, D)`` (zero-padded to the batch maximum NT count).
	"""
	B, SeqLen, D = inputs_embeds.shape
	device = inputs_embeds.device

	nt_mask = (input_ids == self.nt_token_id) # (B, SeqLen)
	batch_idx, nt_pos = torch.nonzero(nt_mask, as_tuple=True)

	if len(nt_pos) == 0:
	empty_d = inputs_embeds.new_zeros((B, 0, D))
	empty_s = inputs_embeds.new_zeros((B, 0, D))
	dummy = inputs_embeds.new_tensor(0.0)
	stats: Dict[str, torch.Tensor] = {
	"raw_norm_mean": inputs_embeds.new_zeros((0,)),
	"raw_norm_std": inputs_embeds.new_zeros((0,)),
	"scaled_norm_mean": inputs_embeds.new_zeros((0,)),
	"scaled_norm_std": inputs_embeds.new_zeros((0,)),
	"cos_mean": inputs_embeds.new_zeros((0,)),
	"cos_std": inputs_embeds.new_zeros((0,)),
	"scales": inputs_embeds.new_zeros((0,)),
	"diff_norm": dummy,
	"step_cos": dummy,
	}
	return inputs_embeds, empty_d, empty_s, stats

	# --- Pass 1: no_grad forward to capture all hidden states ----------
	with torch.no_grad():
	out1 = self.text_model(
	inputs_embeds=inputs_embeds,
	use_cache=False,
	return_dict=True,
	)
	h_all = out1.last_hidden_state # (B, SeqLen, D) - detached in no_grad

	# Context for each NT at position k is h_all[:, k-1, :]
	ctx_pos = (nt_pos - 1).clamp(min=0) # (N_total,)
	contexts = h_all[batch_idx, ctx_pos] # (N_total, D)

	# --- Pass 2: generate deltas (with grad) ---------------------------
	ctx_norm = F.normalize(contexts.float(), dim=-1).to(dtype=inputs_embeds.dtype)
	h_proj = self.init_proj(ctx_norm) # (N_total, D)
	delta_raw = self.delta_proj(h_proj) # (N_total, D)
	if self.latent_use_bounded_delta:
	delta_raw_f32 = delta_raw.float()
	norm_val = delta_raw_f32.norm(dim=-1, keepdim=True)
	delta_dir = (delta_raw_f32 / (norm_val + 1e-6)).to(dtype=delta_raw.dtype)
	scale = self.step_scales().mean() # scalar - shared across all positions
	delta = delta_dir * scale # (N_total, D)
	else:
	delta = delta_raw

	# Build NT embeddings and inject
	# base_nt_ids and nt_embeds are moved down to use the scaled delta


	# --- Pack into (B, max_nt, D) for loss functions -------------------
	nt_counts = nt_mask.sum(dim=1) # (B,)
	max_nt = int(nt_counts.max().item())
	deltas_padded = inputs_embeds.new_zeros((B, max_nt, D))
	states_padded = inputs_embeds.new_zeros((B, max_nt, D))
	for b in range(B):
	b_mask = (batch_idx == b)
	cnt = int(nt_counts[b].item())
	if cnt > 0:
	deltas_padded[b, :cnt] = delta[b_mask]
	states_padded[b, :cnt] = contexts[b_mask]

	# --- Stats for logging (shape matches prefix-mode stats) ------------
	with torch.no_grad():
	delta_for_stats = delta.float()
	raw_norms = delta_raw.norm(dim=-1) # (N_total,)
	scaled_norms = delta_for_stats.norm(dim=-1)
	flat_ctx_norm = F.normalize(contexts, dim=-1)
	word_emb = self.embed_tokens.weight
	vocab_size = word_emb.size(0)
	if vocab_size > 10000:
	idx = torch.randperm(vocab_size, device=device)[:10000]
	word_emb_sub = word_emb[idx]
	else:
	word_emb_sub = word_emb
	word_emb_norm = F.normalize(word_emb_sub, dim=-1)
	sim_matrix = torch.matmul(flat_ctx_norm.to(word_emb_norm.dtype), word_emb_norm.t())
	max_sim, _ = sim_matrix.max(dim=-1)
	cos_mean_val = max_sim.mean()
	cos_std_val = max_sim.std()
	# Compute diff_norm and step_cos across padded deltas
	if max_nt > 1:
	diff = deltas_padded[:, 1:, :] - deltas_padded[:, :-1, :]
	diff_norm_val = diff.norm(dim=-1).mean()
	step_cos_val = F.cosine_similarity(
	deltas_padded[:, 1:, :], deltas_padded[:, :-1, :], dim=-1
	).mean()
	else:
	diff_norm_val = deltas_padded.new_tensor(0.0)
	step_cos_val = deltas_padded.new_tensor(0.0)

	# Inject deltas into embeddings (grad-tracked for training)
	base_nt_ids = torch.full(
	(len(batch_idx),), self.nt_token_id, dtype=torch.long, device=device
	)
	base_nt = self.embed_tokens(base_nt_ids).to(dtype=inputs_embeds.dtype)
	delta_i = delta.to(dtype=inputs_embeds.dtype)
	if self.latent_use_injection_gate:
	gate_input = torch.cat([base_nt, delta_i], dim=-1)
	delta_i = delta_i * self.injection_gate(gate_input)
	if self.latent_use_embedding_anchor:
	nt_embeds = base_nt + delta_i
	else:
	nt_embeds = delta_i

	result_embeds = inputs_embeds.clone()
	result_embeds[batch_idx, nt_pos] = nt_embeds

	n_out = max_nt
	stats = {
	"raw_norm_mean": raw_norms.mean().unsqueeze(0).expand(n_out).detach(),
	"raw_norm_std": raw_norms.std(unbiased=False).unsqueeze(0).expand(n_out).detach(),
	"scaled_norm_mean": scaled_norms.mean().unsqueeze(0).expand(n_out).detach(),
	"scaled_norm_std": scaled_norms.std(unbiased=False).unsqueeze(0).expand(n_out).detach(),
	"cos_mean": torch.full((n_out,), cos_mean_val.item(), device=device),
	"cos_std": torch.full((n_out,), cos_std_val.item(), device=device),
	"scales": scale.unsqueeze(0).expand(n_out).detach(),
	"diff_norm": diff_norm_val.detach(),
	"step_cos": step_cos_val.detach(),
	}
	return result_embeds, deltas_padded, states_padded, stats

	@torch.no_grad()
	def _generate_interleaved(
	self,
	base_embeds: torch.Tensor,
	attention_mask: torch.LongTensor,
	max_new_tokens: int,
	do_sample: bool = False,
	temperature: float = 1.0,
	) -> torch.LongTensor:
	"""Token-by-token generation with interleaved thought injection.

	Before each group of ``self.thought_group_size`` generated word tokens,
	a thought delta is computed from the current sequence's last hidden
	state, and an NT token (with the modified embedding) is prepended.

	NT tokens in the output are included as ``self.nt_token_id`` so that
	callers using ``skip_special_tokens=True`` in ``tokenizer.decode`` will
	strip them automatically (since ``<\|latent\|>`` is a registered special
	token).

	Args:
	base_embeds: ``(1, SeqLen, D)`` audio-fused embeddings of the
	prompt prefix (no NT tokens appended yet).
	attention_mask: ``(1, SeqLen)`` attention mask for base_embeds.
	max_new_tokens: Maximum number of word tokens to generate.
	do_sample: Whether to sample instead of greedy decode.
	temperature: Sampling temperature (ignored if ``do_sample=False``).

	Returns:
	``(1, N_generated)`` LongTensor of token ids including interleaved
	NT tokens and ending with the first EOS/stop token encountered.
	"""
	eos_ids: set = set(self.stop_ids)
	device = base_embeds.device
	safe_temp = float(temperature) if float(temperature) > 0 else 1.0

	cur_embeds = base_embeds.clone()
	cur_attn = attention_mask.clone()
	generated: List[torch.Tensor] = []
	words_generated = 0

	nt_id_t = torch.tensor([[self.nt_token_id]], dtype=torch.long, device=device)
	base_nt_emb = self.embed_tokens(nt_id_t).to(dtype=cur_embeds.dtype) # (1, 1, D)

	def _append_thought() -> None:
	"""Compute a thought from the last hidden state and inject NT."""
	nonlocal cur_embeds, cur_attn
	# Get context = last hidden state of current sequence
	ctx_out = self.text_model(
	inputs_embeds=cur_embeds,
	use_cache=False,
	return_dict=True,
	)
	h_last = ctx_out.last_hidden_state[:, -1:, :] # (1, 1, D)
	h_norm = F.normalize(h_last.float().squeeze(1), dim=-1).to(dtype=cur_embeds.dtype)
	h_proj = self.init_proj(h_norm)
	delta_raw = self.delta_proj(h_proj)
	if self.latent_use_bounded_delta:
	delta_dir = F.normalize(delta_raw, dim=-1)
	scale = self.step_scales().mean()
	delta_scaled = (delta_dir * scale).float().to(dtype=cur_embeds.dtype)
	else:
	delta_scaled = delta_raw.float().to(dtype=cur_embeds.dtype)

	# Use Tanh Value Head for Interleaved mode as well
	if hasattr(self, "value_head"):
	v_raw = self.value_head(h_norm) # (1, 1)
	v_pred = (torch.tanh(v_raw) + 1.0) / 2.0 # (1, 1) -> scaled to [0, 1]
	halt_thresh = float(getattr(self.config, "halt_threshold", 0.0))
	# Skip injecting NT if predicted value < threshold
	if v_pred.squeeze().item() < halt_thresh:
	return

	delta_i = delta_scaled.unsqueeze(1)
	if self.latent_use_injection_gate:
	gate_input = torch.cat([base_nt_emb.expand_as(delta_i), delta_i], dim=-1)
	delta_i = delta_i * self.injection_gate(gate_input)
	if self.latent_use_embedding_anchor:
	nt_emb_mod = base_nt_emb + delta_i
	else:
	nt_emb_mod = delta_i
	cur_embeds = torch.cat([cur_embeds, nt_emb_mod], dim=1)
	cur_attn = torch.cat([cur_attn, cur_attn.new_ones((1, 1))], dim=1)
	generated.append(nt_id_t.clone())

	for _ in range(max_new_tokens + len(generated)):
	# Inject thought before each new group of words
	if words_generated % self.thought_group_size == 0:
	_append_thought()

	# Generate next word token
	out = self.thinker(
	inputs_embeds=cur_embeds,
	attention_mask=cur_attn,
	input_features=None,
	use_cache=False,
	return_dict=True,
	)
	next_logits = out.logits[:, -1, :]
	if do_sample:
	probs = F.softmax(next_logits / safe_temp, dim=-1)
	next_token = torch.multinomial(probs, num_samples=1)
	else:
	next_token = torch.argmax(next_logits, dim=-1, keepdim=True)

	tok_id = int(next_token[0, 0].item())
	generated.append(next_token)
	if tok_id in eos_ids:
	break

	next_emb = self.embed_tokens(next_token)
	cur_embeds = torch.cat([cur_embeds, next_emb], dim=1)
	cur_attn = torch.cat([cur_attn, cur_attn.new_ones((1, 1))], dim=1)
	words_generated += 1

	if not generated:
	return torch.zeros((1, 0), dtype=torch.long, device=device)
	return torch.cat(generated, dim=1)

	# ------------------------------------------------------------------
	# Forward passes
	# ------------------------------------------------------------------

	def _forward_native(
	self,
	input_features: torch.FloatTensor,
	labels: torch.LongTensor,
	feature_attention_mask: torch.LongTensor,
	language_hint: Optional[str] = None,
	) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
	"""Native thinker forward for LoRA / baseline FT."""
	B = input_features.size(0)
	device = input_features.device

	pad_token_id = self.processor.tokenizer.pad_token_id
	if pad_token_id is None:
	pad_token_id = self.processor.tokenizer.eos_token_id or 151643

	input_ids_list = []
	labels_list = []
	label_pad = -100

	audio_lengths = self._get_audio_lengths(input_features, feature_attention_mask)
	system_turn, user_prefix, user_prompt, user_suffix, _ = (
	self._build_chat_template_tensors(device, language=language_hint or "English")
	)

	for i in range(B):
	valid_len = (labels[i] != label_pad).sum().item()
	text_all = labels[i, :valid_len]
	l_audio = audio_lengths[i]

	audio_toks = build_audio_token_seq(
	self.audio_bos_token_id, self.audio_token_id, self.audio_eos_token_id,
	l_audio, device,
	)
	full_ids, full_labels = self._build_full_sequence(
	audio_toks, text_all, system_turn, user_prefix, user_prompt, user_suffix,
	label_pad=label_pad, n_extra_mask=0,
	)
	input_ids_list.append(full_ids)
	labels_list.append(full_labels)

	max_len = max(x.size(0) for x in input_ids_list)
	input_ids_batch = torch.full((B, max_len), pad_token_id, dtype=torch.long, device=device)
	labels_batch = torch.full((B, max_len), label_pad, dtype=torch.long, device=device)
	for i in range(B):
	l = input_ids_list[i].size(0)
	input_ids_batch[i, :l] = input_ids_list[i]
	labels_batch[i, :l] = labels_list[i]

	attention_mask = (input_ids_batch != pad_token_id).long()

	out = self.thinker(
	input_ids=input_ids_batch,
	attention_mask=attention_mask,
	input_features=input_features,
	feature_attention_mask=feature_attention_mask,
	labels=None,
	return_dict=True,
	)
	logits = out.logits

	# One-time audio diagnostic
	if not getattr(self, "_native_diag_done", False):
	self._native_diag_done = True
	import sys as _s
	_d = _s.stderr.write
	_d(f"\n[native-diag] thinker type: {type(self.thinker).__name__}\n")
	_d(f"[native-diag] has base_model (PEFT): {hasattr(self.thinker, 'base_model')}\n")
	n_audio_ph = (input_ids_batch[0] == self.audio_token_id).sum().item()
	feat_valid = (
	feature_attention_mask[0].sum().item()
	if feature_attention_mask is not None
	else input_features.size(2)
	)
	_d(f"[native-diag] audio_token_id={self.audio_token_id} decoded={self.processor.tokenizer.convert_ids_to_tokens(self.audio_token_id)!r}\n")
	_d(f"[native-diag] input_features: {list(input_features.shape)}, valid_mel_frames={feat_valid}\n")
	_d(f"[native-diag] input_ids shape: {list(input_ids_batch.shape)}\n")
	_d(f"[native-diag] input_features shape: {list(input_features.shape)}\n")
	_d(f"[native-diag] audio placeholder tokens in input_ids[0]: {n_audio_ph}\n")
	_d(f"[native-diag] audio_bos/eos in ids[0]: bos={int((input_ids_batch[0]==self.audio_bos_token_id).sum())} eos={int((input_ids_batch[0]==self.audio_eos_token_id).sum())}\n")
	with torch.no_grad():
	try:
	zero_feat = torch.zeros_like(input_features[:1])
	out_noaud = self.thinker(
	input_ids=input_ids_batch[:1],
	attention_mask=attention_mask[:1],
	input_features=zero_feat,
	feature_attention_mask=feature_attention_mask[:1] if feature_attention_mask is not None else None,
	labels=None,
	return_dict=True,
	)
	diff = (logits[0] - out_noaud.logits[0]).abs().mean().item()
	_d(f"[native-diag] logit diff (real vs zero audio): {diff:.6f}\n")
	if diff < 0.01:
	_d(f"[native-diag] ⚠️ AUDIO NOT BEING USED! Logits identical with zero audio.\n")
	else:
	_d(f"[native-diag] ✓ Audio IS affecting logits (diff={diff:.4f})\n")
	except Exception as e:
	_d(f"[native-diag] zero-audio test failed: {e}\n")
	tok_str = self.processor.tokenizer.decode(input_ids_batch[0, :15].tolist())
	_d(f"[native-diag] input_ids[0][:15] decoded: {repr(tok_str)}\n")
	_s.stderr.flush()

	# Pre-shift labels for loss
	labels_for_loss = torch.full_like(labels_batch, label_pad)
	labels_for_loss[:, :-1] = labels_batch[:, 1:]

	D = self._hidden_size
	dummy = logits.new_tensor(0.0)
	stats = {
	"raw_norm_mean": logits.new_zeros((0,)),
	"raw_norm_std": logits.new_zeros((0,)),
	"scaled_norm_mean": logits.new_zeros((0,)),
	"scaled_norm_std": logits.new_zeros((0,)),
	"cos_mean": logits.new_zeros((0,)),
	"cos_std": logits.new_zeros((0,)),
	"scales": logits.new_zeros((0,)),
	"diff_norm": dummy,
	"step_cos": dummy,
	}
	deltas = logits.new_zeros((B, 0, D))
	states = logits.new_zeros((B, 0, D))
	initial_state = logits.new_zeros((B, D))

	return logits, stats, deltas, states, labels_for_loss, initial_state

	def forward(
	self,
	input_features: torch.FloatTensor,
	labels: torch.LongTensor,
	feature_attention_mask: Optional[torch.LongTensor] = None,
	global_step: int = 0,
	language_hint: Optional[str] = None,
	) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
	"""Forward pass for training."""
	B = input_features.size(0)
	device = input_features.device

	if feature_attention_mask is None:
	feature_attention_mask = torch.ones(
	(B, input_features.size(2)), dtype=torch.long, device=device
	)

	# Fast path: use thinker's native forward for LoRA / baseline FT.
	if not self.use_latent and not self.use_soft_prompt:
	return self._forward_native(input_features, labels, feature_attention_mask, language_hint=language_hint)

	audio_lengths = self._get_audio_lengths(input_features, feature_attention_mask)

	new_input_ids_list = []
	new_labels_list = []

	label_pad = -100
	pad_token_id = self.processor.tokenizer.pad_token_id
	if pad_token_id is None:
	pad_token_id = self.processor.tokenizer.eos_token_id or 151643

	system_turn, user_prefix, user_prompt, user_suffix, _ = (
	self._build_chat_template_tensors(device, language=language_hint or "English")
	)

	interleaved_mode = self.use_latent and self.thought_mode == "interleaved"

	# +++ LATENT DROP +++
	drop_latent_this_batch = False
	if self.training and self.use_latent and not interleaved_mode:
	drop_prob = self.latent_drop_prob
	if drop_prob > 0.0 and torch.rand(1).item() < drop_prob:
	drop_latent_this_batch = True

	actual_n_latent = 0 if drop_latent_this_batch else self.n_latent

	for i in range(B):
	valid_len = (labels[i] != label_pad).sum().item()
	text_all = labels[i, :valid_len]
	if drop_latent_this_batch:
	text_all = text_all[text_all != self.nt_token_id]

	l_audio = audio_lengths[i]

	audio_toks = build_audio_token_seq(
	self.audio_bos_token_id, self.audio_token_id, self.audio_eos_token_id,
	l_audio, device,
	)
	if interleaved_mode:
	# NT tokens are scattered throughout text_all; mask all NT positions.
	full_ids, full_labels = self._build_full_sequence(
	audio_toks, text_all, system_turn, user_prefix, user_prompt, user_suffix,
	label_pad=label_pad, n_extra_mask=0, mask_all_nt=True,
	)
	else:
	# Prefix mode: all NT tokens appear contiguously after asst_nl.
	full_ids, full_labels = self._build_full_sequence(
	audio_toks, text_all, system_turn, user_prefix, user_prompt, user_suffix,
	label_pad=label_pad, n_extra_mask=actual_n_latent,
	)
	new_input_ids_list.append(full_ids)
	new_labels_list.append(full_labels)

	max_len = max(x.size(0) for x in new_input_ids_list)
	new_input_ids = torch.full((B, max_len), pad_token_id, dtype=torch.long, device=device)
	new_labels_raw = torch.full((B, max_len), label_pad, dtype=torch.long, device=device)

	for i in range(B):
	l = new_input_ids_list[i].size(0)
	new_input_ids[i, :l] = new_input_ids_list[i]
	new_labels_raw[i, :l] = new_labels_list[i]

	attention_mask = (new_input_ids != pad_token_id).long()

	inputs_embeds = self._get_native_audio_embeds(
	input_ids=new_input_ids,
	input_features=input_features,
	feature_attention_mask=feature_attention_mask,
	attention_mask=attention_mask,
	)

	# +++ LATENT INPUT NOISE (Value Head Regularization) +++
	# Save pre-LR embeddings (with noise) for fair baseline CE comparison
	pre_lr_embeds = None
	if self.training:
	noise_std = self.latent_input_noise_std
	if noise_std > 0.0:
	noise = torch.randn_like(inputs_embeds) * noise_std
	inputs_embeds = inputs_embeds + noise
	# Save the (possibly noisy) embeddings BEFORE LR injection
	if self.use_latent:
	pre_lr_embeds = inputs_embeds.detach().clone()

	# Latent / soft-prompt injection
	if self.use_latent and interleaved_mode:
	# Interleaved: two-pass causal thought generation.
	inputs_embeds, deltas, states, stats = self._compute_causal_thoughts(
	inputs_embeds, new_input_ids
	)
	initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
	elif self.use_latent:
	if actual_n_latent == 0:
	initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
	deltas = inputs_embeds.new_zeros((B, 0, self._hidden_size))
	states = inputs_embeds.new_zeros((B, 0, self._hidden_size))
	stats = {
	"raw_norm_mean": inputs_embeds.new_zeros((0,)),
	"raw_norm_std": inputs_embeds.new_zeros((0,)),
	"scaled_norm_mean": inputs_embeds.new_zeros((0,)),
	"scaled_norm_std": inputs_embeds.new_zeros((0,)),
	"cos_mean": inputs_embeds.new_zeros((0,)),
	"cos_std": inputs_embeds.new_zeros((0,)),
	"scales": inputs_embeds.new_zeros((0,)),
	"diff_norm": inputs_embeds.new_tensor(0.0),
	"step_cos": inputs_embeds.new_tensor(0.0),
	}
	else:
	# Prefix mode: recurrent thought loop from audio-prefix state.
	initial_state, prefix_past_kv, prefix_attn = self._compute_prefix_state(
	inputs_embeds, new_input_ids, strict=True
	)
	deltas, states, stats = self._compute_continuous_thoughts(
	initial_state,
	prefix_past_key_values=prefix_past_kv,
	prefix_attention_mask=prefix_attn,
	)
	inputs_embeds = self._inject_latent_deltas(
	inputs_embeds, new_input_ids, deltas
	)
	elif self.use_soft_prompt:
	nt_mask = (new_input_ids == self.nt_token_id)
	nt_counts = nt_mask.sum(dim=1)
	assert (nt_counts == self.n_latent).all(), (
	f"NT token count mismatch! {nt_counts} vs {self.n_latent}"
	)
	prompt = self.soft_prompt_embed.to(dtype=inputs_embeds.dtype)
	prompt = prompt.unsqueeze(0).expand(B, -1, -1).contiguous()
	batch_idx, seq_idx = torch.nonzero(nt_mask, as_tuple=True)
	prompt_embeds = inputs_embeds.clone()
	prompt_embeds[batch_idx, seq_idx, :] = prompt.view(-1, self._hidden_size)
	inputs_embeds = prompt_embeds
	initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
	deltas = inputs_embeds.new_zeros((B, 0, self._hidden_size))
	states = inputs_embeds.new_zeros((B, 0, self._hidden_size))
	stats = {
	"raw_norm_mean": inputs_embeds.new_zeros((0,)),
	"raw_norm_std": inputs_embeds.new_zeros((0,)),
	"scaled_norm_mean": inputs_embeds.new_zeros((0,)),
	"scaled_norm_std": inputs_embeds.new_zeros((0,)),
	"cos_mean": inputs_embeds.new_zeros((0,)),
	"cos_std": inputs_embeds.new_zeros((0,)),
	"scales": inputs_embeds.new_zeros((0,)),
	"diff_norm": inputs_embeds.new_tensor(0.0),
	"step_cos": inputs_embeds.new_tensor(0.0),
	}
	else:
	initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
	deltas = inputs_embeds.new_zeros((B, 0, self._hidden_size))
	states = inputs_embeds.new_zeros((B, 0, self._hidden_size))
	stats = {
	"raw_norm_mean": inputs_embeds.new_zeros((0,)),
	"raw_norm_std": inputs_embeds.new_zeros((0,)),
	"scaled_norm_mean": inputs_embeds.new_zeros((0,)),
	"scaled_norm_std": inputs_embeds.new_zeros((0,)),
	"cos_mean": inputs_embeds.new_zeros((0,)),
	"cos_std": inputs_embeds.new_zeros((0,)),
	"scales": inputs_embeds.new_zeros((0,)),
	"diff_norm": inputs_embeds.new_tensor(0.0),
	"step_cos": inputs_embeds.new_tensor(0.0),
	}

	out = self.thinker(
	inputs_embeds=inputs_embeds,
	input_ids=new_input_ids,
	attention_mask=attention_mask,
	input_features=None,
	feature_attention_mask=None,
	labels=None,
	use_cache=False,
	return_dict=True,
	)
	logits = out.logits

	new_labels = torch.full_like(new_labels_raw, label_pad)
	new_labels[:, :-1] = new_labels_raw[:, 1:]

	if logits.size(1) != new_labels.size(1):
	if logits.size(1) > new_labels.size(1):
	shift = logits.size(1) - new_labels.size(1)
	logits = logits[:, shift:, :]
	else:
	new_labels = new_labels[:, -logits.size(1):]

	# ---- Baseline CE vs LR CE for Value Head MSE (Delta CE) ----
	# Compute per-sample CE for both baseline (no LR) and LR-enhanced
	# forward passes. The difference (baseline_ce - lr_ce) forms the
	# continuous target for the Value Head via tanh compression.
	if self.use_latent and self.training and deltas.size(1) > 0:
	with torch.no_grad():
	# Use the SAME (possibly noisy) embeddings but WITHOUT LR deltas
	# This ensures a fair comparison: noisy-no-LR vs noisy-with-LR
	if pre_lr_embeds is not None:
	baseline_embeds = pre_lr_embeds
	else:
	baseline_embeds = self._get_native_audio_embeds(
	input_ids=new_input_ids,
	input_features=input_features,
	feature_attention_mask=feature_attention_mask,
	attention_mask=attention_mask,
	)
	baseline_out = self.thinker(
	inputs_embeds=baseline_embeds,
	input_ids=new_input_ids,
	attention_mask=attention_mask,
	input_features=None,
	feature_attention_mask=None,
	labels=None,
	use_cache=False,
	return_dict=True,
	)
	baseline_logits = baseline_out.logits
	if baseline_logits.size(1) != new_labels.size(1):
	if baseline_logits.size(1) > new_labels.size(1):
	shift = baseline_logits.size(1) - new_labels.size(1)
	baseline_logits = baseline_logits[:, shift:, :]
	else:
	bl_labels = new_labels[:, -baseline_logits.size(1):]
	else:
	bl_labels = new_labels

	# Per-SAMPLE CE: for each sample in batch, compute CE on valid tokens.
	# This is used by train.py to form the continuous delta_ce target.
	valid_mask = (bl_labels != -100) # (B, T)
	valid_per_sample = valid_mask.float().sum(dim=1).clamp(min=1) # (B,)

	# For value head accuracy: exclude asr_prefix tokens
	# ("language", "English", "<asr_text>") which are trivially correct
	# and dilute the real transcript accuracy signal.
	asr_prefix_set = set(self.asr_prefix_ids)
	asr_prefix_mask = torch.zeros_like(bl_labels, dtype=torch.bool)
	for tid in asr_prefix_set:
	asr_prefix_mask \|= (bl_labels == tid)
	acc_mask = valid_mask & ~asr_prefix_mask # exclude prefix tokens
	acc_per_sample = acc_mask.float().sum(dim=1).clamp(min=1)

	if not hasattr(self, '_acc_debug_printed'):
	self._acc_debug_printed = True
	total_tokens = bl_labels.size(1)
	ce_valid = valid_per_sample.mean().item()
	acc_valid = acc_per_sample.mean().item()
	sample_acc_toks = bl_labels[0][acc_mask[0]][:20].tolist()
	sample_decoded = self.processor.tokenizer.decode(sample_acc_toks)
	print(f"\n[ACC-DEBUG] total={total_tokens} ce_valid={ce_valid:.0f} acc_valid={acc_valid:.0f}")
	print(f"[ACC-DEBUG] first 20 acc label ids: {sample_acc_toks}")
	print(f"[ACC-DEBUG] decoded: {sample_decoded}")
	if valid_mask.any():
	# Per-sample baseline CE (uses full valid_mask for loss)
	loss_fct_none = torch.nn.CrossEntropyLoss(ignore_index=-100, reduction='none')
	bl_ce_per_token = loss_fct_none(
	baseline_logits.reshape(-1, baseline_logits.size(-1)),
	bl_labels.reshape(-1),
	).view(B, -1) # (B, T)
	baseline_ce_per_sample = bl_ce_per_token.sum(dim=1) / valid_per_sample # (B,)
	stats["baseline_ce"] = baseline_ce_per_sample # (B,)

	# Per-sample LR CE (for fair Delta CE comparison)
	lr_logits_for_ce = logits
	if lr_logits_for_ce.size(1) != bl_labels.size(1):
	lr_logits_for_ce = lr_logits_for_ce[:, -bl_labels.size(1):, :]
	lr_ce_per_token = loss_fct_none(
	lr_logits_for_ce.reshape(-1, lr_logits_for_ce.size(-1)),
	bl_labels.reshape(-1),
	).view(B, -1)
	lr_ce_per_sample = lr_ce_per_token.sum(dim=1) / valid_per_sample
	stats["lr_ce"] = lr_ce_per_sample # (B,)

	# Per-sample accuracy for VALUE HEAD (uses acc_mask, excludes prefix)
	baseline_preds = baseline_logits.argmax(dim=-1) # (B, T)
	lr_preds = logits.argmax(dim=-1) # (B, T)
	if lr_preds.size(1) != bl_labels.size(1):
	lr_preds = lr_preds[:, -bl_labels.size(1):]
	baseline_correct = ((baseline_preds == bl_labels) & acc_mask).float()
	lr_correct = ((lr_preds == bl_labels) & acc_mask).float()
	baseline_acc = baseline_correct.sum(dim=1) / acc_per_sample # (B,)
	lr_acc = lr_correct.sum(dim=1) / acc_per_sample # (B,)
	stats["baseline_acc"] = baseline_acc
	stats["lr_acc"] = lr_acc

	# Error-conditional metrics for value head target (uses acc_mask)
	baseline_wrong = ((baseline_preds != bl_labels) & acc_mask) # (B, T)
	baseline_right = ((baseline_preds == bl_labels) & acc_mask) # (B, T)
	lr_fixes = ((lr_preds == bl_labels) & baseline_wrong).float().sum(dim=1) # (B,)
	lr_breaks = ((lr_preds != bl_labels) & baseline_right).float().sum(dim=1) # (B,)
	baseline_errors = baseline_wrong.float().sum(dim=1).clamp(min=1) # (B,)
	stats["lr_fixes"] = lr_fixes
	stats["lr_breaks"] = lr_breaks
	stats["baseline_errors"] = baseline_errors
	else:
	B_size = logits.size(0)
	stats["baseline_ce"] = logits.new_zeros(B_size)
	stats["baseline_acc"] = logits.new_ones(B_size)
	stats["lr_acc"] = logits.new_ones(B_size)
	stats["lr_fixes"] = logits.new_zeros(B_size)
	stats["lr_breaks"] = logits.new_zeros(B_size)
	stats["baseline_errors"] = logits.new_ones(B_size)

	return logits, stats, deltas, states, new_labels, initial_state

	@torch.no_grad()
	def _generate_with_forward_fallback(
	self,
	start_ids: torch.LongTensor,
	inputs_embeds: torch.FloatTensor,
	attention_mask: torch.LongTensor,
	max_new_tokens: int,
	eos_token_id: Optional[Any],
	do_sample: bool,
	temperature: float,
	) -> torch.LongTensor:
	"""Fallback decoder when ``.generate(inputs_embeds=...)`` is unsupported."""
	if eos_token_id is None:
	eos_ids: List[int] = []
	elif isinstance(eos_token_id, (list, tuple, set)):
	eos_ids = [int(x) for x in eos_token_id]
	else:
	eos_ids = [int(eos_token_id)]

	generated_tokens: List[torch.Tensor] = []
	safe_temp = float(temperature) if float(temperature) > 0 else 1.0

	cur_ids = start_ids.clone()
	cur_embeds = inputs_embeds.clone()
	cur_attn = attention_mask.clone()

	for _ in range(int(max_new_tokens)):
	out = self.thinker(
	inputs_embeds=cur_embeds,
	attention_mask=cur_attn,
	input_features=None,
	use_cache=False,
	return_dict=True,
	)
	next_logits = out.logits[:, -1, :]

	if do_sample:
	probs = F.softmax(next_logits / safe_temp, dim=-1)
	next_token = torch.multinomial(probs, num_samples=1)
	else:
	next_token = torch.argmax(next_logits, dim=-1, keepdim=True)

	generated_tokens.append(next_token)
	if eos_ids and int(next_token[0, 0].item()) in eos_ids:
	break

	next_embed = self.embed_tokens(next_token)
	cur_ids = torch.cat([cur_ids, next_token], dim=1)
	cur_embeds = torch.cat([cur_embeds, next_embed], dim=1)
	cur_attn = torch.cat(
	[
	cur_attn,
	torch.ones(
	(cur_attn.size(0), 1),
	dtype=cur_attn.dtype,
	device=cur_attn.device,
	),
	],
	dim=1,
	)

	if not generated_tokens:
	return start_ids.new_empty((start_ids.size(0), 0))
	return torch.cat(generated_tokens, dim=1)

	@torch.no_grad()
	def generate(
	self,
	input_features: torch.FloatTensor,
	feature_attention_mask: Optional[torch.LongTensor] = None,
	max_new_tokens: int = 128,
	use_baseline: bool = False,
	return_thoughts: bool = False,
	return_stats: bool = False,
	language_hint: Optional[str] = None,
	prompt_text: Optional[str] = None,
	**gen_kwargs: Any,
	) -> Any:
	"""Generate a transcription with or without latent prompting."""
	language_hint = gen_kwargs.pop("language_hint", language_hint)
	prompt_text = gen_kwargs.pop("prompt_text", prompt_text)
	dynamic_halt_threshold = gen_kwargs.pop("dynamic_halt_threshold", getattr(self, "halt_threshold", 0.0))

	B = input_features.size(0)
	assert B == 1, "Evaluation assumes batch size 1."
	device = input_features.device

	if feature_attention_mask is None:
	feature_attention_mask = torch.ones(
	(B, input_features.size(2)), dtype=torch.long, device=device
	)
	audio_lengths = self._get_audio_lengths(input_features, feature_attention_mask)

	pad_token_id = self.processor.tokenizer.pad_token_id
	if pad_token_id is None:
	pad_token_id = self.processor.tokenizer.eos_token_id or 151643

	l_audio = audio_lengths[0]
	audio_toks = build_audio_token_seq(
	self.audio_bos_token_id, self.audio_token_id, self.audio_eos_token_id,
	l_audio, device,
	).unsqueeze(0)

	system_turn, user_prefix, user_prompt, user_suffix, assistant_prefix = (
	self._build_chat_template_tensors(device, language=language_hint or "English")
	)
	prefix = torch.cat(
	[system_turn, user_prefix, audio_toks.squeeze(0), user_prompt, user_suffix, assistant_prefix],
	dim=0,
	).unsqueeze(0)

	latent_active = (not use_baseline) and self.use_latent
	soft_prompt_active = (not use_baseline) and self.use_soft_prompt
	interleaved_active = latent_active and self.thought_mode == "interleaved"

	# For prefix mode: prepend NT tokens to the prompt.
	# For interleaved mode: no upfront NT tokens; they are injected during generation.
	front_prompt_active = latent_active or soft_prompt_active
	if front_prompt_active and not interleaved_active:
	nt_toks = torch.tensor([self.nt_token_id] * self.n_latent, device=device).unsqueeze(0)
	start_ids = torch.cat([prefix, nt_toks], dim=1)
	else:
	start_ids = prefix

	attention_mask = (start_ids != pad_token_id).long()
	inputs_embeds = self._get_native_audio_embeds(
	input_ids=start_ids,
	input_features=input_features,
	feature_attention_mask=feature_attention_mask,
	attention_mask=attention_mask,
	)

	states = None
	stats = {}
	if interleaved_active:
	# Interleaved generate: NT tokens injected token-by-token.
	fallback_do_sample = bool(gen_kwargs.get("do_sample", False))
	fallback_temperature = float(gen_kwargs.get("temperature", 1.0))
	gen_ids = self._generate_interleaved(
	base_embeds=inputs_embeds,
	attention_mask=attention_mask,
	max_new_tokens=max_new_tokens,
	do_sample=fallback_do_sample,
	temperature=fallback_temperature,
	)
	if return_thoughts:
	return gen_ids, states
	return gen_ids

	if latent_active:
	initial_state, prefix_past_kv, prefix_attn = self._compute_prefix_state(
	inputs_embeds,
	start_ids,
	strict=True,
	language=language_hint or "English",
	)
	deltas, states, stats = self._compute_continuous_thoughts(
	initial_state,
	prefix_past_key_values=prefix_past_kv,
	prefix_attention_mask=prefix_attn,
	halt_threshold=dynamic_halt_threshold,
	)

	K = deltas.size(1)
	if K == 0:
	# N=0 skip: remove ALL NT tokens, completely reverting to base generation.
	remove_count = self.n_latent
	start_ids = start_ids[:, :-remove_count]
	inputs_embeds = inputs_embeds[:, :-remove_count, :]
	attention_mask = attention_mask[:, :-remove_count]
	elif K < self.n_latent:
	# N > 0 but early halted: Remove the unneeded NT tokens instead of zero-padding them.
	# This prevents the LLM from attending to "blank" un-updated base_nt embeddings.
	remove_count = self.n_latent - K
	start_ids = start_ids[:, :-remove_count]
	inputs_embeds = inputs_embeds[:, :-remove_count, :]
	attention_mask = attention_mask[:, :-remove_count]
	inputs_embeds = self._inject_latent_deltas(
	inputs_embeds, start_ids, deltas
	)
	else:
	inputs_embeds = self._inject_latent_deltas(
	inputs_embeds, start_ids, deltas
	)
	elif soft_prompt_active:
	nt_mask = (start_ids == self.nt_token_id)
	nt_counts = nt_mask.sum(dim=1)
	assert (nt_counts == self.n_latent).all(), f"Gen NT count mismatch! {nt_counts}"
	prompt = self.soft_prompt_embed.to(dtype=inputs_embeds.dtype)
	batch_idx, seq_idx = torch.nonzero(nt_mask, as_tuple=True)
	inputs_embeds = inputs_embeds.clone()
	inputs_embeds[batch_idx, seq_idx, :] = prompt.view(-1, self._hidden_size)

	if front_prompt_active:
	is_nt = (start_ids == self.nt_token_id)
	nt_counts = is_nt.sum(dim=1)
	assert (nt_counts <= self.n_latent).all(), f"Gen NT count mismatch! {nt_counts}"

	fallback_eos = gen_kwargs.get("eos_token_id", None)
	fallback_do_sample = bool(gen_kwargs.get("do_sample", False))
	fallback_temperature = float(gen_kwargs.get("temperature", 1.0))

	try:
	if latent_active and K == 0:
	# N=0 skip -> we stripped the NT tokens from start_ids.
	# Do NOT pass inputs_embeds so Qwen generates organically from text + audio ids.
	# Must provide input_features so audio encoder handles the raw audio.
	gen_out = self.thinker.generate(
	input_ids=start_ids,
	attention_mask=attention_mask,
	input_features=input_features,
	feature_attention_mask=feature_attention_mask,
	pad_token_id=pad_token_id,
	max_new_tokens=max_new_tokens,
	**gen_kwargs,
	)
	else:
	gen_out = self.thinker.generate(
	inputs_embeds=inputs_embeds,
	input_ids=start_ids,
	attention_mask=attention_mask,
	input_features=None,
	pad_token_id=pad_token_id,
	max_new_tokens=max_new_tokens,
	**gen_kwargs,
	)
	start_len = start_ids.size(1)
	gen_ids = gen_out[:, start_len:]
	except ValueError as e:
	msg = str(e)
	if "inputs_embeds" not in msg:
	raise
	if not self._warned_generate_fallback:
	print(
	"[warn] thinker.generate does not support inputs_embeds; "
	"using fallback decoding loop."
	)
	self._warned_generate_fallback = True
	gen_ids = self._generate_with_forward_fallback(
	start_ids=start_ids,
	inputs_embeds=inputs_embeds,
	attention_mask=attention_mask,
	max_new_tokens=max_new_tokens,
	eos_token_id=fallback_eos,
	do_sample=fallback_do_sample,
	temperature=fallback_temperature,
	)
	if return_stats:
	if return_thoughts and latent_active:
	return gen_ids, states, stats
	return gen_ids, stats

	if return_thoughts and latent_active:
	return gen_ids, states
	return gen_ids