Skip to main content الرئيسية المنشئون adu2021 skillxiv multimodal-diffusion-alignment
multimodal-diffusion-alignment Improve text-image alignment in diffusion transformers through Temperature-Adjusted Cross-modal Attention (TACA), addressing token imbalance and timestep-dependent weighting with parameter-efficient LoRA fine-tuning.
الانتقال إلى التثبيت سوق المهارات اكتشف واستكشف مهارات الذكاء الاصطناعي التي بناها المجتمع.
التثبيت باستخدام Codex أو Claude انسخ هذا Prompt والصقه في Codex أو Claude أو مساعد آخر ليراجع صفحة Skill ويثبّتها لك.
عرض تفاصيل Prompt يتجاوز الأمر المباشر Prompt المخصّص للمراجعة. افحص المصدر قبل تشغيله.
npx skills add https://github.com/ADu2021/skillXiv --skill multimodal-diffusion-alignmentيبقى الأمر في سطر واحد. مرّر أفقيًا لمراجعته كاملًا قبل النسخ.
تفضّل نسخة محلية؟ نزّل الملفات المتاحة حاليًا لدى SkillsMP.
تحميل Zip جاري التحميل... المهن ذات الصلة SOC
استنادا إلى تصنيف SOC المهني
name multimodal-diffusion-alignment title Rethinking Cross-Modal Interaction in Multimodal Diffusion Transformers version 0.0.2 engine skillxiv-v0.0.2-claude-opus-4.6 license MIT url https://arxiv.org/abs/2506.07986 keywords ["diffusion-models","text-to-image","multimodal","cross-modal-attention","lora-finetuning"] description Improve text-image alignment in diffusion transformers through Temperature-Adjusted Cross-modal Attention (TACA), addressing token imbalance and timestep-dependent weighting with parameter-efficient LoRA fine-tuning.
Rethinking Cross-Modal Interaction in Multimodal Diffusion Transformers
Core Concept
Multimodal Diffusion Transformers (MM-DiTs) struggle with text-image alignment due to two fundamental issues: visual tokens vastly outnumber text tokens, causing text guidance to be diluted in attention softmax computation, and attention weights remain static across denoising timesteps despite varying interaction importance. Early denoising prioritizes layout establishment (requiring strong text guidance), while later steps focus on detail refinement. Temperature-Adjusted Cross-modal Attention (TACA) solves both problems through temperature scaling that rebalances modal competition and timestep-dependent weighting that adapts interaction strength across denoising phases.
Architecture Overview
Modality-Specific Temperature Scaling : Amplifies cross-modal attention logits by factor γ > 1Timestep-Dependent Adjustment : Applies temperature only during early denoising (t ≥ t_thresh)Piecewise Temperature Function : γ(t) = γ for early steps, 1.0 for detail refinementLoRA Fine-tuning : Low-rank adaptation to attention layers for artifact suppressionParameter Efficiency : <5% additional parameters, compatible with frozen base models
Implementation
Step 1: Analyze Modal Imbalance in MM-DiT
Understand the core problem before implementing TACA:
import torch
import torch.nn as nn
import torch.nn.functional as F
class ModalImbalanceAnalyzer :
"""Analyze how visual tokens suppress text guidance"""
def __init__ (self, mm_dit_model ):
self .model = mm_dit_model
def analyze_softmax_suppression (self, text_ids, image_ids, timestep ):
"""
Demonstrate how visual token abundance suppresses text attention.
In standard softmax, text token contribution is diluted by visual tokens.
"""
batch_size = text_ids.shape[0 ]
num_text_tokens = text_ids.shape[1
1
with
self
self
True
'attention_logits'
1
0
1
sum
print
f"Token counts: {num_text_tokens} text vs {num_image_tokens} image"
print
f"Text contribution to attention: {text_contribution:.2 %} "
print
f"Expected (ideal): {num_text_tokens / (num_text_tokens + num_image_tokens):.2 %} "
return
'text_logits'
'text_contribution'
def
compute_attention_imbalance_score
self, text_ids, image_ids, timestep
"""
Quantify how much text guidance is suppressed (0 = balanced, 1 = suppressed)
"""
self
1.0
1.0
50.0
'text_contribution'
max
0
return
Step 2: Implement Temperature-Adjusted Attention Create the core TACA mechanism:
class TemperatureAdjustedCrossAttention (nn.Module):
"""
Cross-modal attention with temperature scaling for modal balancing.
Addresses two issues:
1. Visual tokens suppress text via overwhelming softmax
2. Optimal text influence varies across denoising timesteps
"""
def __init__ (self, hidden_dim, num_heads, temperature_base=2.0 ):
super ().__init__()
self .hidden_dim = hidden_dim
self .num_heads = num_heads
self .head_dim = hidden_dim // num_heads
self .temperature_base = temperature_base
self .temperature = nn.Parameter(torch.tensor(temperature_base))
self .timestep_mlp = nn.Sequential(
nn.Linear(1 , 64 ),
nn.SiLU(),
nn.Linear(64 , 1 ),
nn.Sigmoid()
)
self .q_proj = nn.Linear(hidden_dim, hidden_dim)
self .k_proj = nn.Linear(hidden_dim, hidden_dim)
self .v_proj = nn.Linear(hidden_dim, hidden_dim)
self .out_proj = nn.Linear(hidden_dim, hidden_dim)
def forward (self, text_hidden, image_hidden, timestep, attention_mask=None ):
"""
Apply temperature-adjusted cross-modal attention.
Args:
text_hidden: (batch, num_text_tokens, hidden_dim)
image_hidden: (batch, num_image_tokens, hidden_dim)
timestep: scalar or (batch,) - denoising timestep
attention_mask: optional mask
Returns:
output: (batch, num_image_tokens, hidden_dim) - image features updated with text guidance
"""
batch_size = image_hidden.shape[0 ]
if isinstance (timestep, (int , float )):
timestep_tensor = torch.tensor([[timestep]], dtype=torch.float32).to(image_hidden.device)
else :
timestep_tensor = timestep.unsqueeze(1 ).float ()
temp_multiplier = self .timestep_mlp(timestep_tensor)
threshold = 0.5
uses_enhanced_temperature = (timestep_tensor > threshold).float ()
adaptive_temperature = 1.0 + (self .temperature - 1.0 ) * uses_enhanced_temperature * temp_multiplier
q = self .q_proj(image_hidden)
k = self .k_proj(text_hidden)
v = self .v_proj(text_hidden)
q = q.view(batch_size, -1 , self .num_heads, self .head_dim).transpose(1 , 2 )
k = k.view(batch_size, -1 , self .num_heads, self .head_dim).transpose(1 , 2 )
v = v.view(batch_size, -1 , self .num_heads, self .head_dim).transpose(1 , 2 )
scores = torch.matmul(q, k.transpose(-2 , -1 )) / (self .head_dim ** 0.5 )
scores = scores * adaptive_temperature.unsqueeze(1 ).unsqueeze(3 )
if attention_mask is not None :
scores = scores + attention_mask.unsqueeze(1 ).unsqueeze(2 ) * -1e9
attn = F.softmax(scores, dim=-1 )
output = torch.matmul(attn, v)
output = output.transpose(1 , 2 ).contiguous()
output = output.view(batch_size, -1 , self .hidden_dim)
return self .out_proj(output)
Step 3: Apply LoRA Fine-tuning for Artifact Suppression Implement parameter-efficient adaptation:
from peft import get_peft_model, LoraConfig
class TACAWithLoRA (nn.Module):
"""TACA combined with LoRA fine-tuning for artifact suppression"""
def __init__ (self, base_mm_dit, lora_r=8 , lora_alpha=16 ):
super ().__init__()
self .base_model = base_mm_dit
self .replace_cross_attention_with_taca()
lora_config = LoraConfig(
r=lora_r,
lora_alpha=lora_alpha,
target_modules=["q_proj" , "k_proj" , "v_proj" , "out_proj" ],
lora_dropout=0.05 ,
bias="none" ,
task_type="SEQ_2_SEQ_LM"
)
self .model = get_peft_model(self .base_model, lora_config)
def replace_cross_attention_with_taca (self ):
"""Replace all cross-attention modules with TACA"""
hidden_dim = self .base_model.config.hidden_size
num_heads = self .base_model.config.num_attention_heads
for module_name, module in self .base_model.named_modules():
if 'cross_attn' in module_name:
parent_name = '.' .join(module_name.split('.' )[:-1 ])
attr_name = module_name.split('.' )[-1 ]
parent = self .get_module_by_name(self .base_model, parent_name)
taca = TemperatureAdjustedCrossAttention(hidden_dim, num_heads)
setattr (parent, attr_name, taca)
@staticmethod
def get_module_by_name (module, name ):
"""Retrieve module by dot-separated name"""
for component in name.split('.' ):
module = getattr (module, component)
return module
def forward (self, text_ids, image_ids, timestep ):
"""Forward with TACA and LoRA"""
return self .model(text_ids, image_ids, timestep)
def get_trainable_parameters (self ):
"""Return only LoRA parameters for fine-tuning"""
return [p for n, p in self .model.named_parameters() if 'lora' in n]
Step 4: Training and Evaluation Implement fine-tuning on text-image alignment tasks:
class TACATrainer :
def __init__ (self, model, device='cuda' ):
self .model = model.to(device)
self .device = device
def train_taca_lora (self, dataloader, num_epochs=5 , learning_rate=1e-4 ):
"""
Fine-tune TACA + LoRA on T2I-CompBench for text-image alignment.
Objective: improve spatial relationships and shape accuracy in generated images.
"""
optimizer = torch.optim.AdamW(
self .model.get_trainable_parameters(),
lr=learning_rate
)
for epoch in range (num_epochs):
total_loss = 0
for batch in dataloader:
text_ids = batch['text_ids' ].to(self .device)
image_ids = batch['image_ids' ].to(self .device)
timesteps = batch['timesteps' ].to(self .device)
target_ids = batch['target_ids' ].to(self .device)
optimizer.zero_grad()
logits = self .model(text_ids, image_ids, timesteps)
loss = F.cross_entropy(
logits.reshape(-1 , logits.shape[-1 ]),
target_ids.reshape(-1 )
)
loss.backward()
torch.nn.utils.clip_grad_norm_(
self .model.get_trainable_parameters(), 1.0
)
optimizer.step()
total_loss += loss.item()
print (f"Epoch {epoch+1 } : Loss = {total_loss / len (dataloader):.4 f} " )
def evaluate_alignment (self, model_baseline, model_taca, test_dataloader ):
"""
Evaluate text-image alignment improvements.
Metrics: spatial relationship accuracy, shape accuracy
"""
results = {'baseline' : {}, 'taca' : {}}
for model_key, model in [('baseline' , model_baseline), ('taca' , model_taca)]:
spatial_correct = 0
shape_correct = 0
total = 0
model.eval ()
with torch.no_grad():
for batch in test_dataloader:
text_ids = batch['text_ids' ].to(self .device)
image_ids = batch['image_ids' ].to(self .device)
timesteps = batch['timesteps' ].to(self .device)
generated = model(text_ids, image_ids, timesteps)
spatial_acc = evaluate_spatial_accuracy(generated, batch['spatial_labels' ])
shape_acc = evaluate_shape_accuracy(generated, batch['shape_labels' ])
spatial_correct += spatial_acc
shape_correct += shape_acc
total += 1
results[model_key]['spatial_accuracy' ] = spatial_correct / total
results[model_key]['shape_accuracy' ] = shape_correct / total
return results
Practical Guidance
Temperature Value : Start with γ=2.0; higher values increase text influence (may reduce diversity)Timestep Threshold : t_thresh=0.5 (in normalized [0,1] scale) works for most modelsLoRA Rank : r=8 is good balance; larger ranks capture more artifact structureTraining Data : T2I-CompBench or similar with spatial/shape annotationsIntegration : Drop-in replacement for existing MM-DiT cross-attention layersPerformance Gains : FLUX.1-Dev shows 16.4% spatial improvement, 5.9% shape accuracy gainComputational Cost : Negligible overhead; LoRA adds <5% parametersUser Study : Strongly preferred outputs across alignment and quality metrics
Reference
Temperature scaling rebalances softmax competition in highly imbalanced token scenarios
Timestep-aware weighting exploits known structure: early phases emphasize layout, late phases detail
LoRA fine-tuning enables efficient artifact suppression without full model retraining
Cross-modal interaction is fundamentally asymmetric: text guides image, not vice versa
