Skip to content

TensorStream

Structuring Multimodal Data for Causal Transformers

Problem

Generative causal transformers are 1D computers. All instructions need to be serialized into a 1D token stream. For text-only data it's simple because humans process it the same way. All text-only data is stored/processed as 1D token streams. The problem gets challenging with multistream/multimodal data e.g. videos. Not only do videos have separate time-aligned audio and vision streams but also can have multiple metadata streams like transcriptions, narration, or even actions for robotics/video game data. This presents a challenge - we need some protocol for forming this token stream.

The problem of preparing this data is 2-stage. The first is about organizing the raw data into a causal chain. The second is about storing that data in a way it can flexibly be used and ultimately converted into an input to a transformer model.

Design

We solve this problem with a data structure called TensorStream and a library built around it.


First we read in the raw data into packets we call Events which are tagged with a rich set of descriptors that are tracked and updated throughout the pipeline. There’s no complex organization of the data, you just read things in, tag it with the relevant descriptors and throw it on the pile. At this stage the descriptors we care about are the modality and the “time” of the given event which is just some floating point number which communicates our prior idea about ordering.

We solve this problem with a data structure called TensorStream and a library built around it.


First we read in the raw data into packets we call Events which are tagged with a rich set of descriptors that are tracked and updated throughout the pipeline. There’s no complex organization of the data, you just read things in, tag it with the relevant descriptors and throw it on the pile. At this stage the descriptors we care about are the modality and the “time” of the given event which is just some floating point number which communicates our prior idea about ordering.

We solve this problem with a data structure called TensorStream and a library built around it.


First we read in the raw data into packets we call Events which are tagged with a rich set of descriptors that are tracked and updated throughout the pipeline. There’s no complex organization of the data, you just read things in, tag it with the relevant descriptors and throw it on the pile. At this stage the descriptors we care about are the modality and the “time” of the given event which is just some floating point number which communicates our prior idea about ordering.

Once we have this we need to schedule a TensorStream. This takes a list of Events along with a user-specified priority list of modalities and sorts it using the start and end times of the Events resolving conflicts using the modality priority list. 

Once we have this we need to schedule a TensorStream. This takes a list of Events along with a user-specified priority list of modalities and sorts it using the start and end times of the Events resolving conflicts using the modality priority list. 

This priority list is a useful tool when training omni models which need to know how to model p(A|B) but also p(B|A) for (B, A) which can depend on each other. (e.g. A=chunk of frames, B=narration happening at the same time; or A=caption, B=image). During dataloading you can randomize over this priority list and get a very diverse set of valid tasks from the same sample.

This priority list is a useful tool when training omni models which need to know how to model p(A|B) but also p(B|A) for (B, A) which can depend on each other. (e.g. A=chunk of frames, B=narration happening at the same time; or A=caption, B=image). During dataloading you can randomize over this priority list and get a very diverse set of valid tasks from the same sample.

At this point we have the first stage of our problem solved - a causal chain of data prepared for a generative causal transformer - events sorted earlier in the sequence will have their tokens inserted earlier in the context of the transformer. We can move on to the second problem; writing a library around our TensorStreams which allow us to flexibly and intuitively interface with our data in a way that exposes its useful features.


The main difficulty in this stage is token accounting - in dataloading we often need accurate token accounting even though there might be complex transforms outside of it. Different embedders might yield different token counts (e.g. due to different convolution patterns). A parameterized post-embedding transform (like PixelUnshuffle) might also alter it. How do you know how much to pad up or truncate before it even happens? Additionally even if you are truncating often you cant just get rid of the data e.g. SigLIP takes as input full images - you can’t pass in half of an image and expect good results. 


To solve this problem our design allows the data and its “descriptors” to be transformed on independent tracks. The Events are seeded with initial measurements and adjusted in the preprocessors based on the configuration of the model the data is being preprocessed for. Our TensorStreams allow you to inspect the full data or the tokenized data. The 3 additional shape descriptors we have now are:


  • dims_real - the shape of the Event after individual preprocessing. This information is useful for the embedders e.g. position embedding resizing in SigLIP.

  • dims_virtual - the shape of the Event after all transforms are applied (including embedders). This is used for accounting regarding the final token shape of the data. e.g. sequence padding, truncating, buffering.

  • idx_range - pointers to the start and end of the valid tokens of an Event. When we need to slice through an Event but its full data is required for embedding we adjust the idx_range.


