OmniTrace: A Unified Framework for Generation-Time Attribution in Omni-Modal LLMs

1. Problem Statement

Modern Omni-Modal Large Language Models (MLLMs) can generate fluent responses from interleaved inputs of text, images, audio, and video. However, a critical gap exists in attribution: identifying which specific input sources support each generated statement.

Existing attribution methods face three fundamental limitations when applied to decoder-only, open-ended generation:

1. Fixed-Target Bias: Most methods are designed for classification or extractive tasks with a fixed output objective. They fail in autoregressive generation where the output structure evolves dynamically.
2. Single-Modality Focus: Traditional approaches struggle to handle heterogeneous inputs (text, image, audio, video) unified within a single causal token timeline.
3. Token-Level Instability: Raw token-level signals (e.g., attention weights) are noisy and fragmented. They lack semantic coherence at the statement level, making it difficult to explain why a model generated a specific phrase based on a specific image region or audio segment.

The core challenge is to develop a generation-aware, omni-modal, and semantically interpretable attribution framework that operates online during decoding without retraining.

2. Methodology: OmniTrace

OmniTrace is a lightweight, model-agnostic framework that formalizes attribution as a generation-time tracing problem. It converts arbitrary token-level signals into coherent span-level explanations.

Core Algorithm (Algorithm 1)

The process operates in two main stages during autoregressive decoding:

Stage 1: Generation-Time Source Tracing

• Input: An interleaved sequence of tokens $x$ (text, image, audio, video) and a prompt.
• Token-Level Scoring: At each decoding step $t$, the model generates a token $y_t$. Simultaneously, a token-level attribution signal $a_t(i)$ is computed for every input token $i$ in the causal context. This signal can be derived from attention weights (e.g., mean-pooled, raw) or gradient-based scores.
• Source Mapping: The framework aggregates attribution mass over predefined source units $S$ (e.g., text spans, image regions, time intervals). For each generated token, it identifies the source unit $S_j$ with the highest aggregated score:
  $$\hat{s}(t) = \arg \max_{S_j \in S} \sum_{i \in S_j} a_t(i)$$
• Confidence Recording: A confidence score $c_t$ is recorded based on the strength of the attribution mass.

Stage 2: Span-Level Aggregation and Curation

• Segmentation: The generated output sequence is segmented into semantically coherent chunks (spans) $C_k$ (e.g., sentences or phrases).
• Signal Aggregation: Token-level traces within a span are pooled to reduce variance.
• Confidence-Based Source Selection: A curation module selects a concise set of supporting sources $\hat{S}_k$ for each span using:
  • POS-Aware Weighting: Prioritizing content-bearing tokens (nouns, numerals) over function words.
  • Confidence Shaping: Amplifying high-certainty alignments and suppressing weak signals.
  • Run-Level Coherence: Favoring temporally contiguous source assignments to ensure stability.
  • Threshold Filtering: Removing low-mass assignments to prevent noise.

Design Properties

• Generation-Aware: Computes attribution dynamically at each decoding step, not post-hoc on a fixed target.
• Omni-Modal: Operates on a unified token timeline spanning all modalities.
• Model-Agnostic: Compatible with various scoring mechanisms (attention, gradients) and model architectures.

3. Key Contributions

1. Formalization: Defined attribution for open-ended multimodal generation as a dynamic tracing problem, highlighting the limitations of static, single-modality approaches.
2. Framework: Introduced OmniTrace, a plug-and-play protocol that transforms noisy token-level signals into stable, span-level, cross-modal explanations without model retraining.
3. Empirical Validation: Demonstrated that generation-aware span-level attribution yields significantly more stable and interpretable results than naive baselines across diverse tasks and models.

4. Experimental Results

The authors evaluated OmniTrace on 759 examples across visual, audio, and video tasks using Qwen2.5-Omni and MiniCPM-o-4.5.

• Performance: OmniTrace significantly outperformed post-hoc baselines (Self-Attribution, Embedding-based similarity, Random).
  • Example: On Qwen2.5-Omni for Visual Summarization, OmniTrace (OTAttMean) achieved 75.66 F1, compared to 9.25 for Self-Attribution.
  • Audio/Video: Achieved strong Time-F1 scores (e.g., 83.12 for Audio QA), proving effectiveness in continuous temporal domains where embedding baselines fail.
• Ablation Studies:
  • Filtering is Critical: Removing confidence weighting or run-level coherence caused a drastic drop in Image F1 (from ~76 to ~20), indicating visual grounding is highly sensitive to noise filtering.
  • ASR Segmentation: High-quality Automatic Speech Recognition (ASR) segmentation is essential; using raw tokens without semantic segmentation led to severe performance degradation.
  • Multimodal Synergy: Combining visual and audio inputs yielded the best attribution performance, confirming the need for joint context.
• Faithfulness: In multiple-choice QA, OmniTrace showed 93.84% Top-1 consistency, meaning the attribution signals correctly identified the source corresponding to the model's chosen answer, validating that the tracing reflects the model's internal decision process.

5. Significance and Conclusion

OmniTrace addresses a critical transparency gap in the next generation of AI models. By treating attribution as a structured, generation-time problem, it provides:

• Scalability: A method that works online during decoding without heavy computational overhead or retraining.
• Interpretability: It bridges the gap between raw token probabilities and human-understandable explanations (e.g., "This sentence is based on the video segment from 14-16s").
• Trust: It enables debugging and verification of omni-modal models, ensuring that generated content is grounded in actual input evidence rather than hallucination.

The framework establishes a foundational protocol for transparency in omni-modal LLMs, suggesting that future models should integrate such tracing mechanisms natively to ensure reliability in complex, multi-evidence reasoning tasks.