GaussianFlow

GaussianFlow calculates the predicted flow of Gaussians between two time steps ($t_1$ and $t_2$). To accelerate training, variables related to $t_1$ (such as conic_2D) should be detached from the gradient graph using .detach(). This ensures that the backward pass only updates variables at $t_2$, as $t_1$ variables were already updated in the previous step ($t_1 - 1$).

To implement this, you need:

1. Gaussian parameters at $t_1$ (e.g., proj_2D, gs_per_pixel, weight_per_gs_pixel, x_mu, and conic_2D).
2. Gaussian parameters at $t_2$ (e.g., proj_2D, conic_2D, and conic_2D_inv).
3. A CUDA implementation for calculating Gaussian flow variables (refer to the diff-gaussian-rasterization repository for the underlying CUDA kernels).

### We detach the variables related to t_1 in calculation of GaussianFlow such that the gradient backward 
### only works for variables at t_2 while keeping variables at t_1 unchanged because
### variables at t_1 have been updated at t_1 - 1 with the same logic. 

# Gaussian parameters at t_1
proj_2D_t_1 = render_t_1["proj_2D"]
gs_per_pixel = render_t_1["gs_per_pixel"].long() 
weight_per_gs_pixel = render_t_1["weight_per_gs_pixel"]
x_mu = render_t_1["x_mu"]
cov2D_inv_t_1 = render_t_1["conic_2D"].detach()

# Gaussian parameters at t_2
proj_2D_t_2 = render_t_2["proj_2D"]
cov2D_inv_t_2 = render_t_2["conic_2D"]
cov2D_t_2 = render_t_2["conic_2D_inv"]

# ... (calculation of B_t_2 and B_inv_t_1) ...

# full formulation of GaussianFlow
cov_multi = (B_t_2_B_inv_t_1[gs_per_pixel] @ x_mu.permute(0,2,3,1).unsqueeze(-1).detach()).squeeze()
predicted_flow_by_gs = (cov_multi + proj_2D_next[gs_per_pixel] - proj_2D[gs_per_pixel].detach() - x_mu.permute(0,2,3,1).detach()) * weights.detach()

To supervise the predicted Gaussian flow, use a loss function that compares the predicted flow against a pseudo-ground truth optical flow. It is recommended to use a mask to filter out noise by only considering regions with large motion.

1. Create a large_motion_msk using a threshold (e.g., flow_thresh = 0.1) on the norm of the optical_flow.
2. Calculate the $L_2$ norm of the difference between the optical_flow and the summed predicted_flow_by_gs within the masked regions.
3. Add the resulting Lflow to your total loss, scaled by a flow_weight (e.g., 1.0 or 0.1).

# flow supervision loss 
large_motion_msk = torch.norm(optical_flow, p=2, dim=-1) >= flow_thresh  # flow_thresh = 0.1 or other value to filter out noise
Lflow = torch.norm((optical_flow - predicted_flow_by_gs.sum(0))[large_motion_msk], p=2, dim=-1).mean() 
loss = loss + flow_weight * Lflow # flow_weight could be 1, 0.1, ... whatever you want.