AHAMask: Task Specification in Audio LLMs

• AHAMask is a method that defines binary attention-head masks as task specifiers in large audio language models, bypassing natural language instructions.
• It employs a lightweight mechanism with fewer than 2K additional parameters to activate task-dependent functional pathways in the decoder.
• Empirical results show AHAMask achieves comparable or improved performance on single and composite tasks by mitigating instruction sensitivity.

AHAMask is a method for reliable task specification in large audio LLMs (LALMs) that dispenses with natural-language instructions by masking selected attention heads in the decoder-only LLM backbone. In the formulation reported in "AHAMask: Reliable Task Specification for Large Audio LLMs without Instructions" (Guo et al., 1 Sep 2025), the method learns task-specific binary masks over attention heads while freezing all original LALM weights, so that inference can be performed by applying a pre-learned mask rather than prompting with instructions. The reported motivation is instruction sensitivity: different instructions of the same intention can yield drastically different outcomes in current LALMs. AHAMask is presented as both a practical task-specification mechanism and evidence that LALMs contain task-dependent "functional pathways" in their attention heads (Guo et al., 1 Sep 2025).

1. Architectural setting and notation

AHAMask is defined for LALMs built by prepending an audio encoder to a pretrained decoder-only Transformer, referred to as the LLM backbone. After the audio encoder produces a sequence of embeddings, these embeddings are tokenized and then fed into the Transformer decoder (Guo et al., 1 Sep 2025).

Within each Transformer layer, the relevant component is an $H$-head Multi-Head Attention (MHA) block. If $X \in \mathbb{R}^{l \times d}$ denotes the layer input, then for each head $i = 1 \ldots H$ the model forms $Q = X W_Q^{(i)}$, $K = X W_K^{(i)}$, and $V = X W_V^{(i)}$, each of shape $l \times d_{\text{head}}$. The per-head output is

$$\mathrm{Att}_i(Q,K,V) = \mathrm{Softmax}\!\left(\frac{QK^T}{\sqrt{d_{\text{head}}}}\right)V.$$

The standard MHA output is

$$\mathrm{MHA}(X) = \sum_{i=1}^H \mathrm{Att}_i(Q,K,V) W_O^{(i)}.$$

The method therefore operates at the granularity of attention heads rather than layers, tokens, or instructions. A plausible implication is that task control is being relocated from the prompt channel to a structural subnetwork-selection mechanism inside the decoder stack.

2. Attention-head masking formulation

The central object in AHAMask is a binary mask vector $m = [m_1,\ldots,m_H] \in \{0,1\}^H$, with one bit per head (Guo et al., 1 Sep 2025). Instead of summing over all heads, the masked attention block uses a weighted sum,

$X \in \mathbb{R}^{l \times d}$0

In practice, because the $X \in \mathbb{R}^{l \times d}$1 terms may be folded into $X \in \mathbb{R}^{l \times d}$2, the implementation is also written as

$X \in \mathbb{R}^{l \times d}$3

The intended semantics are explicit: each mask $X \in \mathbb{R}^{l \times d}$4 is associated with a task $X \in \mathbb{R}^{l \times d}$5, and applying that mask is meant to activate only those heads that form the "functional pathway" for the task. Rather than issuing an instruction such as "please transcribe" or "please classify," inference proceeds by selecting the task mask itself (Guo et al., 1 Sep 2025).

This formulation is notable for its parameter economy. The number of trainable parameters is equal to the attention head count in the LLM backbone, and the paper summarizes the approach as a lightweight, binary-mask mechanism with fewer than $X \in \mathbb{R}^{l \times d}$6K parameters in the reported settings. This suggests that the method assumes substantial task functionality is already stored in pretrained or instruction-tuned LALMs and need only be selectively gated.

3. Mask learning and training regime

AHAMask treats the binary mask as a latent parameterization learned through a real-valued logit matrix $X \in \mathbb{R}^{l \times d}$7, where $X \in \mathbb{R}^{l \times d}$8 is the number of layers. A Gumbel-Sigmoid relaxation is used during training: Gumbel noise $X \in \mathbb{R}^{l \times d}$9 is sampled, a soft mask is computed as $i = 1 \ldots H$0, a hard mask is used in the forward pass via $i = 1 \ldots H$1, and gradients are passed through $i = 1 \ldots H$2 with a straight-through estimator (Guo et al., 1 Sep 2025).

Only $i = 1 \ldots H$3 is trainable; all original LALM weights are frozen. At inference, the mask is deterministically recovered as $i = 1 \ldots H$4. The training objective is standard next-token cross-entropy loss $i = 1 \ldots H$5 over the target text for paired $i = 1 \ldots H$6 examples, with no instruction tokens appended during training.

The single-task datasets reported for mask learning are as follows.

Task	Dataset
ASR	LibriSpeech
GR	LibriSpeech clean100
SER	IEMOCAP
ASV	VoxCeleb1
AAC	AudioCaps
S2TT	CoVoST2 en→zh
OSR	Libri2Mix

Composite multi-hop tasks, including formats such as "GR|ASR" or JSON style, use LibriSpeech. For these composite settings, an instruction-following-rate penalty is added only at evaluation time, with no additional loss during training (Guo et al., 1 Sep 2025).

4. Instruction-free task specification

The operational claim of AHAMask is that task execution can be triggered without any natural-language instruction. For each downstream task $i = 1 \ldots H$7, a pre-learned mask $i = 1 \ldots H$8 is applied directly. In this scheme, the mask itself functions as the task specifier (Guo et al., 1 Sep 2025).

