Clarifications on attention mask shapes in `attn_encoder.py` and use of noisy `x0_sample` in `train.py`

[yueyin19960520]

Title: Clarifications on attention mask shapes in attn_encoder.py and use of noisy x0_sample in train.py

Body:

Hi! Thanks for open-sourcing FlowER — I’ve been reading through the code and had two small questions about masking and the training target. 🙏

1) Attention mask shapes and which mask is used

File: FlowER/model/attn_encoder.py

• Around line 93, the comment says the mask has shape (batch, query_len, key_len).
• Later, around lines 154–155, scores are masked_fill’d with a mask that (per the comment) has shape (B, 1, 1, T_values) just for letting the code keep processing.
• In forward() (around lines 336 and 339), it looks like the code uses MASK (from around line 316) rather than MATRIX_MASKS (from around line 318).

This feels a bit inconsistent: the documented shape is (B, Q, K), while the applied mask seems to use a broadcast-friendly (B, 1, 1, T_values). Also, forward() appears to use MASK instead of MATRIX_MASKS.

Questions / suggestions:

• In forward(), should MATRIX_MASKS be used instead of MASK, or is the current use of MASK intentional?

Note: I realize this may not affect final outputs because padding is masked later anyway; I’m mainly looking to understand the intended convention and avoid confusion for future readers. I’m happy to open a small PR to standardize comments/names if that helps.

---

2) Why use noisy x0_sample in ut computation?

File: FlowER/train.py

• Around line 191, ut is computed as:

ut = flow.compute_conditional_vector_field(x0_sample, x1)

Questions

• What’s the rationale for using the noisy x0_sample here instead of the clean x0? Is this for regularization/noise conditioning (e.g., to stabilize training or match the objective), or to ensure an unbiased estimate under the training distribution? If there’s a relevant paper/section that motivates this choice, a pointer would be great.

[FongMunHong]

Hey, thanks for reaching out!

1. You are spot on that my dimensions are a little off, since it's an artifact since I imported the code from Graph2SMILES. Would be happy if you can create a pull request to make a change on the annotation of the dimensions!

As for the difference between mask and matrix_masks, they serve different purposes, the first is used for the atom representations, the latter for the bond-electron matrix. Since matrix_masks is derived from mask, thus, it's natural not to pass it around.

2. As for question about x0_sample, this boils down to the fundamentals of flow matching. During inference, we know that flow matching uses an ODE to integrate the model to create a conditional probability path towards the target distribution. And since it's an ODE, it is deterministic. This means the only way to create stochasticity is by injecting noise to the prior distribution x0. Thus, during training it is required for the model to see such stochastic priors in order for the inference to work. Read more about this in the method section of the paper and also in the peer review file.

[yueyin19960520]

Thanks a lot for your detailed response!

Regarding the first point — yes, I understand that the mask input corresponds to the node mask, while matrix_masks is the bond–electron mask. The main issue I noticed is in lines 152–153 of attn_encoder.py:

mask = mask.unsqueeze(1)
scores = scores.masked_fill(mask, -1e18)

I believe this should instead be:

matrix_masks = ~(~mask * ~mask.transpose(1, 2)).bool()
scores = scores.masked_fill(matrix_masks, -1e18)

This modification ensures that the masking is correctly applied to the pairwise interactions in the bond–electron matrix rather than just along node dimensions.

As for the second question — I fully understand now. Your explanation of x0_sample and its role in introducing stochasticity during training for flow matching inference was very helpful.

Thanks again for clarifying everything!

[FongMunHong]

Hi, sorry for the late reply.

As for the attn_encoder.py , I believe what was intended is a node mask, since the final output will be with respect to the atoms, not the bond electron matrix. The bond electron matrix is intended to be treated as a "positional encoding" in this case.

You can always run inference and check the dimensions. It should work as intended.
Let me know if there is any confusion.