Your work is excellent, could you please explain how to achieve visualization of the affinity matrix and attention masks