Questions about "NeRAF audio-visual joint training improves vision performance"

[yangbang18]

Thanks for sharing. Table 2 in the paper shows that NeRAF audio-visual joint training improves vision performance. After reading the code, I have the following puzzle.

It is noted that the 3D grid in NeRAF is initialized as a tensor rather than parameters that require gradients:

def reset_grid(self,device=None):
        device = self.device if device is None else device
        self.grid = torch.zeros((7, int((self.grid_size[1] - self.grid_size[0]) / self.grid_step),
                                 int((self.grid_size[3] - self.grid_size[2]) / self.grid_step),
                                 int((self.grid_size[5] - self.grid_size[4]) / self.grid_step)),dtype=torch.float32,device=device)
        # Add coordinates
        grid_coordinates = torch.meshgrid(torch.arange(self.grid_size[0]+self.grid_step/2, self.grid_size[1], self.grid_step), torch.arange(self.grid_size[2]+self.grid_step/2, self.grid_size[3], self.grid_step), torch.arange(self.grid_size[4]+self.grid_step/2, self.grid_size[5], self.grid_step), indexing='ij')
        grid_coordinates = torch.stack(grid_coordinates, dim=0)
        self.grid[4:,:,:,:] = grid_coordinates

Then, when updating 3D grid values during training, the 3D grid is detached from the computation graph:

self.grid = self.grid.detach()

self.grid[0, xs, ys, zs] = color[:, 0].float().squeeze()
self.grid[1, xs, ys, zs] = color[:, 1].float().squeeze()
self.grid[2, xs, ys, zs] = color[:, 2].float().squeeze()
self.grid[3, xs, ys, zs] = alpha.float().squeeze()

Although such a 3D grid guides acoustic modeling, it seems that the audio loss can not be backpropagated to NeRF. So, why does NeRAF audio-visual joint training improve vision performance?

[AmandineBtto]

You're right that the 3D grid isn't a trainable parameter. It's a buffer for storing geometric and appearance information from NeRF, used to guide the acoustic field.

The reason we detach past grid values is that we only want to update the grid based on the current batch of coordinates, not backpropagate through stored values from previous steps.

We did carefully verify that the audio loss still backpropagates through the vision model.
If you want to confirm this yourself, you can try a quick test in get_train_loss_dict() (inside NeRAF_pipeline.py):
After computing both the vision and audio loss dictionaries, but before any backward pass (i.e. just before the return), insert something like:

# Check gradients before backward
for param in self.model.parameters():
    if param.grad is not None:
        print(param.grad.norm())
    else:
        print("no grad for vision model")

# Backward only audio loss
audio_loss = loss_audio_dict['audio_sc_loss'] + loss_audio_dict['audio_mag_loss']
audio_loss.backward()

# Check gradients again
for param in self.model.parameters():
    if param.grad is not None:
        print(param.grad.norm())
    else:
        print("no grad for vision model")

exit()

You'll see that before the backward pass, the vision model has no gradients. But after audio_loss.backward(), gradients appear. This confirms that the audio loss is influencing the vision model, even though the grid is detached.

Also, just to note: in get_param_groups() (inside NeRAF_model.py), we explicitly add the vision model parameters to the optimizer.

[yangbang18]

I have checked the gradients, you are right 👍

One more question about the function _generate_fixed_viewing_directions in NeRAF_Model.py:

def _generate_fixed_viewing_directions(self) -> torch.Tensor:
        phis = [math.pi / 3, 0, -math.pi]
        thetas = [k * math.pi / 3 for k in range(0, 6)]
        viewdirs = []

        for phi in phis:
            for theta in thetas:
                viewdirs.append(torch.Tensor([
                    math.cos(phi) * math.sin(theta),
                    math.cos(phi) * math.sin(theta),
                    math.sin(theta)
                ]))
        viewdirs = torch.stack(viewdirs, dim=0)
        return viewdirs

1. If phis are elevation angles, should they be [math.pi / 3, 0, -math.pi / 3] instead?
2. Should the viewdir calculation be:

[
math.sin(phi) * math.cos(theta),  # x
math.sin(phi) * math.sin(theta), # y
math.cos(phi) # z
]

[AmandineBtto]

The purpose of using multiple viewing directions is simply to obtain a view-independent averaged color. The exact directions don’t matter much for this step. For reference, I directly reused the code snippet from here: https://github.com/AIBluefisher/DReg-NeRF/blob/cd5c01beab1e754ebd002a40df2e64ccdfe2955e/conerf/register/sample_grid.py#L132

That said, after looking at the code again, I agree that the viewdir computation may not be correct (the x and y components should likely differ). However, since the method only requires that the viewing directions be different, not that they correspond to any specific orientation, this issue doesn’t affect the overall approach.