Once each Event is tagged with its descriptors we can pass it through packing, collation, and embedding all of which can be implemented very intuitively with minimal overhead.


That's it! With this information we’re able to build a strong library around TensorStreams which allows us to express usually inconvenient operations in simple intuitive terms. Lets walk through our training pipeline and inspect some of the operations that are simplified using our TensorStreams library.

At this point we have the first stage of our problem solved - a causal chain of data prepared for a generative causal transformer - events sorted earlier in the sequence will have their tokens inserted earlier in the context of the transformer. We can move on to the second problem; writing a library around our TensorStreams which allow us to flexibly and intuitively interface with our data in a way that exposes its useful features.


The main difficulty in this stage is token accounting - in dataloading we often need accurate token accounting even though there might be complex transforms outside of it. Different embedders might yield different token counts (e.g. due to different convolution patterns). A parameterized post-embedding transform (like PixelUnshuffle) might also alter it. How do you know how much to pad up or truncate before it even happens? Additionally even if you are truncating often you cant just get rid of the data e.g. SigLIP takes as input full images - you can’t pass in half of an image and expect good results. 


To solve this problem our design allows the data and its “descriptors” to be transformed on independent tracks. The Events are seeded with initial measurements and adjusted in the preprocessors based on the configuration of the model the data is being preprocessed for. Our TensorStreams allow you to inspect the full data or the tokenized data. The 3 additional shape descriptors we have now are:


  • dims_real - the shape of the Event after individual preprocessing. This information is useful for the embedders e.g. position embedding resizing in SigLIP.

  • dims_virtual - the shape of the Event after all transforms are applied (including embedders). This is used for accounting regarding the final token shape of the data. e.g. sequence padding, truncating, buffering.

  • idx_range - pointers to the start and end of the valid tokens of an Event. When we need to slice through an Event but its full data is required for embedding we adjust the idx_range.


Once each Event is tagged with its descriptors we can pass it through packing, collation, and embedding all of which can be implemented very intuitively with minimal overhead.


That's it! With this information we’re able to build a strong library around TensorStreams which allows us to express usually inconvenient operations in simple intuitive terms. Lets walk through our training pipeline and inspect some of the operations that are simplified using our TensorStreams library.

Examples

Sequence packing

One of the most useful utilities is being able to slice and concatenate TensorStreams on the sequence dimension for applications like sequence packing. Using these we can pack samples with similar complexity to common text-only packing implementations:

One of the most useful utilities is being able to slice and concatenate TensorStreams on the sequence dimension for applications like sequence packing. Using these we can pack samples with similar complexity to common text-only packing implementations:

def pack_sequences(tensor_stream_iter: Iterable, L: int):
  buf_stream = TensorStream([])
  for tensor_stream in tensor_stream_iter():
    # you can concatenate TensorStream as if they were lists
    buf_stream += tensor_stream

    while len(buf) >= L + 1:
      # TensorStream.slice handles updates to all Event descriptors.
      # e.g. slicing through image tokens will duplicate the Event in the resulting
      # TensorStream and adjust the idx_ranges accordingly
      packed_data, buf = buf[:L+1], buf[L+1:]
      yield packed_data

Embedding

Once the data leaves the dataloader we need to embed it for the transformer backbone. Here we need to group events in our TensorStream by the modality, embed each modality sub-stream, and emit the embedded tensors for final assembly.

Once the data leaves the dataloader we need to embed it for the transformer backbone. Here we need to group events in our TensorStream by the modality, embed each modality sub-stream, and emit the embedded tensors for final assembly.

