support returning attention weights from transformer forward pass

[TahaZahid05]

Description

Context & Goals

Exposing attention weights from transformer models is crucial for interpretability, diagnostic audits, and verifying model alignment (e.g. visualizing the diagonal reverse-mapping learned in the string-reversal task).

This PR adds the optional capability to return the self-attention weights from the forward pass of the model, blocks, and attention layers, while maintaining maximum performance optimizations during standard training.

Key Changes

1. Attention Layer Update (trainite/models/transformer.py):
   • Added an output_attentions boolean flag to the Attention.forward signature.
   • When output_attentions=True, it bypasses nn.functional.scaled_dot_product_attention (which does not expose internal weights) and manually calculates the attention weights matrix via Query-Key dot products and Softmax, applying both causal and padding masks.
   • Returns a 3-tuple: (output, context, attention_weights).
2. SDPA Unification & Refactoring:
   • Cleaned up the scaled_dot_product_attention invocation code in Attention.forward by dynamically setting the mask and is_causal variables first. This avoids duplicate code lines and maintains hardware-optimized FlashAttention execution when output_attentions is False.
3. Block & Model Layer Propagation:
   • Propagated the output_attentions flag down from TransformerModel through TransformerBlock to Attention.
   • TransformerModel returns a list of attention weight tensors (one per layer/block) along with the final model logits when output_attentions=True.
4. Test Suite Alignment (tests/models/transformer_test.py):
   • Updated the attention test assertions to unpack the newly introduced 3-value tuple returned by Attention.forward.

---

[aaishwarymishra]

@TahaZahid05 I am not sure if its necessary we are doing extra calculation , outside the method which is specialized for this task.

[TahaZahid05]

@aaishwarymishra how exactly is this specialized? It is just giving attention matrix which can be used for any task?

[aaishwarymishra]

My concern is we are doing extra calculation for attention, for which we have a specialized method ig

[TahaZahid05]

@aaishwarymishra SPDA does not return attention matrix. the manual calculation is necessary to support return of attention matrix?

[aaishwarymishra]

Is it necessary though? Like we actually need it? If yes you can check if some other pytorch method returns it ig.

[TahaZahid05]

@aaishwarymishra i have checked and i don't think there is any other function that returns attention matrix, we need to do it manually.

[aaishwarymishra]

Why do we need it though

[TahaZahid05]

Well, after training if user wants to visualize model's attention matrix on inputs, how would they do so without an option for attention matrix to be returned in the code?

[aaishwarymishra]

They can add it themselves ig, I am not sure about this if we do manual extra calculation then what's the point of using spda

[TahaZahid05]

manual calculation is only happening when output_attentions is True. for general training it stays False. I agree they can add it themselves, my intent for adding this in trainite was actually support for debug mode meaning in future if we provide the user with a debug mode and pre-built component of visualizing attention matrix then we would need this.

[aaishwarymishra]

hmm lets wait for the meeting ig.

[aaishwarymishra]

@TahaZahid05
hi, I have a proposal https://docs.pytorch.org/docs/2.12/generated/torch.nn.functional.scaled_dot_product_attention.html gives the implementation of what we can do is just copy paste this in attention layer, the advantage of using native attention if for speed right, we can have a simple if in the method that checks, if we are training the model we will simply route to the pytorch implementation, during inference we can use manual method.

[TahaZahid05]

Converting this to a draft for now. We can work on this if our direction of debug mode aligns with having pre-made templates of debugging like visualizing attention matrix to add to Trainite.

Would open another PR to megre attention matrix code for our string-reversal example