如何使用PyTorch仅对矩阵的上三角元素执行Softmax操作？

To solve this problem, we need to apply the Softmax function exclusively to the upper triangular elements (including the diagonal) of a batch of square matrices, while leaving the lower triangular elements unchanged. Here's a step-by-step implementation using PyTorch:

Approach #

1. Identify Upper Triangular Indices: Use torch.triu_indices() to get the row and column indices of all upper triangular elements (including the diagonal).
2. Extract Elements: Pull out these upper triangular elements from the input tensor into a 1D vector per batch.
3. Apply Softmax: Compute the Softmax over the extracted upper triangular elements.
4. Reconstruct Tensor: Create a copy of the original tensor, then replace the upper triangular positions with the Softmaxed values.

Code Implementation #

import torch

def softmax_upper_triangular(input_tensor):
    batch_size, height, width = input_tensor.shape
    assert height == width, "Input must contain square matrices (height == width)"
    
    # Get indices for upper triangular elements (including diagonal)
    row_indices, col_indices = torch.triu_indices(height, width)
    
    # Extract the upper triangular elements from each batch
    upper_elements = input_tensor[:, row_indices, col_indices]
    
    # Apply Softmax along the dimension of the extracted elements
    softmaxed_upper = torch.softmax(upper_elements, dim=1)
    
    # Create output tensor by copying input and updating upper triangular values
    output_tensor = input_tensor.clone()
    output_tensor[:, row_indices, col_indices] = softmaxed_upper
    
    return output_tensor

# Test with your input tensor
input_tensor = torch.tensor([[[1.9392, -1.9266, 0.9664],
                              [0.0000, -1.9266, 0.9664],
                              [0.0000, -0.0000, 0.9664]]])

result = softmax_upper_triangular(input_tensor)
print(result)

Output #

Running this code will produce exactly the desired output:

tensor([[[0.4596, 0.0096, 0.1737],
         [0.0000, 0.0096, 0.1737],
         [0.0000, -0.0000, 0.1737]]])

Explanation #

• Triangular Indices: torch.triu_indices(3,3) returns indices ([0,0,0,1,1,2], [0,1,2,1,2,2]), which correspond to all upper triangular positions in a 3x3 matrix.
• Softmax Calculation: The extracted upper elements are [1.9392, -1.9266, 0.9664, -1.9266, 0.9664, 0.9664]. Applying Softmax normalizes these values so their sum is 1, resulting in the values seen in the output.
• Preserving Lower Triangular: By cloning the original tensor first, we retain all lower triangular elements (like the 0s in positions (1,0) and (2,0), (2,1)) exactly as they were.

内容的提问来源于stack exchange，提问作者shu.liu