Rolling along all Dimension

[leanderloew]

Hey, first of all thanks for this implementation! I noticed a small issue:

shifted_x = torch.roll(x, shifts=(-shift_pl, -shift_lat, -shift_lon), dims=(1, 2, 3))

We are rolling along all dimensions, depending on the exact window / shift configuration without padding. Thus it can happen that we "wrap" around not only on longitudes (where it makes sense), but also on latitudes or on pressure levels (where it doesn't make sense). We can fix this by guaranteeing a padding of at least shift size for lat/pl dims.

[robert-DL]

Hi Leander,

thanks for bringing that up!

regarding pressure levels, it isn't an issue as pressure levels are sorted ascending, i.e. from low pressure level (higher in the atmosphere) to high pressure levels (lower in the atmosphere). By shifting -1, the surface level is close to 1000 hpa, which is fine. For the latitude I am not completely sure, I try to figure that out.

EDIT:

geoarches/geoarches/backbones/archesweather_layers.py

Lines 438 to 472 in 9c2fc25

	x = self.pad(x.permute(0, 4, 1, 2, 3)).permute(0, 2, 3, 4, 1)
	_, Pl_pad, Lat_pad, Lon_pad, _ = x.shape
	shift_pl, shift_lat, shift_lon = self.shift_size
	if self.roll:
	shifted_x = torch.roll(x, shifts=(-shift_pl, -shift_lat, -shift_lon), dims=(1, 2, 3))
	x_windows = window_partition(shifted_x, self.window_size)
	# B*num_lon, num_pl*num_lat, win_pl, win_lat, win_lon, C
	else:
	shifted_x = x
	x_windows = window_partition(shifted_x, self.window_size)
	# B*num_lon, num_pl*num_lat, win_pl, win_lat, win_lon, C
	win_pl, win_lat, win_lon = self.window_size
	x_windows = x_windows.view(
	x_windows.shape[0], x_windows.shape[1], win_pl * win_lat * win_lon, C
	)
	# B*num_lon, num_pl*num_lat, win_pl*win_lat*win_lon, C
	attn_windows = self.attn(
	x_windows, mask=self.attn_mask
	) # B*num_lon, num_pl*num_lat, win_pl*win_lat*win_lon, C
	attn_windows = attn_windows.view(
	attn_windows.shape[0], attn_windows.shape[1], win_pl, win_lat, win_lon, C
	)
	if self.roll:
	shifted_x = window_reverse(attn_windows, self.window_size, Pl_pad, Lat_pad, Lon_pad)
	# B * Pl * Lat * Lon * C
	x = torch.roll(shifted_x, shifts=(shift_pl, shift_lat, shift_lon), dims=(1, 2, 3))
	else:
	shifted_x = window_reverse(attn_windows, self.window_size, Pl_pad, Lat_pad, Lon_pad)
	x = shifted_x

Padding is applied here, but the padding values are zero if I am not mistaken, as the window size divides the input resolutions without a remainder (see init method of the EarthSpecificBlock). I will further look into that.

[gcouairon]

Hello,
Indeed the high latitudes tokens can attend to those at lower lattitudes and similarly for the high pressure levels to the low pressure levels. but since those are pretty uncorrelated, the model will just learn to ignore irrelevant tokens without needing to pass in an explicit attention mask. I don't expect the model to perform better with an explicit attention mask