Skip connection indexing in Swin3DTransformerBackbone.forward: mismatch with paper's backbone figure?

[MatthiasBey]

Hi, I am currently working on latent space data of Aurora for interpretability research at LMU. I found that spatial plots of latent states of encoding stage 1 strongly resemble those of decoding stage 3, which I didn't expect.

I checked the forward method of Swin3DTransformerBackbone in swin3d.py (

aurora/aurora/model/swin3d.py

Lines 1026 to 1055 in 639b4e8

	skips = []
	# Saved contexts at higher resolutions to be reused in the backbone U-Net decoder
	saved_cs = []
	for i, layer in enumerate(self.encoder_layers):
	saved_cs.append(c)
	x, x_unscaled = layer(x, c, all_enc_res[i], rollout_step=rollout_step)
	# There should be a `context_down_layer` for every encoder layer except the last.
	if self.stochastic and i < self.num_encoder_layers - 1:
	c = self.context_down_layers[i](c, all_enc_res[i])
	skips.append(x_unscaled)
	for i, layer in enumerate(self.decoder_layers):
	index = self.num_decoder_layers - i - 1
	x, _ = layer(
	x,
	c,
	all_enc_res[index],
	padded_outs[index - 1],
	rollout_step=rollout_step,
	)
	if 0 < i < self.num_decoder_layers - 1:
	# For the intermediate stages, we use additive skip connections.
	x = x + skips[index - 1]
	elif i == self.num_decoder_layers - 1:
	# For the last stage, we perform concatentation like in Pangu.
	x = torch.cat([x, skips[0]], dim=-1)
	if self.stochastic:
	c = saved_cs[index - 1]
	return x

) and to me it seems the skip connections are different than shown in the backbone figure of the paper. I am running version v1.8.0, in which this decoder loop logic is the same as shown here.

Tracing the decoder loop (for i, layer in enumerate(self.decoder_layers), index = self.num_decoder_layers - i - 1):

• i=0 (index=2): No skip connection is applied at all. The 0 < i condition in the additive-skip check excludes this case, and i == num_decoder_layers - 1 doesn't match either.
• i=1 (index=1): skips[0] (output of encoding stage 1) is added additively to the input of decoding stage 3.
• i=2 (index=0): skips[0] is concatenated again to the output of decoding stage 3.

So skips[0] is used twice: once additively before decoding stage 3, and once via concatenation after decoding stage 3. Meanwhile skips[1] (encoding stage 2's output) is not used at all.

Is this intentional, or should skips[1] be used somewhere in the decoder as well? I might be misreading the loop, so let me know if I'm missing something.

[ishan-1010]

I read through swin3d.py at the referenced commit to check this independently, and the trace holds. With the three-layer encoder and decoder configuration used in the reproduction above, the additive skip branch at i=1 (index=1, line 1048) evaluates skips[index - 1], which resolves to skips[0]. So skips[0] gets used twice, once there and once via the final concatenation on line 1051, and skips[1] is never referenced anywhere in the decoder loop.

The project documentation describes this backbone as "a 3D Perceiver encoder-decoder with a 3D Swin Transformer U-Net backbone." The standard U-Net convention pairs decoder stage k with encoder stage k at the same resolution, which would be skips[index], not skips[index - 1]. That reads like an off-by-one rather than an intentional asymmetry, though I have only seen the architecture description and not the paper's actual figure, so I cannot rule out the offset being deliberate for a reason not visible from the code alone.

If it turns out to be a real bug, I am happy to put together a PR that swaps the index and adds a test confirming the fix actually changes decoder output. Right now, since skips[1] is unused, corrupting its values would not change anything, which is itself a useful way to verify current behavior before touching it. Let me know if this reading is right.

[wesselb]

Hey Matthias Beylich (@MatthiasBey) and Ishan Katoch (@ishan-1010), thanks for opening an issue! Very sorry for the slow response from our side on this issue.

That skips[0] is used twice is intentional. The last stage of the U-Net decoder doesn't split the patches and should be viewed as a final processing step, so there is no dimension mismatch. Roughly this is the mental image to have in mind in the case that there are four stages (apologies for my hastily put together ASCII diagram):

                ┌───────────────────────────────────────┐             
                │                                       ▼             
stage 0 ────────┼────────────────►┌───────►stage 3  ─────────►        
                ▼                 │                                   
             stage 1            stage 2                           
                │                 ▲                                   
                ├────────────────►│                                   
                ▼                 │                                   
             stage 2            stage 1                               
                │                 ▲                                   
                │                 │                                   
                ▼                 │                                   
             stage 3            stage 0                               
                │                 ▲                                   
                └─────────────────┘                                   

See also Figure B1 in the supplement of the paper.

Now do note that there is no skip connection in the bottom of the U. This is a consequence of the 0 < i in the first if-statement. One could arguably remove this condition to fully connect the U.

If it turns out to be a real bug, I am happy to put together a PR that swaps the index and adds a test confirming the fix actually changes decoder output.

I appreciate the willingness to put together a PR, but we should be careful not to break the existing checkpoints.

Hopefully this clarifies!