Backward gradient output is zero except class token in Transformer LayerNorm

yojayc: <p>I added a backward hook to the norm layer before mlp_head, but the grad_output was all zero except the class token.<br>
Can anyone explain this strange thing?</p>
<blockquote>
<p>Here is the gradients output<br>
LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
grad_output:<br>
(tensor([[[ 0.0218,  0.0294,  0.0096,  …, -0.0287,  0.0241,  0.0031],<br>
[ 0.0000,  0.0000,  0.0000,  …,  0.0000,  0.0000,  0.0000],<br>
[ 0.0000,  0.0000,  0.0000,  …,  0.0000,  0.0000,  0.0000],<br>
[ 0.0000,  0.0000,  0.0000,  …,  0.0000,  0.0000,  0.0000],<br>
[ 0.0000,  0.0000,  0.0000,  …,  0.0000,  0.0000,  0.0000],<br>
[ 0.0000,  0.0000,  0.0000,  …,  0.0000,  0.0000,  0.0000]],<br>
[[ 0.0218,  0.0294,  0.0096,  …, -0.0287,  0.0241,  0.0031],<br>
[ 0.0000,  0.0000,  0.0000,  …,  0.0000,  0.0000,  0.0000],<br>
[ 0.0000,  0.0000,  0.0000,  …,  0.0000,  0.0000,  0.0000],<br>
[ 0.0000,  0.0000,  0.0000,  …,  0.0000,  0.0000,  0.0000],<br>
[ 0.0000,  0.0000,  0.0000,  …,  0.0000,  0.0000,  0.0000],<br>
[ 0.0000,  0.0000,  0.0000,  …,  0.0000,  0.0000,  0.0000]]],<br>
device=‘cuda:0’), None, None)</p>
</blockquote>
<blockquote>
<p>Here is the backward hook function<br>
def hook_fn_backward(module, grad_input, grad_output):<br>
print(module)<br>
print(‘grad_output:\n’, grad_output[0], grad_output[0].shape)<br>
print(‘grad_input:\n’, grad_input[0], grad_input[0].shape)</p>
</blockquote>
<blockquote>
<p>Here is the network structure<br>
STAM(<br>
(to_patch_embedding): Sequential(<br>
(0): Rearrange(‘b f c (h p1) (w p2) → b f (h w) (p1 p2 c)’, p1=32, p2=32)<br>
(1): Linear(in_features=3072, out_features=512, bias=True)<br>
)<br>
(dropout): Dropout(p=0.0, inplace=False)<br>
(space_transformer): Transformer(<br>
(layers): ModuleList(<br>
(0): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(1): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(2): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(3): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(4): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(5): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(6): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(7): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(8): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(9): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(10): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(11): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
)<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
)<br>
(time_transformer): Transformer(<br>
(layers): ModuleList(<br>
(0): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(1): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(2): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(3): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(4): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
(5): ModuleList(<br>
(0): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): Attention(<br>
(to_qkv): Linear(in_features=512, out_features=1536, bias=False)<br>
(to_out): Sequential(<br>
(0): Linear(in_features=512, out_features=512, bias=True)<br>
(1): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(1): PreNorm(<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
(fn): FeedForward(<br>
(net): Sequential(<br>
(0): Linear(in_features=512, out_features=2048, bias=True)<br>
(1): GELU()<br>
(2): Dropout(p=0.0, inplace=False)<br>
(3): Linear(in_features=2048, out_features=512, bias=True)<br>
(4): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
)<br>
)<br>
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)<br>
)<br>
(mlp_head): Linear(in_features=512, out_features=100, bias=True)<br>
)</p>
</blockquote>