TensorTEE: Unifying Heterogeneous TEE Granularity for Efficient Secure Collaborative Tensor Computing

Training a large model securely across a CPU and an accelerator is slow not because encryption is expensive but because the two sides protect memory in different sized chunks — TensorTEE makes both of them think in whole tensors.

Why this problem is worth solving

Nobody trains a large model on one chip any more, and increasingly nobody trains one on hardware they fully control. That combination is what makes this paper’s setting real. In the collaborative arrangement the authors study — DeepSpeed’s ZeRO-Offload — the accelerator handles the forward and backward passes while the CPU keeps the high-precision weights and optimiser state and runs the Adam update. Splitting the job this way is what lets a ten-billion-parameter GPT-2 train on a single V100. But now the model, the gradients and the training data all cross a bus between two chips, several times per batch.

Protecting that traffic is what a

What people did before, and where it broke

• Protect CPU memory per cacheline, the SGX way. Each 64-byte block gets its own 56-bit version number and integrity tag, and the version numbers themselves are protected by a Merkle tree. This supports arbitrary random access, but the metadata is 11% storage overhead and every read triggers extra memory accesses plus a tree walk. On the memory-bound Adam update the paper measures 3.65x slowdown at eight threads.
• Protect accelerator memory per tensor. Because accelerator workloads stream through arrays predictably, one version number can cover a whole tensor. This is what MGX and Securator do, and it removes the Merkle tree entirely. But the integrity tags stay fine-grained, and they still cost close to 10% of the accelerator’s scarce memory and around 12% in performance.
• Declare the tensors in software. SoftVN keeps tensor structure in an on-chip table that the program explicitly populates. It works for simple kernels, but the table has few entries, so the programmer has to manage creation and release by hand, and complex tiled kernels overrun it.
• Just make the integrity tags coarser. Moving from 64-byte to 512-byte or 4 KB tags shrinks storage nicely. The catch is that a tag can only be checked once the whole block has arrived, so decrypted data sits waiting on a verification that has not finished. At 4 KB granularity the pipeline bubbles cost 13%.
• Predict rather than wait. Plutus uses predictive execution to avoid the extra memory accesses for integrity tags, but leaves the storage overhead untouched.
• Relay transfers through unprotected memory. Graviton-style protocols handle the cross-device hop by having the sender decrypt into a non-secure buffer, re-encrypt it for the wire, and having the receiver decrypt and store. Correct, but it triples the memory traffic and, because a single encryption engine supplies only 8 GB/s where 20 GB/s is needed, forces transfer and computation to run one after the other.
• The gap. Every design optimised one side of the boundary. Nobody had noticed that the boundary itself was the expensive part, and that the two granularities could be made to agree.

How TensorTEE unifies the two granularities

The idea

The cacheline was never the natural unit for this workload. The whole design rests on one measurement: during optimiser updates, tensors grow to megabytes while their number grows only to a few hundred. A protection scheme that assigns metadata per 64-byte block is paying for a flexibility this workload never uses. Raising the unit to the tensor collapses the metadata to something that fits on chip, which removes both the extra memory accesses and the Merkle tree in one move. The same observation — that accelerator and tensor workloads have regular enough structure that tracking them at tensor granularity beats tracking them at hardware-page or cacheline granularity — is what drives tensor-level migration in unified memory systems, covered elsewhere on this site, and it is striking that the argument transfers unchanged from a performance problem to a security one.

Infer the structure instead of asking for it. SoftVN had the right idea and the wrong interface: it required the programmer to declare tensors, which does not scale past simple kernels. TensorTEE’s answer is to detect them in hardware, and the placement of the detector is the clever part. Putting it in the memory controller but wiring it to receive requests directly from the cores means it sees virtual addresses in their original order — before cache evictions and prefetchers scramble the pattern, and before physical page boundaries chop up the contiguity. That one placement decision is what makes automatic detection viable, and it converges fast: after a single iteration of Adam the table is essentially complete, and a 256x256 matrix multiply tiled 64x64 reaches a 98.8% precise-hit rate after one pass. The broader move — noticing that the hardware was already carrying the signal it needed and simply had to be wired to look at it, rather than demanding that software supply the answer — recurs whenever a design replaces an external annotation with an internal observation, covered elsewhere on this site, and it is usually the difference between a mechanism that works on the benchmark and one that works on arbitrary code.

Let the check run late, and stop the data at the door. Coarse integrity tags create a dilemma — bigger tags save storage but force decrypted data to wait for a verification that cannot finish until the whole block arrives. TensorTEE resolves it by refusing to wait: computation proceeds on decrypted-but-unverified data while the tag is regenerated in parallel. The security argument is that a bare accelerator has no operating system and no file access, so temporarily corrupted intermediate values cannot be turned into anything worse, provided nothing tainted ever leaves. Poison bits propagate through dependent tensors and a verification barrier blocks communication until they clear, while accelerator code is deliberately excluded from the relaxation and verified the strict way. The pattern of allowing unverified work to proceed and enforcing correctness only where results become observable is the same reasoning that makes speculative execution safe when a hardware contract bounds what can escape, covered elsewhere on this site.

