ARCHITECTURE.md

# How ROMA Works — Architecture & Implementation

This document explains ROMA (["Real-time Omni-Multimodal Assistant"](https://arxiv.org/abs/2601.10323))
in plain terms: what problem it solves, the model design, and **where each piece lives in this
codebase**. It is written from the actual code, with clickable pointers to the relevant files.

---

## 1. The problem in one paragraph

A normal video-LLM is **reactive**: you give it a whole video + a question, it answers once.
ROMA is **streaming and proactive**: it watches audio + video as they arrive, *second by second*,
and decides **on its own when to speak** — to fire an alert the moment a condition is met, to
narrate an event right when it finishes, or to answer a spoken question at the right time. Two
hard sub-problems fall out of this:

1. **Granularity mismatch** — audio is *dense* (a continuous waveform), video is *sparse*
   (a few discrete frames per second). They must be fused on one shared timeline.
2. **When to speak** — the model must continuously judge "should I respond *now*?" without a user
   pressing enter. ROMA's answer is a tiny extra classifier called the **Speak Head**.

---

## 2. The base model: Qwen2.5-Omni-7B

ROMA is **not trained from scratch**. It starts from **Qwen2.5-Omni-7B**, an omni-modal model with
two cooperating sub-models (you can see both referenced in the merge tooling at
[scripts/merger_new_module.py](scripts/merger_new_module.py#L23-L31)):

```
Qwen2_5OmniModel
├── thinker   (Qwen2_5OmniThinkerForConditionalGeneration)
│              understands audio + video + text, generates TEXT
│              ← ROMA adds the "Speak Head" here
└── talker    (turns the thinker's output into SPEECH tokens; uses spk_dict.pt voices)
```

- The **thinker** is the multimodal brain: it ingests interleaved audio/video/text tokens and
  produces text. ROMA's new modules are bolted onto the thinker.
- The **talker** produces audio so ROMA can *speak* its responses. ROMA leaves it essentially as-is
  (the `spk_dict.pt` speaker dictionary is carried along during merging,
  [merger_new_module.py:73-79](scripts/merger_new_module.py#L73-L79)).

> Two repos, one model. The **model internals** (the Qwen2.5-Omni classes, the Speak-Head forward
> pass, the interleaved-RoPE position function) live in a **custom `transformers` fork**,
> `git+https://github.com/Eureka-Maggie/transformers.git@roma_patch` (pinned in
> [requirements.txt:205](requirements.txt#L205)). **This repo** (a fork of *LLaMA-Factory*) holds
> everything *around* the model: data formatting, the streaming chat template, training, and the
> demo/inference glue. When you see `model.thinker.gate_mixer` below, the *class* is defined in the
> fork; the *call site* is in this repo.

---

## 3. The key idea #1 — "synchronized multimodal units" (one packet per second)

Instead of feeding one giant audio blob and one giant video blob, ROMA slices the stream into
**1-second units** and interleaves the two modalities *inside each second*. This is done in the
multimodal plugin, [src/llamafactory/data/mm_plugin.py](src/llamafactory/data/mm_plugin.py#L1836-L1990).

How the alignment works:

- A common clock: `MODEL_TIME_UNITS_PER_SECOND = 25`
  ([mm_plugin.py:1859](src/llamafactory/data/mm_plugin.py#L1859)). Everything is converted to these
  units so audio and video share one timeline.
- **Audio** is encoded to ~25 tokens/second (the dense signal, downsampled by the audio encoder —
  see the length formula at [mm_plugin.py:1844-1847](src/llamafactory/data/mm_plugin.py#L1844-L1847)).
- **Video** frames (sampled at `video_fps = 2`) become a grid of tokens each; each frame's tokens
  are stamped with their real time `frame_index × video_sec_per_grid × 25`
  ([mm_plugin.py:1887-1898](src/llamafactory/data/mm_plugin.py#L1887-L1898)).
- `processor.get_chunked_index(...)` then cuts both token streams into per-second chunks
  ([mm_plugin.py:1917-1928](src/llamafactory/data/mm_plugin.py#L1917-L1928)), and each second is
  emitted as **one packet** with this exact layout
  ([mm_plugin.py:1965-1972](src/llamafactory/data/mm_plugin.py#L1965-L1972)):

```
<|vision_bos|><|audio_bos|>  [ video tokens for this second ][ audio tokens for this second ]  <|audio_eos|><|vision_eos|>
```

That single, repeated structure is the "synchronized multimodal unit". Dense audio and discrete
video for the *same* second sit side-by-side, so the model never has to guess which audio goes with
which frame. Positions are then assigned with an **interleaved RoPE** index
(`model.thinker.get_interleaved_rope_index(...)`, called at
[gradio/proactive_gradio.py:225-232](gradio/proactive_gradio.py#L225-L232)) that orders audio and
video by time rather than by modality.

---

## 4. The key idea #2 — the Speak Head ("when to speak", decoupled from "what to say")

This is ROMA's headline contribution. It is a **tiny binary classifier** attached to the thinker
that, every second, outputs a probability "should I speak now?". Critically, it is **separate from**
the thinker's normal language-model head that decides *what* words to say — that's the "decoupling
of response initiation from generation" the paper describes.

It has two parts (the structure is visible where the merge tool re-creates them,
[merger_new_module.py:85-105](scripts/merger_new_module.py#L85-L105), and where inference calls
them, [gradio/proactive_gradio.py:259-277](gradio/proactive_gradio.py#L259-L277)):

1. **`gate_mixer`** — a learnable mixer over the last few transformer layers. It holds `K` logits
   (`K = len(gate_layer_ids)`, default the **last 4 layers** `[-4, -3, -2, -1]`) and returns
   `softmax(logits)` as mixing weights. It blends the **last token's** hidden state across those
   layers:

   ```
   h_mix = Σ_k  w_k · hidden_state[layer_k][:, -1, :]        # w = softmax(gate_mixer.logits)
   ```

2. **`gate_head`** — a small MLP that maps `h_mix` to a single number (a logit). The released model
   uses the "pro" variant: `gate_head_pro_fc1 → activation → gate_head_pro_fc2`
   ([proactive_gradio.py:271-275](gradio/proactive_gradio.py#L271-L275)). A sigmoid turns the logit
   into a probability:

   ```
   p_speak = sigmoid( gate_head_pro_fc2( act( gate_head_pro_fc1( h_mix ) ) ) )
   ```

Then a simple rule fires the response
([proactive_gradio.py:283-288](gradio/proactive_gradio.py#L283-L288)):

```
if p_speak > THRESHOLD:   ->  speak (alert / narrate / answer)
else:                     ->  stay silent
```

The threshold is task-dependent: **0.6** for proactive alerts
([proactive_gradio.py:20](gradio/proactive_gradio.py#L20)) and **0.975** for narration
([narration_gradio.py](gradio/narration_gradio.py)) — narration is stricter so it only speaks at
clear event boundaries.

Why this design is nice: the gate is *lightweight* (a few thousand parameters reading existing
hidden states), so checking "should I speak?" every second is cheap, and it doesn't disturb the
thinker's generation quality.

---

## 5. Putting it together — the real-time inference loop

The clearest end-to-end implementation is the proactive demo,
[gradio/proactive_gradio.py](gradio/proactive_gradio.py#L88-L292). One full pass:

```
load Qwen2_5OmniModel (bf16, flash_attention_2)          # proactive_gradio.py:27-34
build per-second multimodal units via the template       # :122-144  (uses streaming_turn template)
past_key_values = None                                   # KV cache → makes it incremental/streaming
for each 1-second chunk:
    wait until this second actually arrives (real-time)  # :195-197  (time.sleep keeps it ~1 fps)
    slice this second's video tokens + audio mel frames  # :200-222  (audio: 100 mel frames/sec)
    compute interleaved-RoPE positions, shift by KV pos  # :225-239
    out = model.thinker(..., past_key_values, use_cache, output_hidden_states)   # :254-255
    p_speak = SpeakHead(out.hidden_states)               # :259-277  (gate_mixer + gate_head)
    past_key_values = out.past_key_values                # :279  carry the cache forward
    if p_speak > THRESHOLD: emit alert                   # :283-288
```

Two things make it *streaming* rather than batch:
- **KV cache** (`past_key_values`): each second only the *new* chunk's tokens are forwarded; the
  past is reused. Cost per step stays roughly constant instead of growing with video length.
- **Real-time pacing**: the loop sleeps so it advances ~1 second of input per wall-clock second
  ([:195-197](gradio/proactive_gradio.py#L195-L197)), mimicking a live feed.

The three demos differ only in the "what happens when the gate fires" part:
- [proactive_gradio.py](gradio/proactive_gradio.py) — fire an **alert** when a spoken condition is met.
- [narration_gradio.py](gradio/narration_gradio.py) — **narrate** the event that just ended.
- [mme_gradio.py](gradio/mme_gradio.py) — **answer** a multimodal question (reactive).

---

## 6. How it's trained

Training config: [yamls/train.yaml](yamls/train.yaml). It is **full supervised fine-tuning** of
Qwen2.5-Omni-7B with DeepSpeed ZeRO-3, FlashAttention-2, Liger kernels, bf16, the vision tower
frozen, lr `1e-5`, `max_steps: 6000`, on a **streaming** dataset (`streaming: true`,
interleaved `abl_all_1, abl_all_2`). The entry point on a multi-GPU node is
[sh/train.sh](sh/train.sh) → `launcher.py` → `run_exp`; the single-GPU debug path is
[debug_sft_singlegpu.py](debug_sft_singlegpu.py).

What the model learns comes from how the **labels** are built, in
[src/llamafactory/data/mm_plugin.py](src/llamafactory/data/mm_plugin.py#L1901-L1990) under the
`streaming_mix` template ([template.py:1630-1647](src/llamafactory/data/template.py#L1630-L1647)).
Each second gets a target:

- The dataset gives "say *this text* at time *t*" pairs; these are bucketed into
  `answers_at_second` ([mm_plugin.py:1901-1914](src/llamafactory/data/mm_plugin.py#L1901-L1914)).
- A second **with** a target → the model should *speak* that content (gate label ≈ 1).
- A second **with no** target → for proactive alerts the gold output is literally `"no"`
  (encoded in the system prompt,
  [template.py:1636-1641](src/llamafactory/data/template.py#L1636-L1641)) → gate label ≈ 0.

So one objective trains **both** behaviors at once: the **Speak Head** learns the binary
speak/stay-silent decision per second, while the **thinker's LM head** learns to produce the right
content when it *does* speak. The system prompt also encodes the task rules — narrate only at event
transitions; for alerts, output the specified text (or `"alert"`) only when the condition holds,
otherwise `"no"`. The paper's "two-stage streaming curriculum" governs the order/mix in which these
streaming examples are presented.

The new gate parameters are flagged as trainable add-ons via `additional_target: gate_head,gate_mixer`
(and the freeze-mode variant `gate_head_pro_fc1,gate_head_pro_fc2,gate_mixer`) in
[yamls/train.yaml](yamls/train.yaml#L13-L24).

---

## 7. From trained weights to the released checkpoint

After full fine-tuning you have a `thinker` that contains the new `gate_*` modules. The merge tool
[scripts/merger_new_module.py](scripts/merger_new_module.py#L152-L205) (`save_full`):

1. Loads the fine-tuned thinker (and, if needed, back-fills the `gate_head` / `gate_mixer` tensors
   straight out of the safetensors shards — [:108-149](scripts/merger_new_module.py#L108-L149)).
2. Drops it into a fresh top-level `Qwen2_5OmniModel` (`base_model.thinker = thinker`).
3. Saves the whole thing as sharded safetensors + processor, copying `spk_dict.pt` along.

That merged artifact is what you download from HuggingFace (`EurekaTian/ROMA`) and point the demos at
via `whole_model/model`.

---

## 8. Mental model / cheat-sheet

| Concept | What it is | Where in the code |
|---|---|---|
| Base model | Qwen2.5-Omni-7B (`thinker` + `talker`) | fork `transformers@roma_patch`; used in [merger_new_module.py](scripts/merger_new_module.py#L23-L31) |
| Synchronized unit | 1-second packet interleaving video + audio tokens | [mm_plugin.py:1965-1972](src/llamafactory/data/mm_plugin.py#L1965-L1972) |
| Shared clock | 25 model-time-units per second | [mm_plugin.py:1859](src/llamafactory/data/mm_plugin.py#L1859) |
| Interleaved RoPE | time-ordered positions for audio+video | call at [proactive_gradio.py:225](gradio/proactive_gradio.py#L225) |
| **Speak Head** | gate_mixer (layer blend) + gate_head (MLP→sigmoid) → p(speak) | [proactive_gradio.py:259-277](gradio/proactive_gradio.py#L259-L277), [merger_new_module.py:85-105](scripts/merger_new_module.py#L85-L105) |
| Streaming loop | KV-cache + real-time pacing, gate checked each second | [proactive_gradio.py:193-292](gradio/proactive_gradio.py#L193-L292) |
| Streaming template + labels | per-second targets; "no" when silent | [template.py:1630-1647](src/llamafactory/data/template.py#L1630-L1647), [mm_plugin.py:1901-1990](src/llamafactory/data/mm_plugin.py#L1901-L1990) |
| Training recipe | full SFT, ZeRO-3, fa2, streaming dataset | [yamls/train.yaml](yamls/train.yaml) |

> Note on accuracy: line numbers point at the code as cloned. The Speak-Head *module classes*,
> the Qwen2.5-Omni model code, and `get_interleaved_rope_index` are defined in the
> `Eureka-Maggie/transformers@roma_patch` fork, not in this repo — this repo calls into them.
> If you want to read the gate's exact `forward`/init, look in that fork's
> `modeling_qwen2_5_omni.py`.