feat(gui): Add Attention Maps visualization tool#653
Conversation
- Implement `AttentionMapsTool` for visualizing transformer self-attention mechanisms. - Provide interactive inputs for Queries (Q), Keys (K), and Values (V) using `egui::Grid`. - Visualize resulting Attention Weights and final Output matrices using custom egui heatmaps. - Integrate tool into the existing `AiTab` via the `AiTool` strategy pattern. - Update `todo_gui.md` to mark the feature as complete. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
There was a problem hiding this comment.
Pull request overview
Adds an “Attention Maps” interactive visualization tool to the AI tab in math_explorer_gui, completing the corresponding Transformers roadmap item and leveraging the existing scaled_dot_product_attention implementation in math_explorer.
Changes:
- Mark “Attention Maps” as implemented in
todo_gui.md. - Register a new
AttentionMapsToolunder the AI tab tool selector. - Introduce
attention_maps.rsimplementing an interactive Q/K/V editor with heatmap rendering for attention weights and output.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| todo_gui.md | Marks “Attention Maps” feature as completed in the GUI roadmap. |
| math_explorer_gui/src/tabs/ai/mod.rs | Exposes and registers the new Attention Maps tool in the AI tab. |
| math_explorer_gui/src/tabs/ai/attention_maps.rs | Implements the new tool UI and hooks it up to scaled_dot_product_attention. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| egui::ScrollArea::vertical().show(ui, |ui| { | ||
| ui.horizontal(|ui| { | ||
| ui.vertical(|ui| { | ||
| let mut changed = false; | ||
| ui.group(|ui| { | ||
| changed |= Self::draw_matrix_input(ui, "Queries (Q)", &mut self.q_matrix); | ||
| }); | ||
| ui.group(|ui| { | ||
| changed |= Self::draw_matrix_input(ui, "Keys (K)", &mut self.k_matrix); | ||
| }); | ||
| ui.group(|ui| { | ||
| changed |= Self::draw_matrix_input(ui, "Values (V)", &mut self.v_matrix); | ||
| }); | ||
|
|
||
| if changed { | ||
| self.recalculate(); | ||
| } | ||
| }); | ||
|
|
||
| ui.separator(); | ||
|
|
||
| ui.vertical(|ui| { | ||
| ui.group(|ui| { | ||
| Self::draw_heatmap(ui, "Attention Weights (softmax(Q * K^T / sqrt(d_k)))", &self.attention_weights); | ||
| }); | ||
|
|
||
| ui.add_space(20.0); | ||
|
|
||
| ui.group(|ui| { | ||
| Self::draw_heatmap(ui, "Output (Weights * V)", &self.output_matrix); | ||
| }); | ||
| }); | ||
| }); | ||
| }); |
This PR adds an "Attention Maps" interactive visualization to the AI tab of the
math_explorer_gui, fulfilling the corresponding task intodo_gui.md.It adheres to the
AiTooltrait and encapsulates its logic withinattention_maps.rs. The tool allows users to manually manipulate small (4x4) Query, Key, and Value matrices and see the resulting Attention Weights and Output update in real time via a custom-drawn heatmap. The underlying mathematical calculation leverages the existingscaled_dot_product_attentionfunction from themath_explorercore library.PR created automatically by Jules for task 17698362353464505205 started by @fderuiter