批量构建大型下三角矩阵的内存友好实现方案咨询

Understanding the Core Problem #

You're stuck building a massive 90000×90000 all-ones lower triangular matrix—full initialization isn't feasible due to memory limits. Instead, you need to create row-wise batches (like 10 batches of 9000×90000 each) where every batch maps perfectly to the corresponding rows of the full lower triangular matrix.

Key Observation: Lower Triangular Matrix Structure #

For any global row index i (0-based), that row has 1s in columns 0 to i, and 0s everywhere else. When working with a batch starting at start_row and ending at end_row, each relative row k in the batch corresponds to global row start_row + k—so it needs 1s up to column start_row + k.

Solution 1: Batch-wise Matrix Initialization (Vectorized) #

We can use TensorFlow's broadcasting to generate each batch efficiently, no slow Python loops required:

import tensorflow as tf

N = 90000  # Total rows/columns of the full matrix
batch_size = 9000  # Rows per batch
num_batches = N // batch_size

for batch_idx in range(num_batches):
    # Define the global row range for this batch
    start_row = batch_idx * batch_size
    end_row = start_row + batch_size
    
    # Generate global row indices for the current batch
    batch_rows = tf.range(start_row, end_row)  # Shape: [batch_size]
    
    # Generate all column indices of the full matrix
    all_cols = tf.range(N)  # Shape: [N]
    
    # Create a boolean mask: True where column <= row (matches lower triangular pattern)
    # Expand batch_rows to [batch_size, 1] to broadcast against the [N] column array
    mask = all_cols <= tf.expand_dims(batch_rows, axis=1)  # Shape: [batch_size, N]
    
    # Convert mask to float32 to get the batch matrix
    batch_matrix = tf.cast(mask, tf.float32)
    
    # Use this batch matrix for your vector-matrix multiplication
    # Example: If you have a vector v of shape [N], compute the batch result
    # v = tf.random.normal([N])
    # batch_result = tf.matmul(batch_matrix, tf.expand_dims(v, axis=1))  # Shape: [batch_size, 1]
    
    # Store or process the batch result as needed

Why This Works #

• We never load the full 90000×90000 matrix into memory—only one batch exists at a time.
• The broadcasting operation is fully optimized by TensorFlow, making it fast and efficient.

Solution 2: Optimized for Vector-Matrix Multiplication (No Explicit Matrix) #

If your end goal is to compute y = L @ v (where L is the lower triangular matrix), you don't need to create the matrix at all!

Multiplying an all-ones lower triangular matrix with a vector v is exactly the same as computing the prefix sum of v:
y[i] = sum_{j=0 to i} v[j]

This approach is way more memory-efficient—here's how to implement it batch-wise:

import tensorflow as tf

N = 90000
batch_size = 9000
num_batches = N // batch_size

# Example input vector
v = tf.random.normal([N])

# Compute the full cumulative sum (O(N) time, minimal memory usage)
cumulative_v = tf.cumsum(v)

# Extract results in batches
for batch_idx in range(num_batches):
    start_row = batch_idx * batch_size
    end_row = start_row + batch_size
    batch_result = cumulative_v[start_row:end_row]
    
    # Process the batch result here

Why This Is the Better Option #

• No matrix is ever created—memory usage is just for the input vector and its cumulative sum (both O(N) space, negligible compared to O(N²) for the full matrix).
• Computation is drastically faster than matrix multiplication, since it's a simple cumulative sum operation.

Why Your Previous Approach Failed #

tf.linalg.LinearOperatorLowerTriangular treats each element in the batch as an independent lower triangular matrix. When you passed tf.ones((3,3,9)), it generated a small lower triangular matrix for every 3×9 slice—so every row in the batch started 1s from column 0, which doesn't match the global lower triangular structure you need.

内容的提问来源于stack exchange，提问作者Devanshi Sukhija