From 9296033cecd806342ef49b25436138f4e122ed68 Mon Sep 17 00:00:00 2001
From: Junst <arsol970812@gmail.com>
Date: Mon, 22 Sep 2025 13:22:18 +0900
Subject: [PATCH] Fix spatial dimension processing and similarity computation
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

🔧 **Critical Fix for Horizontal Line Issue**:

1. **Improved Attention Layer Selection**:
   - Prefer outputs with proper spatial structure [B, H, W, C]
   - Higher score for square spatial format
   - Better scoring for focused attention (smaller seq_len)

2. **Fixed Similarity Computation**:
   - Direct similarity computation between concept embedding and attention features
   - Proper dimension handling (embedding_dim vs attention_dim)
   - Correct spatial reshaping from [1, seq_len] to [1, spatial_h, spatial_w]

3. **Enhanced Spatial Processing**:
   - Better handling of [1, 1024, 9216] attention output
   - Proper reshaping to [1, 32, 32] spatial format
   - Maintain spatial structure throughout computation

The horizontal line issue should now be resolved as we're properly
maintaining spatial dimensions and using direct similarity computation
instead of complex matrix operations that lose spatial structure.
---
 concept_attention.py | 56 ++++++++++++++++++++++++++++++++++----------
 1 file changed, 44 insertions(+), 12 deletions(-)

diff --git a/concept_attention.py b/concept_attention.py
index 92cbefc..bd547a2 100644
--- a/concept_attention.py
+++ b/concept_attention.py
@@ -630,9 +630,21 @@ class ConceptAttentionProcessor:
                 elif 'norm' in key.lower():
                     score += 10   # Norm layers are less preferred
                 
-                # Prefer larger outputs (more spatial information)
+                # Prefer outputs with proper spatial structure
                 if hasattr(output, 'shape') and len(output.shape) >= 3:
-                    score += output.shape[1] * output.shape[2]  # Spatial dimensions
+                    # Check if it's already in spatial format [B, H, W, C] or [B, C, H, W]
+                    if len(output.shape) == 4:
+                        if output.shape[1] == output.shape[2]:  # Square spatial format
+                            score += 200  # Much higher score for proper spatial format
+                        else:
+                            score += output.shape[1] * output.shape[2]  # Spatial dimensions
+                    else:
+                        # For [B, seq_len, dim] format, prefer smaller seq_len (more focused)
+                        seq_len = output.shape[1]
+                        if seq_len <= 1024:  # Prefer smaller, more focused attention
+                            score += 1000 // seq_len  # Higher score for smaller seq_len
+                        else:
+                            score += 50  # Lower score for very large seq_len
                 
                 if score > best_score:
                     best_score = score
@@ -697,19 +709,39 @@ class ConceptAttentionProcessor:
                 
                 # Compute attention-weighted concept relevance
                 # This follows the original ConceptAttention paper approach
-                # Project concept embedding to attention space
-                concept_projection = torch.matmul(embedding_norm, attention_norm.transpose(-1, -2))
                 
-                # Ensure we maintain spatial dimensions properly
-                if len(concept_projection.shape) == 3:
-                    # concept_projection is [batch, seq_len, attention_dim]
-                    # We want to sum over the attention dimension but keep spatial structure
-                    similarity = torch.sum(concept_projection, dim=-1)  # [batch, seq_len]
+                # Method 1: Direct similarity computation
+                # Reshape attention to [batch, seq_len, dim] for easier computation
+                attention_flat = attention_spatial.view(batch_size, seq_len, dim)
+                
+                # Compute cosine similarity between concept embedding and attention features
+                # embedding_norm: [1, 1, embedding_dim]
+                # attention_flat: [1, seq_len, dim]
+                
+                # Project embedding to attention dimension if needed
+                if embedding_dim != dim:
+                    if embedding_dim < dim:
+                        # Pad embedding
+                        padding = torch.zeros(1, 1, dim - embedding_dim, device=embedding.device, dtype=embedding.dtype)
+                        embedding_expanded = torch.cat([embedding, padding], dim=-1)
+                    else:
+                        # Truncate embedding
+                        embedding_expanded = embedding[:, :, :dim]
                 else:
-                    # Fallback: take mean over last dimension
-                    similarity = torch.mean(concept_projection, dim=-1)
+                    embedding_expanded = embedding
                 
-                logger.info(f"Concept attention similarity shape: {similarity.shape}")
+                # Compute similarity: [1, seq_len]
+                similarity = torch.sum(embedding_expanded * attention_flat, dim=-1)
+                
+                # Reshape to spatial format [1, spatial_h, spatial_w]
+                spatial_size = int(np.sqrt(seq_len))
+                if spatial_size * spatial_size == seq_len:
+                    similarity = similarity.view(1, spatial_size, spatial_size)
+                    logger.info(f"Reshaped similarity to spatial format: {similarity.shape}")
+                else:
+                    logger.warning(f"Cannot reshape similarity {similarity.shape} to square format")
+                
+                logger.info(f"Final similarity shape: {similarity.shape}")
                 
                 # Get target image dimensions
                 h, w = image_shape[1], image_shape[2]