The paper distinguishes between single tasks and composite tasks. For single tasks, the model is expected to perform transcription, classification, captioning, translation, speaker verification, or overlap-aware speech recognition depending on the selected mask. For composite tasks, separate masks are learned for each composite format, for example $i = 1 \ldots H$9, and inference again applies exactly that mask with no text instruction. The method therefore does not claim to infer arbitrary task compositions from a symbolic rule; instead, it learns a dedicated mask per composite format.

This design directly addresses instruction sensitivity. The abstract states that current LALMs can produce drastically different outcomes for different instructions of the same intention, whereas AHAMask is proposed as a way to specify the intended task reliably without relying on instruction wording (Guo et al., 1 Sep 2025). A plausible implication is that mask selection reduces variance introduced by prompt phrasing and shifts the burden of task disambiguation from language understanding to internal routing.

5. Empirical results

The reported experiments cover single tasks on SALMONN and Qwen2Audio variants and composite tasks such as GR|ASR. The backbone sizes are stated as SALMONN with 1.6 K heads and Qwen2Audio variants with 1 K heads (Guo et al., 1 Sep 2025).

For single tasks on SALMONN, the paper reports the following comparison between natural instructions and AHAMask without instructions.

Setting	Metric	Instruction → AHAMask
ASR	WER clean\	other
GR	ACC	96.79% → 98.05%
SER	ACC	69.70% → 70.02%
AAC	METEOR/ROUGE-L	20.60/40.42 → 24.15/48.71
S2TT	BLEU-4	34.48 → 33.90
OSR	WER	23.72 → 23.89

These numbers support the paper’s claim that selective attention-head masks achieve comparable or even better performance than using instructions on single tasks (Guo et al., 1 Sep 2025). In the same section, Qwen2Audio-Instruct is reported to yield comparable WER and ACC on all seven tasks when comparing instruction-based inference with AHAMask. Random masks or swapped masks from a different model are reported to fail completely. The base-model Qwen2Audio, which has no instruction tuning, cannot perform most tasks without a mask, but with AHAMask it matches or exceeds the instructed model on GR, ASV, AAC, and OSR.

For composite tasks, the clearest quantitative example given is SALMONN on GR|ASR. Relative to instruction-based prompting, AHAMask changes Instruction Following Rate from 98.6% to 99.1%, GR ACC from 68.0% to 97.8%, and ASR WER from 3.52 to 2.21. Qwen2Audio variants are reported to show similar gains in IFR, WER, and ACC (Guo et al., 1 Sep 2025).

Taken at face value, these results indicate that the method is not merely instruction-competitive on isolated tasks but can also improve compositional task execution in settings where following the intended output format is itself part of the problem.

6. Functional pathways and interpretive analysis

A substantial part of the AHAMask contribution is interpretive. The paper states that LALMs exhibit certain "functional pathways" in their attention heads, and the masking results are used as evidence for this internal organization (Guo et al., 1 Sep 2025).

Several analyses are reported. First, mask similarity measured by Jaccard similarity correlates with task similarity: ASR and OSR have the highest overlap, while paralinguistic tasks such as GR, SER, and ASV overlap less with generative tasks. Second, when heads are sorted by learned logit scores and included by gradually relaxing a threshold, downstream metrics improve smoothly and monotonically as more high-weight heads are added. The paper interprets this as evidence that the functional pathway is an emergent, collective subnetwork rather than a single "master head."

A concrete ablation-style example is given for SER. Below a certain head index there is no meaningful output; then incorrect emotion labels appear; after passing a critical index, correct classification emerges; beyond that point, heads drift into captioning or transcription. This description situates task behavior along a continuum of head inclusion rather than an all-or-nothing switch.

The paper also reports an "all roads lead to Rome" observation: different random seeds produce distinct masks of similar size for the same task, yet these masks achieve near-identical accuracy, and their intersection often yields an even smaller necessary head set with negligible performance drop. This suggests that task-relevant computation may be distributed across multiple substitutable head subsets rather than uniquely localized.

A plausible implication is that AHAMask connects practical task control with mechanistic interpretability. However, the claims in the paper are specifically about attention-head pathways in decoder-only LALMs and should not be generalized beyond the reported models, tasks, and datasets without further evidence.

7. Scope, limitations, and relation to instruction-based control

AHAMask is framed as a method for task specification, not as a replacement for all forms of prompting. Its masks are learned from paired audio-text supervision for downstream tasks, and composite tasks require separate masks for each composite format rather than on-the-fly composition. In that sense, the method provides an alternative interface for invoking known task behaviors already latent in the model (Guo et al., 1 Sep 2025).

The reported failure of random masks and swapped masks from a different model underscores that the learned masks are model-specific. Likewise, the observation that base-model Qwen2Audio cannot perform most tasks without a mask, but can with AHAMask, indicates that the mask is not a generic pruning heuristic; it is a task-conditioned control signal tied to a particular backbone and training regime. This suggests that transferability across architectures is limited in the present formulation.

Within its reported scope, AHAMask contributes three main claims. First, decoder-only LALMs internally store task pathways in their attention heads. Second, these pathways can be selectively activated via a lightweight binary mask with the number of trainable parameters equal to the attention head count in the backbone. Third, the resulting instruction-free inference matches or exceeds natural-language instructions on a broad suite of single and composite audio tasks (Guo et al., 1 Sep 2025).