Following a gradient from accelerator to CPU

This is required because the CPU cannot interpret the metadata of the accelerator even with the same key — the two enclaves use different hardware engines and different metadata formats, so the ciphertext is mutually unreadable.

Why it works

• Metadata shrinks until it fits on chip. One version number per tensor instead of per cacheline → a few hundred entries covers an entire optimiser step → the Merkle tree and its traversal disappear entirely, along with the extra memory access on every read.
• Detection is automatic and converges within one iteration. The analyser sees virtual addresses straight from the cores, unshuffled by caches → dimensional streaming patterns are visible and partial entries merge in multiple directions → tensor structure is essentially complete after one pass, and training runs tens of thousands of them.
• Correctness of the shared version number is enforced, not assumed. A bitmap with an updating flag and a bit-state marker → asserts that each cacheline in a tensor is written exactly once before the version increments → the on-chip tensor version stays consistent with the off-chip per-cacheline view, so the scheme composes with ordinary protection.
• Verification stops blocking computation. Tag regeneration runs in parallel with the kernel instead of gating it → the 13% pipeline-bubble cost of 4 KB tags falls to 2.5% → and storage overhead for tags becomes negligible rather than the roughly 10% the accelerator could not spare.
• Matching granularity removes the relay entirely. Both sides now interpret the same ciphertext and metadata → no decrypt, stage, re-encrypt sequence through unprotected memory → the 84% of DRAM bandwidth that re-encryption consumed is returned, which is what lets transfer and computation finally overlap.

What it costs (trade-offs and scope)

• It requires new hardware in the memory controller. TenAnalyzer, the Tensor Filter, the bitmap cache and the poison bits are all additions to the chip. They are small — 24 KB of storage, about 0.0072 square millimetres at 7 nm — but this is a silicon proposal, not something deployable on current parts.
• The table has a fixed capacity. The Meta Table holds 512 entries, and the paper is explicit that once a workload touches more than 512 tensors at once the benefit gradually fades. The design suits workloads with few, large tensors, which is exactly what it was measured on.
• Delayed verification is a deliberate weakening. Intermediate tensors can be tampered with for a window before detection. The security argument rests on the accelerator being a bare accelerator with no operating system, no file I/O and non-delayed verification of its own code — a reasonable model for a TPU-like device, but one that would not survive on a more general accelerator.
• The tensor tag is an XOR of cacheline tags. The authors argue this does not meaningfully reduce forgery resistance since the output space stays at 56 bits, still requiring around two to the fifty-sixth blind attempts. It is a simple construction rather than a proven one, chosen because it is order-insensitive and therefore survives tiling.
• The first iteration is slower. Until the Meta Table is populated, TensorTEE carries higher overhead than the steady state. Training runs tens of thousands of iterations so this amortises away, but a short or one-shot workload would not see the benefit.
• Everything is simulated. The evaluation uses an enhanced Gem5 for the CPU and a custom cycle-accurate simulator for a TPUv3-like accelerator, and side channels, denial of service and adversarial attacks are all explicitly out of scope.

The evaluation

Check yourself

Sources

The paper is “TensorTEE: Unifying Heterogeneous TEE Granularity for Efficient Secure Collaborative Tensor Computing” by Husheng Han, Xinyao Zheng, Yuanbo Wen, Yifan Hao, Erhu Feng, Ling Liang, Jianan Mu, Xiaqing Li, Tianyun Ma, Pengwei Jin, Xinkai Song, Zidong Du, Qi Guo and Xing Hu, published at ASPLOS 2024 Volume 4, pages 282-297, DOI 10.1145/3622781.3674168. The authors are with the Institute of Computing Technology of the Chinese Academy of Sciences and collaborating institutions including Shanghai Jiao Tong University and Cambricon Technologies.

This explainer is written from the full text, read via the arXiv version at arxiv.org/abs/2407.08903. Every quantitative claim above — the 4.0x average and 5.5x maximum speedup, the 2.1% overhead versus unprotected execution, the 3.65x SGX slowdown at eight threads, the 13.0% cost of 4 KB integrity tags against 2.5% with delayed verification, the 18.7x communication improvement, the 84% bandwidth figure for re-encryption, the 98.8% hit rate on a tiled matrix multiply, the 512-entry table limit and the 24 KB and 0.0072 square millimetre hardware cost — comes from the paper’s own sections 3, 4 and 6. Readers should note that the baseline is a combination the authors assembled (SGX-style CPU protection with MGX-style accelerator protection) rather than a published unified system, since no unified heterogeneous design existed to compare against.