File size: 2,508 Bytes
262fa3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """Loss functions for latent-reasoning ASR training."""
from typing import Any
import torch
import torch.nn as nn
import torch.nn.functional as F
def _fmt(x: torch.Tensor) -> str:
"""Format a one-dimensional tensor for logging."""
return " | ".join([f"d{k}:{x[k].item():.3f}" for k in range(x.numel())])
def contrastive_loss(view_a: torch.Tensor, view_b: torch.Tensor, temperature: float) -> torch.Tensor:
"""InfoNCE contrastive loss over batch (symmetric)."""
assert view_a.dim() == 2 and view_b.dim() == 2, "contrastive_loss expects (B, D) inputs"
B = view_a.size(0)
if B <= 1:
return view_a.new_tensor(0.0)
z1 = F.normalize(view_a, dim=-1)
z2 = F.normalize(view_b, dim=-1)
logits = torch.matmul(z1, z2.t()) / float(temperature)
labels = torch.arange(B, device=view_a.device)
loss_a = F.cross_entropy(logits, labels)
loss_b = F.cross_entropy(logits.t(), labels)
return 0.5 * (loss_a + loss_b)
def trajectory_regularization_loss(thoughts: torch.Tensor, alpha: float = 0.5) -> torch.Tensor:
"""Trajectory regularization loss.
Encodes the intuition that the latent reasoning trajectory should:
1. Stay anchored to the initial acoustic state h_0 (semantic grounding)
2. Evolve smoothly between consecutive steps (temporal coherence)
Args:
thoughts: shape (B, N, D), where thoughts[:, 0, :] is the initial state h_0
alpha: weight balancing anchor vs smoothness (default: 0.5)
Returns:
Scalar loss combining anchor and smoothness terms
"""
N = thoughts.size(1)
# Guard: empty trajectory (e.g. n_latent=0) – return 0.
if N == 0:
return thoughts.new_tensor(0.0)
# Extract initial acoustic anchor h_0
h_0 = thoughts[:, 0:1, :] # shape: (B, 1, D)
# 1. Acoustic Anchor Loss (distance of all states from h_0)
anchor_loss = (thoughts - h_0).pow(2).sum(dim=-1).mean()
# 2. Step-wise Smoothness Loss (distance between adjacent states).
# When N==1 (e.g. DEQ single fixed-point), there are no consecutive pairs,
# so smooth_loss is zero – avoids mean() on an empty tensor returning NaN.
if N > 1:
diffs = thoughts[:, 1:, :] - thoughts[:, :-1, :] # shape: (B, N-1, D)
smooth_loss = diffs.pow(2).sum(dim=-1).mean()
else:
smooth_loss = thoughts.new_tensor(0.0)
# Combine final regularization loss
total_trajectory_loss = alpha * anchor_loss + (1.0 - alpha) * smooth_loss
return total_trajectory_loss
|