Question about fixbe=False, blockwise hybrid attention, and commented token_levels path in I2IA_mi

[Lcyanstars]

Hi, thanks for releasing the code and for the detailed clarification in #6.

I am trying to understand the intended implementation of the I2IA/JD3P training attention mask, especially the relationship between:

• --fixbe False
• Emu3SFTDatasetI2IA_mi
• --use_blockwise_attn_mask True
• the commented-out token_levels / build_blockwise_causal_attn_mask path

According to the README, --use_blockwise_attn_mask True should use blockwise hybrid attention:

• bidirectional within the future-image block
• bidirectional within the action block
• causal across blocks

However, after reading the current code, I am unsure whether the current implementation actually matches this intended behavior.

Code path I inspected

The CALVIN I2IA training script uses:

--with_i_ia True
--mask_image True
--actions True
--actions_format fast
--use_blockwise_attn_mask True

This selects Emu3SFTDatasetI2IA_mi in train/datasets.py.

In Emu3SFTDatasetI2IA_mi, there appears to be an intended token_levels implementation, but it is currently commented out:

# token_levels = torch.full((self.tokenizer.model_max_length,), fill_value=-1, dtype=torch.long)
# ...
# token_levels[cur_idx:cur_idx+len(masked_image_input_ids)] = cur_level
# ...
# token_levels[cur_idx:action_end] = cur_level
# ...
# "token_levels": token_levels,

In reference/Emu3/emu3/mllm/modeling_emu3.py, the use_blockwise_attn_mask branch currently calls:

attention_mask = build_blockwise_attn_mask(
    input_ids,
    boi_token_id=151852,
    eoi_token_id=151853,
    pad_token_id=151643,
    boa_token_id=151844,
    eoa_token_id=151845,
    include_boundary_as_image=True,
    return_bool_mask=False,
    past_key_values_length=past_key_values_length,
).to(dtype=inputs_embeds.dtype)

while the more explicit token_levels implementation is also commented out:

# attention_mask = build_blockwise_causal_attn_mask(
#     token_levels,
#     past_key_values_length=past_key_values_length,
# ).to(dtype=inputs_embeds.dtype)

Concern

build_blockwise_attn_mask(input_ids) treats both image and action boundaries as one generic region:

is_boi = (sequence == boi_token_id) | (sequence == boa_token_id)
is_eoi = (sequence == eoi_token_id) | (sequence == eoa_token_id)
...
is_image = (level > 0) | is_boi | is_eoi

Since Emu3SFTDatasetI2IA_mi appends the condition image, masked future image, and masked action contiguously, all image/action tokens can become one continuous non-text region. In that case, the future-image block and the action block may be merged into the same bidirectional block, rather than being separated by causal block boundaries.

I tested this with a minimal synthetic sequence:

import torch
from reference.Emu3.emu3.mllm.modeling_emu3 import (
    build_blockwise_attn_mask,
    build_blockwise_causal_attn_mask,
)

BOI=151852; EOI=151853; EOF=151847; BOA=151844; EOA=151845

# text, condition image, future image, action
seq = torch.tensor([[
    10, 11,
    BOI, 101, EOF, EOI,
    BOI, 201, EOF, EOI,
    BOA, 401, 402, EOA,
]], dtype=torch.long)

# intended explicit blocks: text=0, condition image=1, future image=2, action=3
levels = torch.tensor([[
    0, 0,
    1, 1, 1, 1,
    2, 2, 2, 2,
    3, 3, 3, 3,
]], dtype=torch.long)

q_fut_img = 7
k_action = 11
q_action = 11
k_fut_img = 7

m1 = build_blockwise_attn_mask(seq, return_bool_mask=True)[0, 0]
m2 = build_blockwise_causal_attn_mask(levels)[0, 0]

print("current build_blockwise_attn_mask fut_img->action allowed:", bool(m1[q_fut_img, k_action]))
print("current build_blockwise_attn_mask action->fut_img allowed:", bool(m1[q_action, k_fut_img]))
print("token_levels mask fut_img->action allowed:", bool(m2[q_fut_img, k_action] == 0))
print("token_levels mask action->fut_img allowed:", bool(m2[q_action, k_fut_img] == 0))

The result is:

current build_blockwise_attn_mask fut_img->action allowed: True
current build_blockwise_attn_mask action->fut_img allowed: True
token_levels mask fut_img->action allowed: False
token_levels mask action->fut_img allowed: True

This makes me wonder whether the current build_blockwise_attn_mask(input_ids) path is intended, or whether the commented token_levels path was the intended implementation for the
paper/README description.

Question about fixbe=False

In #6, you mentioned that the correct training setting is self.fixbe=False, and that masking BOI/EOI can make the attention pattern switch from blockwise attention to causal attention, resulting in a mixture of causal and blockwise attention.

I would like to clarify this point. In the current implementation, because block boundaries are inferred from the masked input_ids, masking BOI/EOI/BOA/EOA seems able to create several different malformed block parses, not only a clean switch to causal attention. For example:

• masking BOI can make later image/action tokens be treated as text-like causal tokens until an EOI appears
• masking EOI can make an image/action region extend further than intended
• masking BOA/EOA can change action block parsing
• if image/action boundaries are not separated explicitly, future image and action may already be in the same bidirectional block even with fixed boundaries

Could you clarify the intended behavior?

Specific questions

1. For I2IA/JD3P training, should future-image tokens be prevented from attending to action tokens, with action tokens allowed to attend to future-image tokens?
2. Is the current build_blockwise_attn_mask(input_ids) implementation expected to reproduce the paper results, or should the commented token_levels + build_blockwise_causal_attn_mask path be enabled?
3. If fixbe=False is intentionally used to mix causal and blockwise attention, how should BOI/EOI/BOA/EOA masking avoid unintended block merging or leakage between future-image and action tokens?
4. Would a patch that returns token_levels from Emu3SFTDatasetI2IA_mi and uses build_blockwise_causal_attn_mask(token_levels) be consistent with the intended blockwise hybrid attention described in the README?

Thanks again for the clarification and for maintaining the repository.