def embed_stream(self, tensor_stream: TensorStream) -> dict[ModalityType, torch.Tensor]:
  per_modality_stream: dict[ModalityType, TensorStream] = group_streams(
    tensor_stream, group_fn=lambda ev: ev.modality
  )
  # TensorStream.compact concatenates all tokens on first dimension.
  # we do this to batch embedding
  per_modality_compact_stream = {k: v.compact() for k, v in per_modality_stream.items()}

  embedded_stream = {}
  for modality in per_modality_compact_stream:
    tensor = per_modality_compact_stream[modality]
    input_data = (tensor,)
    if modality == VisionType:
      # NaViT embedding requires image dimensions
      # we use `virtual=False` because we want the true image dimensions
      token_dims = [event.dims(virtual=False) for event in per_modality_stream[VisionType]]
      input_data = (tensor, token_dims)
    embedded_stream[modality] = self.embed_fns[modality](*input_data)
  return embedded_stream

Interleaving

After embedding we need to interleave the data into a (T, D) shaped tensor as input to our transformer. Event level attributes allow us to create arbitrary masks over the input sequence, in this example modality and token validity, which lets us reconstruct the interleaved tensor through indexing. 

After embedding we need to interleave the data into a (T, D) shaped tensor as input to our transformer. Event level attributes allow us to create arbitrary masks over the input sequence, in this example modality and token validity, which lets us reconstruct the interleaved tensor through indexing. 

def interleave_compact_embeddings(
  compact_embedded_streams: dict[ModalityType, torch.Tensor], source_s: TensorStream
) -> torch.Tensor:
  T, D = len(source_s), next(iter(compact_embedded_streams.values())).shape[-1]

  per_modality_stream = group_streams(source_s, group_fn=lambda ev: ev.modality)
  mod = modality_mask(source_s)
  interleaved_tensor = torch.empty((T, D))

  for mod_type, compact_tensor in compact_embedded_streams.items():
    # the compact streams have multiple events concatenated on the sequence dimension. 
    # some of those events could be sliced so we create a token validity mask
    valid_token_mask = valid_token_mask(per_modality_stream[mod_type])
    assert valid_token_mask.sum() == (mod == mod_type.value).sum()
    # embeddings are compacted in order; we can insert using indexing
    interleaved_tensor[mod == mod_type.value] = compact_tensor[valid_token_mask]
  return interleaved_tensor

# example simple library functions
def event_mask(
  s: TensorStream,
  tag_fn: Callable[[Event], int | None],
) -> torch.Tensor:
  def to_label(ev: Event) -> Any:
    label = tag_fn(ev)
    return [label] * ev.num_tokens()
  # TensorStream.map_compact fuses map with compact to save time on Tensor creation
  return s.map_compact(to_label).squeeze(-1)

def modality_mask(ts: TensorStream) -> torch.Tensor:
  # creates a (T_virtual,) shaped tensor s.t. mask[i] = modality.value
  return event_mask(ts, lambda ev: ev.modality.value)

def valid_token_mask(ts: TensorStream) -> torch.Tensor:
  # returns a (T_real,) shaped tensor s.t. mask[i] = 1.0 if valid
  # e.g. if the TensorStream has 1 Event with 5 tokens but its idx_range=(1,3):
  # returns: torch.Tensor([0, 1, 1, 0, 0]).bool()
  return ts.map_compact(
    lambda ev: [
      int(ev.idx_range[0] <= i < ev.idx_range[1]) 
      for i in range(ev.num_tokens(virtual=False))
    ]
  ).bool()

Conclusion

TensorStream localizes the complexity of multimodal data handling into a single abstraction. What's exposed is simple and familiar. Users just see lists they can slice and concatenate. Events they can group and map over. Tensors they can index into. The interface matches the mental model perfectly because we contained complexity rather than letting it leak throughout the codebase.

TensorStream localizes the complexity of multimodal data handling into a single abstraction. What's exposed is simple and familiar. Users just see lists they can slice and concatenate. Events they can group and map over. Tensors they can index into. The interface matches the mental model perfectly because we contained complexity rather than letting it leak throughout the codebase.

@article{perceptron2025tensorstream,
  author = {Maciej Kilian and PerceptronAI},
  title = {TensorStream: Structuring Multimodal Data for Causal Transformers},
  journal = {PerceptronAI Research Blog},
  year = {2025},
  note = {https://perceptron.inc/blog/tensorstream},
}
Footnote

What we’re describing here is our internal framework for handling TensorStreams. For our release of Isaac 0.1 we shared a min-cut of this which is required to get inference working. If there’s enough interest we can invest some time in releasing the full version.