Adding time embedding for score / flow matching

[ListIndexOutOfRange]

Hey ! Thank you for your great work ! The code is of great quality, I truly appreciate that !

I'd like to use this arch as a U-Net like model for point cloud generation using score or flow matching. Hence, I'd like to add a time embedding. There are many ways to do so, and I wanted to have your opinion.

Currently, I'm doing it inside the pooling layers. It looks like that, inside the SerializedPooling forward method:

h = self.proj(point.feat)[indices]
if 'time_embedding' in point and self.time_proj is not None:
    time_embedding_proj = self.time_proj(point.time_embedding)[indices]
    h = h + time_embedding_proj
    point_dict = Dict(
            # instead of:
            # feat=torch_scatter.segment_csr(self.proj(point.feat)[indices], idx_ptr, reduce=self.reduce),
            feat=torch_scatter.segment_csr(h, idx_ptr, reduce=self.reduce),
            ...

where time_proj is a small sequential:

self.proj = nn.Linear(in_channels, out_channels)
self.time_proj = nn.Linear(time_embedding_dim, out_channels)

and time_embedding is the output of a small MLP passed through before the actual encoder, defined in the PTV3 init as:

self.embed_time = nn.Sequential(
    nn.Linear(enc_channels[0], time_embedding_dim),
    act_layer(),
    nn.Linear(time_embedding_dim, time_embedding_dim),
)

Does it look good to you ? I'd really appreciate your feedback.

Thanks again !

[Gofinge]

Hi, sorry for the late response. I was super busy in the past half year. Back to your question, I think injecting time embedding directly before pooling, which is relatively weak in reasoning compared with attention, might cause some issues. For example, assume 2 point in the same grid have 2 different time embedding, let's assume 1, 3 as examples. However, after mean pooling, the time embedding would be 2, which is strange. Here is two potential solutions:

1. Only inject time embedding before the attention, let attention handle this thing.
2. Separately pooling for points belongs to different time stamps.

[DDXDaniel]

@Gofinge Thanks for your great work! I still have some questions about the pooling. In my opinion, different time embeddings implies that the points come from different point clouds. So how can 2 points with 2 different time embeddings are in the same grid since they actually come from different point clouds? And why pooling points without regarding whether they are from the same point clouds in your code? Thanks and Looking forward to your reply!

Hi, sorry for the late response. I was super busy in the past half year. Back to your question, I think injecting time embedding directly before pooling, which is relatively weak in reasoning compared with attention, might cause some issues. For example, assume 2 point in the same grid have 2 different time embedding, let's assume 1, 3 as examples. However, after mean pooling, the time embedding would be 2, which is strange. Here is two potential solutions:

1. Only inject time embedding before the attention, let attention handle this thing.
2. Separately pooling for points belongs to different time stamps.