メニュー
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.
デフォルトでは、最初にソースを確認する Prompt が選択されています。直接コマンドに切り替えるか、ローカルコピーをダウンロードすることもできます。
インストールを決める前に、SKILL.md と SkillsMP に表示されている付属ファイルをお読みください。
Codex または Claude でインストール この Prompt をコピーして Codex、Claude、または他のアシスタントに貼り付けると、Skill ページを確認してインストールできます。
直接コマンドでは確認用 Prompt が省略されます。実行前にソースを確認してください。
npx skills add https://github.com/ADu2021/skillXiv --skill multimodal-diffusion-alignmentコマンドは1行のまま表示されます。コピー前に横へスクロールして全体を確認してください。
ローカルで確認しますか？SkillsMP が現在取得できるファイルをダウンロードできます。
Convert arXiv papers into ready-to-use agent skills using category-aware extraction. First classifies the paper into one or more of 11 research categories, then applies a specialized extraction pipeline for each category — because different types of papers produce different types of usable knowledge. A single paper can yield multiple skills if it spans categories. Use this skill whenever the user wants to turn a paper into a skill, extract practical techniques from research, build a skill library from papers, convert arXiv papers into reusable agent instructions, or batch-process multiple papers into skills. Also trigger when someone asks about extracting actionable knowledge from papers, making research practical for LLM agents, or systematically converting academic contributions into structured agent capabilities.
Establish regret bounds for behavior cloning with discretized actions combining statistical error and quantization error terms. Prove smoothness requirements for safe quantizer design, show that learning-based quantizers fail these requirements, and propose model-based augmentation to reduce error dependence from H² to H.
Dynamically allocate LoRA ranks per-layer during fine-tuning instead of using fixed uniform ranks. Learn optimal rank for each layer and subject via variational framework with discretized exponential distribution, reducing memory footprint while maintaining fidelity and text-alignment.
SOC 職業分類に基づく
SKILL.md を表示中
| 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. |
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.
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[]
num_image_tokens = image_ids.shape[]
torch.no_grad():
embeddings = .model.embed_tokens(text_ids, image_ids, timestep)
hidden = .model.forward_with_hooks(embeddings, return_attention=)
attention_logits = hidden[]
standard_probs = F.softmax(attention_logits, dim=-)
text_prob_per_position = standard_probs.mean(dim=(, ))
text_contribution = text_prob_per_position.().item()
()
()
()
{
: attention_logits,
: text_contribution
}
():
result = .analyze_softmax_suppression(text_ids, image_ids, timestep)
expected = / ( + )
actual = result[]
imbalance = (, (expected - actual) / expected)
imbalance
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
# Learnable temperature parameter (can be fixed or trainable)
self.temperature = nn.Parameter(torch.tensor(temperature_base))
# Timestep embedding for adaptive temperature
self.timestep_mlp = nn.Sequential(
nn.Linear(1, 64),
nn.SiLU(),
nn.Linear(64, 1),
nn.Sigmoid() # Output in [0, 1]
)
# Standard attention components
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]
# Compute temperature multiplier based on timestep
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) # (batch, 1)
# Piecewise temperature function
# Early steps (high t): use full temperature
# Late steps (low t): use temperature=1.0 for fine details
threshold = 0.5 # Normalized timestep threshold
uses_enhanced_temperature = (timestep_tensor > threshold).float()
adaptive_temperature = 1.0 + (self.temperature - 1.0) * uses_enhanced_temperature * temp_multiplier
# (batch, 1)
# Project text and image to Q, K, V space
q = self.q_proj(image_hidden) # Image queries: (B, N_img, D)
k = self.k_proj(text_hidden) # Text keys: (B, N_txt, D)
v = self.v_proj(text_hidden) # Text values: (B, N_txt, D)
# Reshape for multi-head attention
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)
# Compute attention scores
scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
# (batch, num_heads, num_image, num_text)
# Apply temperature scaling to cross-modal logits
# This rebalances softmax competition between high-cardinality image tokens
# and low-cardinality text tokens
scores = scores * adaptive_temperature.unsqueeze(1).unsqueeze(3)
# Apply attention mask if provided
if attention_mask is not None:
scores = scores + attention_mask.unsqueeze(1).unsqueeze(2) * -1e9
# Softmax over text tokens
attn = F.softmax(scores, dim=-1)
# Apply attention
output = torch.matmul(attn, v)
# (batch, num_heads, num_image, head_dim)
# Merge heads
output = output.transpose(1, 2).contiguous()
output = output.view(batch_size, -1, self.hidden_dim)
return self.out_proj(output)
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
# Replace cross-attention layers with TACA
self.replace_cross_attention_with_taca()
# Apply LoRA to attention layers
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:
# Get parent module
parent_name = '.'.join(module_name.split('.')[:-1])
attr_name = module_name.split('.')[-1]
parent = self.get_module_by_name(self.base_model, parent_name)
# Replace with TACA
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]
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 for LoRA parameters only
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()
# Forward pass
logits = self.model(text_ids, image_ids, timesteps)
# Reconstruction loss (match target image tokens)
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):.4f}")
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)
# Generate image
generated = model(text_ids, image_ids, timesteps)
# Evaluate against ground truth
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